perf: reduce calls to GetWorkspaceByAgentID in GetWorkspaceAgentByID (#21046)

This PR piggy backs on the agent API cached workspace added in an earlier PR to provide a fast path for avoiding `GetWorkspaceByAgentID` calls in dbauthz's `GetWorkspaceAgentByID`. This query is not the most expensive, but has a significant call volume at ~16 million calls per week.

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

coderd/agentapi/api.go

@@ -197,6 +197,7 @@ func New(opts Options, workspace database.Workspace) *API {
 		AgentFn:          api.agent,
 		ConnectionLogger: opts.ConnectionLogger,
 		Database:         opts.Database,
+		Workspace:        api.cachedWorkspaceFields,
 		Log:              opts.Log,
 	}
 

coderd/agentapi/connectionlog.go

@@ -14,11 +14,13 @@ import (
 	"github.com/coder/coder/v2/coderd/connectionlog"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/db2sdk"
+	"github.com/coder/coder/v2/coderd/database/dbauthz"
 )
 
 type ConnLogAPI struct {
 	AgentFn          func(context.Context) (database.WorkspaceAgent, error)
 	ConnectionLogger *atomic.Pointer[connectionlog.ConnectionLogger]
+	Workspace        *CachedWorkspaceFields
 	Database         database.Store
 	Log              slog.Logger
 }
@@ -51,14 +53,31 @@ func (a *ConnLogAPI) ReportConnection(ctx context.Context, req *agentproto.Repor
 		}
 	}
 
+	// Inject RBAC object into context for dbauthz fast path, avoid having to
+	// call GetWorkspaceByAgentID on every metadata update.
+	rbacCtx := ctx
+	var ws database.WorkspaceIdentity
+	if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok {
+		ws = dbws
+		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))
+		}
+	}
+
 	// Fetch contextual data for this connection log event.
-	workspaceAgent, err := a.AgentFn(ctx)
+	workspaceAgent, err := a.AgentFn(rbacCtx)
 	if err != nil {
 		return nil, xerrors.Errorf("get agent: %w", err)
 	}
