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/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,