From 3a09e365ef7b2b889660993bcbadf6bcaff3b653 Mon Sep 17 00:00:00 2001 From: Callum Styan Date: Tue, 16 Dec 2025 20:28:22 +0000 Subject: [PATCH] perf(agentapi): skip GetWorkspaceAgentByID when RBAC object is cached MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Optimize BatchUpdateMetadata to avoid an unnecessary database call by skipping the GetWorkspaceAgentByID query when the workspace RBAC object is already cached in context. Previously, even with the RBAC fast path working, we would: 1. Inject workspace RBAC object into context 2. Call AgentFn() which calls GetWorkspaceAgentByID (1 DB call for auth) 3. Call UpdateWorkspaceAgentMetadata (uses RBAC from context) Now with the optimization: 1. Inject workspace RBAC object into context 2. Skip AgentFn() and use the already-known agent ID directly 3. Call UpdateWorkspaceAgentMetadata (uses RBAC from context) This eliminates one GetWorkspaceAgentByID call per metadata update. The AgentFn() call was only needed for authorization, but since we already have the RBAC object in context and we know the agent ID from a.opts.AgentID, we can skip the database fetch entirely. The slow path (when RBAC object is not cached) still calls AgentFn() to fetch the agent and perform authorization the traditional way. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Sonnet 4.5 --- coderd/agentapi/api.go | 1 + coderd/agentapi/metadata.go | 25 ++++++++++++++++++++----- 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/coderd/agentapi/api.go b/coderd/agentapi/api.go index 59ce37884440c..20f47d411a515 100644 --- a/coderd/agentapi/api.go +++ b/coderd/agentapi/api.go @@ -174,6 +174,7 @@ func New(opts Options, workspace database.Workspace) *API { } api.MetadataAPI = &MetadataAPI{ + AgentID: opts.AgentID, AgentFn: api.agent, Workspace: api.cachedWorkspaceFields, Database: opts.Database, diff --git a/coderd/agentapi/metadata.go b/coderd/agentapi/metadata.go index 1d4e23ab88932..3bc7a1114006f 100644 --- a/coderd/agentapi/metadata.go +++ b/coderd/agentapi/metadata.go @@ -18,6 +18,7 @@ import ( ) type MetadataAPI struct { + AgentID uuid.UUID AgentFn func(context.Context) (database.WorkspaceAgent, error) Workspace *CachedWorkspaceFields Database database.Store @@ -51,25 +52,39 @@ func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.B // call GetWorkspaceByAgentID on every metadata update. var err error rbacCtx := ctx + hasRBACObject := false if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok { rbacCtx, err = dbauthz.WithWorkspaceRBAC(ctx, dbws.RBACObject()) if err != nil { // Don't error level log here, will exit the function. We want to fall back to GetWorkspaceByAgentID. //nolint:gocritic a.Log.Debug(ctx, "Cached workspace was present but RBAC object was invalid", slog.F("err", err)) + } else { + hasRBACObject = true } } - workspaceAgent, err := a.AgentFn(rbacCtx) - if err != nil { - return nil, err + // Fast path: If we have the RBAC object in context, we can skip fetching + // the agent from the database since we only need the agent ID, which we + // already have. The RBAC object will be used for authorization in the + // UpdateWorkspaceAgentMetadata call. + var agentID uuid.UUID + if hasRBACObject { + agentID = a.AgentID + } else { + // Slow path: Fetch the agent from the database for authorization. + workspaceAgent, err := a.AgentFn(rbacCtx) + if err != nil { + return nil, err + } + agentID = workspaceAgent.ID } var ( collectedAt = a.now() allKeysLen = 0 dbUpdate = database.UpdateWorkspaceAgentMetadataParams{ - WorkspaceAgentID: workspaceAgent.ID, + WorkspaceAgentID: agentID, // These need to be `make(x, 0, len(req.Metadata))` instead of // `make(x, len(req.Metadata))` because we may not insert all // metadata if the keys are large. @@ -134,7 +149,7 @@ func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.B if err != nil { return nil, xerrors.Errorf("marshal workspace agent metadata channel payload: %w", err) } - err = a.Pubsub.Publish(WatchWorkspaceAgentMetadataChannel(workspaceAgent.ID), payload) + err = a.Pubsub.Publish(WatchWorkspaceAgentMetadataChannel(agentID), payload) if err != nil { return nil, xerrors.Errorf("publish workspace agent metadata: %w", err) }