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/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)
+	})
+}