-	workspace, err := a.Database.GetWorkspaceByAgentID(ctx, workspaceAgent.ID)
-	if err != nil {
-		return nil, xerrors.Errorf("get workspace by agent id: %w", err)
+	if ws.Equal(database.WorkspaceIdentity{}) {
+		workspace, err := a.Database.GetWorkspaceByAgentID(ctx, workspaceAgent.ID)
+		if err != nil {
+			return nil, xerrors.Errorf("get workspace by agent id: %w", err)
+		}
+		ws = database.WorkspaceIdentityFromWorkspace(workspace)
 	}
 
 	// Some older clients may incorrectly report "localhost" as the IP address.
@@ -74,10 +93,10 @@ func (a *ConnLogAPI) ReportConnection(ctx context.Context, req *agentproto.Repor
 	err = connLogger.Upsert(ctx, database.UpsertConnectionLogParams{
 		ID:               uuid.New(),
 		Time:             req.GetConnection().GetTimestamp().AsTime(),
-		OrganizationID:   workspace.OrganizationID,
-		WorkspaceOwnerID: workspace.OwnerID,
-		WorkspaceID:      workspace.ID,
-		WorkspaceName:    workspace.Name,
+		OrganizationID:   ws.OrganizationID,
+		WorkspaceOwnerID: ws.OwnerID,
+		WorkspaceID:      ws.ID,
+		WorkspaceName:    ws.Name,
 		AgentName:        workspaceAgent.Name,
 		Type:             connectionType,
 		Code:             code,

coderd/agentapi/connectionlog_test.go

@@ -117,6 +117,7 @@ func TestConnectionLog(t *testing.T) {
 				AgentFn: func(context.Context) (database.WorkspaceAgent, error) {
 					return agent, nil
 				},
+				Workspace: &agentapi.CachedWorkspaceFields{},
 			}
 			api.ReportConnection(context.Background(), &agentproto.ReportConnectionRequest{
 				Connection: &agentproto.Connection{

coderd/agentapi/metadata.go

@@ -47,7 +47,20 @@ func (a *MetadataAPI) BatchUpdateMetadata(ctx context.Context, req *agentproto.B
 		maxErrorLen = maxValueLen
 	)
 
-	workspaceAgent, err := a.AgentFn(ctx)
+	// Inject RBAC object into context for dbauthz fast path, avoid having to
+	// call GetWorkspaceByAgentID on every metadata update.
+	var err error
+	rbacCtx := ctx
+	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))
+		}
+	}
+
+	workspaceAgent, err := a.AgentFn(rbacCtx)
 	if err != nil {
 		return nil, err
 	}
@@ -109,18 +122,6 @@ 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.
-	rbacCtx := ctx
-	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))
-		}
-	}
-
 	err = a.Database.UpdateWorkspaceAgentMetadata(rbacCtx, dbUpdate)
 	if err != nil {
 		return nil, xerrors.Errorf("update workspace agent metadata in database: %w", err)

coderd/agentapi/metadata_test.go

@@ -295,6 +295,7 @@ func TestBatchUpdateMetadata(t *testing.T) {
 			now  = dbtime.Now()
 			// Set up consistent IDs that represent a valid workspace->agent relationship
 			workspaceID = uuid.MustParse("12345678-1234-1234-1234-123456789012")
+			templateID  = uuid.MustParse("aaaabbbb-cccc-dddd-eeee-ffffffff0000")
 			ownerID     = uuid.MustParse("87654321-4321-4321-4321-210987654321")
 			orgID       = uuid.MustParse("11111111-1111-1111-1111-111111111111")
 			agentID     = uuid.MustParse("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa")
@@ -358,8 +359,48 @@ func TestBatchUpdateMetadata(t *testing.T) {
 			OrganizationID: orgID,
 		})
 
-		// Create context with system actor so authorization passes
-		ctx := dbauthz.AsSystemRestricted(context.Background())
+		// Create roles with workspace permissions
+		userRoles := rbac.Roles([]rbac.Role{
+			{
+				Identifier: rbac.RoleMember(),
+				User: []rbac.Permission{
+					{
+						Negate:       false,
+						ResourceType: rbac.ResourceWorkspace.Type,
+						Action:       policy.WildcardSymbol,
+					},
+				},
+				ByOrgID: map[string]rbac.OrgPermissions{
+					orgID.String(): {
+						Member: []rbac.Permission{
+							{
+								Negate:       false,
+								ResourceType: rbac.ResourceWorkspace.Type,
+								Action:       policy.WildcardSymbol,
+							},
+						},
+					},
+				},
+			},
+		})
+
+		agentScope := rbac.WorkspaceAgentScope(rbac.WorkspaceAgentScopeParams{
+			WorkspaceID: workspaceID,
+			OwnerID:     ownerID,
+			TemplateID:  templateID,
+			VersionID:   uuid.New(),
+		})
+
+		ctx := dbauthz.As(context.Background(), rbac.Subject{
+			Type:         rbac.SubjectTypeUser,
+			FriendlyName: "testuser",
+			Email:        "testuser@example.com",
+			ID:           ownerID.String(),
+			Roles:        userRoles,
+			Groups:       []string{orgID.String()},
+			Scope:        agentScope,
+		}.WithCachedASTValue())
+
 		resp, err := api.BatchUpdateMetadata(ctx, req)
 		require.NoError(t, err)
 		require.NotNil(t, resp)
@@ -376,6 +417,7 @@ func TestBatchUpdateMetadata(t *testing.T) {
 			pub         = &fakePublisher{}
 			now         = dbtime.Now()
 			workspaceID = uuid.MustParse("12345678-1234-1234-1234-123456789012")
+			templateID  = uuid.MustParse("aaaabbbb-cccc-dddd-eeee-ffffffff0000")
 			ownerID     = uuid.MustParse("87654321-4321-4321-4321-210987654321")
 			orgID       = uuid.MustParse("11111111-1111-1111-1111-111111111111")
 			agentID     = uuid.MustParse("bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb")
@@ -445,12 +487,53 @@ func TestBatchUpdateMetadata(t *testing.T) {
 			OrganizationID: uuid.Nil, // Invalid: fails dbauthz fast path validation
 		})
 
-		// Create context with system actor so authorization passes
-		ctx := dbauthz.AsSystemRestricted(context.Background())
+		// Create roles with workspace permissions
+		userRoles := rbac.Roles([]rbac.Role{
+			{
+				Identifier: rbac.RoleMember(),
+				User: []rbac.Permission{
+					{
+						Negate:       false,
+						ResourceType: rbac.ResourceWorkspace.Type,
+						Action:       policy.WildcardSymbol,
+					},
+				},
+				ByOrgID: map[string]rbac.OrgPermissions{
+					orgID.String(): {
+						Member: []rbac.Permission{
+							{
+								Negate:       false,
+								ResourceType: rbac.ResourceWorkspace.Type,
+								Action:       policy.WildcardSymbol,
+							},
+						},
+					},
+				},
+			},
+		})
+
+		agentScope := rbac.WorkspaceAgentScope(rbac.WorkspaceAgentScopeParams{
+			WorkspaceID: workspaceID,
+			OwnerID:     ownerID,
+			TemplateID:  templateID,
+			VersionID:   uuid.New(),
+		})
+
+		ctx := dbauthz.As(context.Background(), rbac.Subject{
+			Type:         rbac.SubjectTypeUser,
+			FriendlyName: "testuser",
+			Email:        "testuser@example.com",
+			ID:           ownerID.String(),
+			Roles:        userRoles,
+			Groups:       []string{orgID.String()},
+			Scope:        agentScope,
+		}.WithCachedASTValue())
+
 		resp, err := api.BatchUpdateMetadata(ctx, req)
 		require.NoError(t, err)
 		require.NotNil(t, resp)
 	})
