Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

WIP perf: significantly reduce DB calls toGetWorkspaceByAgentID via caching workspace information in Agent API#20662

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to ourterms of service andprivacy statement. We’ll occasionally send you account related emails.

Already on GitHub?Sign in to your account

Open
cstyan wants to merge2 commits intomain
base:main
Choose a base branch
Loading
fromcallum/workspace-agent-db-call-pt2
Open
Show file tree
Hide file tree
Changes fromall commits
Commits
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
131 changes: 125 additions & 6 deletionscoderd/agentapi/api.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -2,6 +2,7 @@ package agentapi

import (
"context"
"database/sql"
"io"
"net"
"net/url"
Expand All@@ -21,6 +22,7 @@ import (
"github.com/coder/coder/v2/coderd/appearance"
"github.com/coder/coder/v2/coderd/connectionlog"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/coderd/database/pubsub"
"github.com/coder/coder/v2/coderd/externalauth"
"github.com/coder/coder/v2/coderd/notifications"
Expand All@@ -36,6 +38,35 @@ import (
"github.com/coder/quartz"
)

const workspaceCacheRefreshInterval = 5 * time.Minute

// CachedWorkspaceFields contains workspace data that is safe to cache for the
// duration of an agent connection. These fields are used to reduce database calls
// in high-frequency operations like stats reporting and metadata updates.
//
// IMPORTANT: ACL fields (GroupACL, UserACL) are NOT cached because they can be
// modified in the database and we must use fresh data for authorization checks.
//
// Prebuild Safety: When a prebuild is claimed, the owner_id changes in the database
// but the agent connection persists. Currently we handle this by periodically refreshing
// the cached fields (every 5 minutes) to pick up changes like prebuild claims.
type CachedWorkspaceFields struct {
// Identity fields
ID uuid.UUID
OwnerID uuid.UUID
OrganizationID uuid.UUID
TemplateID uuid.UUID

// Display fields for logging/metrics
Name string
OwnerUsername string
TemplateName string

// Lifecycle fields needed for stats reporting
AutostartSchedule sql.NullString
DormantAt sql.NullTime
}

// API implements the DRPC agent API interface from agent/proto. This struct is
// instantiated once per agent connection and kept alive for the duration of the
// session.
Expand All@@ -54,6 +85,8 @@ type API struct {
*SubAgentAPI
*tailnet.DRPCService

cachedWorkspaceFields CachedWorkspaceFields

mu sync.Mutex
}

Expand DownExpand Up@@ -92,14 +125,31 @@ type Options struct {
UpdateAgentMetricsFn func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric)
}

func New(opts Options) *API {
func New(opts Options, workspace database.Workspace) *API {
if opts.Clock == nil {
opts.Clock = quartz.NewReal()
}

api := &API{
opts: opts,
mu: sync.Mutex{},
cachedWorkspaceFields: CachedWorkspaceFields{
ID: workspace.ID,
OwnerID: workspace.OwnerID,
OrganizationID: workspace.OrganizationID,
TemplateID: workspace.TemplateID,
Name: workspace.Name,
OwnerUsername: workspace.OwnerUsername,
TemplateName: workspace.TemplateName,
AutostartSchedule: sql.NullString{
String: workspace.AutostartSchedule.String,
Valid: workspace.AutostartSchedule.Valid,
},
DormantAt: sql.NullTime{
Time: workspace.DormantAt.Time,
Valid: workspace.DormantAt.Valid,
},
},
mu: sync.Mutex{},
}

api.ManifestAPI = &ManifestAPI{
Expand DownExpand Up@@ -139,6 +189,7 @@ func New(opts Options) *API {

api.StatsAPI = &StatsAPI{
AgentFn: api.agent,
WorkspaceFn: api.workspace,
Database: opts.Database,
Log: opts.Log,
StatsReporter: opts.StatsReporter,
Expand All@@ -162,10 +213,11 @@ func New(opts Options) *API {
}

api.MetadataAPI = &MetadataAPI{
AgentFn: api.agent,
Database: opts.Database,
Pubsub: opts.Pubsub,
Log: opts.Log,
AgentFn: api.agent,
RBACContextFn: api.rbacContext,
Database: opts.Database,
Pubsub: opts.Pubsub,
Log: opts.Log,
}

api.LogsAPI = &LogsAPI{
Expand DownExpand Up@@ -205,6 +257,10 @@ func New(opts Options) *API {
Database: opts.Database,
}

// Start background cache refresh loop to handle workspace changes
// like prebuild claims where owner_id and other fields may be modified in the DB.
go api.startCacheRefreshLoop(opts.Ctx)

return api
}

Expand DownExpand Up@@ -254,6 +310,69 @@ func (a *API) agent(ctx context.Context) (database.WorkspaceAgent, error) {
return agent, nil
}

func (a *API) workspace() (database.Workspace) {
a.mu.Lock()
defer a.mu.Unlock()

return database.Workspace{
ID: a.cachedWorkspaceFields.ID,
OwnerID: a.cachedWorkspaceFields.OwnerID,
OrganizationID: a.cachedWorkspaceFields.OrganizationID,
TemplateID: a.cachedWorkspaceFields.TemplateID,
Name: a.cachedWorkspaceFields.Name,
OwnerUsername: a.cachedWorkspaceFields.OwnerUsername,
TemplateName: a.cachedWorkspaceFields.TemplateName,
AutostartSchedule: a.cachedWorkspaceFields.AutostartSchedule,
DormantAt: a.cachedWorkspaceFields.DormantAt,
}
}

func (a *API) rbacContext(ctx context.Context) context.Context {
workspace := a.workspace()
return dbauthz.WithWorkspaceRBAC(ctx, workspace.RBACObject())
}

// refreshCachedWorkspace periodically updates the cached workspace fields.
// This ensures that changes like prebuild claims (which modify owner_id, name, etc.)
// are eventually reflected in the cache without requiring agent reconnection.
func (a *API) refreshCachedWorkspace(ctx context.Context) {
a.mu.Lock()
defer a.mu.Unlock()

ws, err := a.opts.Database.GetWorkspaceByID(ctx, a.cachedWorkspaceFields.ID)
if err != nil {
a.opts.Log.Warn(ctx, "failed to refresh cached workspace fields", slog.Error(err))
return
}

// Update fields that can change during workspace lifecycle (e.g., prebuild claim)
a.cachedWorkspaceFields.OwnerID = ws.OwnerID
a.cachedWorkspaceFields.Name = ws.Name
a.cachedWorkspaceFields.OwnerUsername = ws.OwnerUsername
a.cachedWorkspaceFields.AutostartSchedule = ws.AutostartSchedule
a.cachedWorkspaceFields.DormantAt = ws.DormantAt

a.opts.Log.Debug(ctx, "refreshed cached workspace fields",
slog.F("workspace_id", ws.ID),
slog.F("owner_id", ws.OwnerID),
slog.F("name", ws.Name))
}

// startCacheRefreshLoop runs a background goroutine that periodically refreshes
// the cached workspace fields. This is primarily needed to handle prebuild claims
// where the owner_id and other fields change while the agent connection persists.
func (a *API) startCacheRefreshLoop(ctx context.Context) {
// Refresh every 5 minutes. This provides a reasonable balance between:
// - Keeping cache fresh for prebuild claims and other workspace updates
// - Minimizing unnecessary database queries
_ = a.opts.Clock.TickerFunc(ctx, workspaceCacheRefreshInterval, func() error {
a.refreshCachedWorkspace(ctx)
return nil
}, "cache_refresh")

<-ctx.Done()
}

func (a *API) publishWorkspaceUpdate(ctx context.Context, agent *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
a.opts.PublishWorkspaceUpdateFn(ctx, a.opts.OwnerID, wspubsub.WorkspaceEvent{
Kind: kind,
Expand Down
12 changes: 8 additions & 4 deletionscoderd/agentapi/metadata.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -17,10 +17,11 @@ import (
)

type MetadataAPI struct {
AgentFn func(context.Context) (database.WorkspaceAgent, error)
Database database.Store
Pubsub pubsub.Pubsub
Log slog.Logger
AgentFn func(context.Context) (database.WorkspaceAgent, error)
RBACContextFn func(context.Context) context.Context
Database database.Store
Pubsub pubsub.Pubsub
Log slog.Logger

TimeNowFn func() time.Time // defaults to dbtime.Now()
}
Expand DownExpand Up@@ -107,6 +108,9 @@ func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.B
)
}

// Inject RBAC object into context for dbauthz fast path, avoid having to
// call GetWorkspaceByAgentID on every metadata update.
ctx = a.RBACContextFn(ctx)
err = a.Database.UpdateWorkspaceAgentMetadata(ctx, dbUpdate)
if err != nil {
return nil, xerrors.Errorf("update workspace agent metadata in database: %w", err)
Expand Down
Loading
Loading
[8]ページ先頭
©2009-2025 Movatter.jp