perf: reduce DB calls to GetWorkspaceByAgentID via caching workspace info (#20662)

---------

Signed-off-by: Callum Styan <callumstyan@gmail.com>

coderd/database/dbauthz/dbauthz.go

@@ -5556,6 +5556,22 @@ func (q *querier) UpdateWorkspaceAgentLogOverflowByID(ctx context.Context, arg d
 }
 
 func (q *querier) UpdateWorkspaceAgentMetadata(ctx context.Context, arg database.UpdateWorkspaceAgentMetadataParams) error {
+	// Fast path: Check if we have an RBAC object in context.
+	// This is set by the workspace agent RPC handler to avoid the expensive
+	// GetWorkspaceByAgentID query for every metadata update.
+	// NOTE: The cached RBAC object is refreshed every 5 minutes in agentapi/api.go.
+	if rbacObj, ok := WorkspaceRBACFromContext(ctx); ok {
+		// Errors here will result in falling back to the GetWorkspaceAgentByID query, skipping
+		// the cache in case the cached data is stale.
+		if err := q.authorizeContext(ctx, policy.ActionUpdate, rbacObj); err == nil {
+			return q.db.UpdateWorkspaceAgentMetadata(ctx, arg)
+		}
+		q.log.Debug(ctx, "fast path authorization failed, using slow path",
+			slog.F("agent_id", arg.WorkspaceAgentID))
+	}
+
+	// Slow path: Fallback to fetching the workspace for authorization if the RBAC object is not present (or is invalid)
+	// in the request context.
 	workspace, err := q.db.GetWorkspaceByAgentID(ctx, arg.WorkspaceAgentID)
 	if err != nil {
 		return err

coderd/database/dbauthz/workspace_rbac_context.go

@@ -0,0 +1,41 @@
+package dbauthz
+
+import (
+	"context"
+
+	"github.com/google/uuid"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/v2/coderd/rbac"
+)
+
+func isWorkspaceRBACObjectEmpty(rbacObj rbac.Object) bool {
+	// if any of these are true then the rbac.Object work a workspace is considered empty
+	return rbacObj.Owner == "" || rbacObj.OrgID == "" || rbacObj.Owner == uuid.Nil.String() || rbacObj.OrgID == uuid.Nil.String()
+}
+
+type workspaceRBACContextKey struct{}
+
+// WithWorkspaceRBAC attaches a workspace RBAC object to the context.
+// RBAC fields on this RBAC object should not be used.
+//
+// This is primarily used by the workspace agent RPC handler to cache workspace
+// authorization data for the duration of an agent connection.
+func WithWorkspaceRBAC(ctx context.Context, rbacObj rbac.Object) (context.Context, error) {
+	if rbacObj.Type != rbac.ResourceWorkspace.Type {
+		return ctx, xerrors.New("RBAC Object must be of type Workspace")
+	}
+	if isWorkspaceRBACObjectEmpty(rbacObj) {
+		return ctx, xerrors.Errorf("cannot attach empty RBAC object to context: %+v", rbacObj)
+	}
+	if len(rbacObj.ACLGroupList) != 0 || len(rbacObj.ACLUserList) != 0 {
+		return ctx, xerrors.New("ACL fields for Workspace RBAC object must be nullified, the can be changed during runtime and should not be cached")
+	}
+	return context.WithValue(ctx, workspaceRBACContextKey{}, rbacObj), nil
+}
+
+// WorkspaceRBACFromContext attempts to retrieve the workspace RBAC object from context.
+func WorkspaceRBACFromContext(ctx context.Context) (rbac.Object, bool) {
+	obj, ok := ctx.Value(workspaceRBACContextKey{}).(rbac.Object)
+	return obj, ok
+}

coderd/database/modelmethods.go

@@ -1,6 +1,7 @@
 package database
 
 import (
+	"database/sql"
 	"encoding/hex"
 	"slices"
 	"sort"
@@ -796,3 +797,60 @@ func (s UserSecret) RBACObject() rbac.Object {
 func (s AIBridgeInterception) RBACObject() rbac.Object {
 	return rbac.ResourceAibridgeInterception.WithOwner(s.InitiatorID.String())
 }
+
+// WorkspaceIdentity contains the minimal workspace fields needed for agent API metadata/stats reporting
+// and RBAC checks, without requiring a full database.Workspace object.
+type WorkspaceIdentity struct {
+	// Add any other fields needed for IsPrebuild() if it relies on workspace fields
+	// 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
+}
+
+func (w WorkspaceIdentity) RBACObject() rbac.Object {
+	return Workspace{
+		ID:                w.ID,
+		OwnerID:           w.OwnerID,
+		OrganizationID:    w.OrganizationID,
+		TemplateID:        w.TemplateID,
+		Name:              w.Name,
+		OwnerUsername:     w.OwnerUsername,
+		TemplateName:      w.TemplateName,
+		AutostartSchedule: w.AutostartSchedule,
+	}.RBACObject()
+}
+
+// IsPrebuild returns true if the workspace is a prebuild workspace.
+// A workspace is considered a prebuild if its owner is the prebuild system user.
+func (w WorkspaceIdentity) IsPrebuild() bool {
+	return w.OwnerID == PrebuildsSystemUserID
+}
+
+func (w WorkspaceIdentity) Equal(w2 WorkspaceIdentity) bool {
+	return w.ID == w2.ID && w.OwnerID == w2.OwnerID && w.OrganizationID == w2.OrganizationID &&
+		w.TemplateID == w2.TemplateID && w.Name == w2.Name && w.OwnerUsername == w2.OwnerUsername &&
+		w.TemplateName == w2.TemplateName && w.AutostartSchedule == w2.AutostartSchedule
+}
+
+func WorkspaceIdentityFromWorkspace(w Workspace) WorkspaceIdentity {
+	return WorkspaceIdentity{
+		ID:                w.ID,
+		OwnerID:           w.OwnerID,
+		OrganizationID:    w.OrganizationID,
+		TemplateID:        w.TemplateID,
+		Name:              w.Name,
+		OwnerUsername:     w.OwnerUsername,
+		TemplateName:      w.TemplateName,
+		AutostartSchedule: w.AutostartSchedule,
+	}
+}