+
 	// Test RBAC slow path - no RBAC object in context
 	// This test verifies that when no RBAC object is present in context, the dbauthz layer
 	// falls back to the slow path and calls GetWorkspaceByAgentID.
@@ -463,6 +546,7 @@ func TestBatchUpdateMetadata(t *testing.T) {
 			pub         = &fakePublisher{}
 			now         = dbtime.Now()
 			workspaceID = uuid.MustParse("12345678-1234-1234-1234-123456789012")
+			templateID  = uuid.MustParse("aaaabbbb-cccc-dddd-eeee-ffffffff0000")
 			ownerID     = uuid.MustParse("87654321-4321-4321-4321-210987654321")
 			orgID       = uuid.MustParse("11111111-1111-1111-1111-111111111111")
 			agentID     = uuid.MustParse("dddddddd-dddd-dddd-dddd-dddddddddddd")
@@ -523,8 +607,48 @@ func TestBatchUpdateMetadata(t *testing.T) {
 			},
 		}
 
-		// Create context with system actor so authorization passes
-		ctx := dbauthz.AsSystemRestricted(context.Background())
+		// Create roles with workspace permissions
+		userRoles := rbac.Roles([]rbac.Role{
+			{
+				Identifier: rbac.RoleMember(),
+				User: []rbac.Permission{
+					{
+						Negate:       false,
+						ResourceType: rbac.ResourceWorkspace.Type,
+						Action:       policy.WildcardSymbol,
+					},
+				},
+				ByOrgID: map[string]rbac.OrgPermissions{
+					orgID.String(): {
+						Member: []rbac.Permission{
+							{
+								Negate:       false,
+								ResourceType: rbac.ResourceWorkspace.Type,
+								Action:       policy.WildcardSymbol,
+							},
+						},
+					},
+				},
+			},
+		})
+
+		agentScope := rbac.WorkspaceAgentScope(rbac.WorkspaceAgentScopeParams{
+			WorkspaceID: workspaceID,
+			OwnerID:     ownerID,
+			TemplateID:  templateID,
+			VersionID:   uuid.New(),
+		})
+
+		ctx := dbauthz.As(context.Background(), rbac.Subject{
+			Type:         rbac.SubjectTypeUser,
+			FriendlyName: "testuser",
+			Email:        "testuser@example.com",
+			ID:           ownerID.String(),
+			Roles:        userRoles,
+			Groups:       []string{orgID.String()},
+			Scope:        agentScope,
+		}.WithCachedASTValue())
+
 		resp, err := api.BatchUpdateMetadata(ctx, req)
 		require.NoError(t, err)
 		require.NotNil(t, resp)

coderd/agentapi/stats.go

@@ -10,6 +10,7 @@ import (
 	"cdr.dev/slog"
 	agentproto "github.com/coder/coder/v2/agent/proto"
 	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbauthz"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/workspacestats"
 	"github.com/coder/coder/v2/codersdk"
@@ -43,7 +44,21 @@ func (a *StatsAPI) UpdateStats(ctx context.Context, req *agentproto.UpdateStatsR
 		return res, nil
 	}
 
-	workspaceAgent, err := a.AgentFn(ctx)
+	// Inject RBAC object into context for dbauthz fast path, avoid having to
+	// call GetWorkspaceAgentByID on every stats update.
+
+	rbacCtx := ctx
+	if dbws, ok := a.Workspace.AsWorkspaceIdentity(); ok {
+		var err error
+		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))
+		}
+	}
+
+	workspaceAgent, err := a.AgentFn(rbacCtx)
 	if err != nil {
 		return nil, err
 	}

coderd/database/dbauthz/dbauthz.go

@@ -3560,6 +3560,21 @@ func (q *querier) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context
 }
 
 func (q *querier) GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (database.WorkspaceAgent, error) {
+	// Fast path: Check if we have a workspace RBAC object in context.
+	// In the agent API this is set at agent connection time to avoid the expensive
+	// GetWorkspaceByAgentID query for every agent operation.
+	// 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 GetWorkspaceByAgentID,
+		// in case the cached data is stale.
+		if err := q.authorizeContext(ctx, policy.ActionRead, rbacObj); err == nil {
+			return q.db.GetWorkspaceAgentByID(ctx, id)
+		}
+		q.log.Debug(ctx, "fast path authorization failed for GetWorkspaceAgentByID, using slow path",
+			slog.F("agent_id", id))
+	}
+
+	// Slow path: Fallback to fetching the workspace for authorization
 	if _, err := q.GetWorkspaceByAgentID(ctx, id); err != nil {
 		return database.WorkspaceAgent{}, err
 	}

coderd/database/dbauthz/dbauthz_test.go

@@ -4805,3 +4805,81 @@ func TestGetLatestWorkspaceBuildByWorkspaceID_FastPath(t *testing.T) {
 		require.Equal(t, build, result)
 	})
 }
+
+func TestGetWorkspaceAgentByID_FastPath(t *testing.T) {
+	t.Parallel()
+
+	agentID := uuid.New()
+	ownerID := uuid.New()
+	wsID := uuid.New()
+	orgID := uuid.New()
+
+	agent := database.WorkspaceAgent{
+		ID:   agentID,
+		Name: "test-agent",
+	}
+
+	workspace := database.Workspace{
+		ID:             wsID,
+		OwnerID:        ownerID,
+		OrganizationID: orgID,
+	}
+
+	wsIdentity := database.WorkspaceIdentity{
+		ID:             wsID,
+		OwnerID:        ownerID,
+		OrganizationID: orgID,
+	}
+
+	actor := rbac.Subject{
+		ID:     ownerID.String(),
+		Roles:  rbac.RoleIdentifiers{rbac.RoleOwner()},
+		Groups: []string{orgID.String()},
+		Scope:  rbac.ScopeAll,
+	}
+
+	authorizer := &coderdtest.RecordingAuthorizer{
+		Wrapped: (&coderdtest.FakeAuthorizer{}).AlwaysReturn(nil),
+	}
+
+	t.Run("WithWorkspaceRBAC", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := dbauthz.As(context.Background(), actor)
+		ctrl := gomock.NewController(t)
+		mockDB := dbmock.NewMockStore(ctrl)
+
+		rbacObj := wsIdentity.RBACObject()
+		ctx, err := dbauthz.WithWorkspaceRBAC(ctx, rbacObj)
+		require.NoError(t, err)
+
+		mockDB.EXPECT().Wrappers().Return([]string{})
+		// GetWorkspaceByAgentID should NOT be called
+		mockDB.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID).Return(agent, nil)
+
+		q := dbauthz.New(mockDB, authorizer, slogtest.Make(t, nil), coderdtest.AccessControlStorePointer())
+
+		result, err := q.GetWorkspaceAgentByID(ctx, agentID)
+		require.NoError(t, err)
+		require.Equal(t, agent, result)
+	})
+
+	t.Run("WithoutWorkspaceRBAC", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := dbauthz.As(context.Background(), actor)
+		ctrl := gomock.NewController(t)
+		mockDB := dbmock.NewMockStore(ctrl)
+
+		mockDB.EXPECT().Wrappers().Return([]string{})
+		// GetWorkspaceByAgentID SHOULD be called
+		mockDB.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agentID).Return(workspace, nil)
+		mockDB.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID).Return(agent, nil)
+
+		q := dbauthz.New(mockDB, authorizer, slogtest.Make(t, nil), coderdtest.AccessControlStorePointer())
+
+		result, err := q.GetWorkspaceAgentByID(ctx, agentID)
+		require.NoError(t, err)
+		require.Equal(t, agent, result)
+	})
+}