perf(coderd): reduce GetWorkspaceBuildByJobID calls in ExtractWorkspaceAgentParam

Callum Styan · Callum Styan · commit 73627a092b5b · 2025-12-16T20:36:45.000Z
The ExtractWorkspaceAgentParam middleware was making 3 calls to
GetWorkspaceBuildByJobID per request:
1. GetWorkspaceResourceByID → GetProvisionerJobByID cascade
2. Explicit GetProvisionerJobByID call
3. Explicit GetWorkspaceBuildByJobID call

This optimization:
- Adds GetWorkspaceResourceWithJobByIDNoAuth that fetches resource+job
  in a single query while still performing proper authorization
- Updates ExtractWorkspaceAgentParam middleware to use it
- Reduces GetWorkspaceBuildByJobID calls from 3 to 1

This affects all endpoints using this middleware, including:
- GET /workspaceagents/{agent}/listening-ports (the primary target)
- GET /workspaceagents/{agent}/
- GET /workspaceagents/{agent}/logs
- And other agent-scoped endpoints
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -3942,6 +3942,27 @@ func (q *querier) GetWorkspaceResourceWithJobByID(ctx context.Context, id uuid.U
 	return q.db.GetWorkspaceResourceWithJobByID(ctx, id)
 }
 
+// GetWorkspaceResourceWithJobByIDNoAuth is an optimized version that fetches
+// both the resource and job information in a single query, but still performs
+// proper authorization checks. This reduces duplicate GetWorkspaceBuildByJobID
+// calls compared to calling GetWorkspaceResourceByID + GetProvisionerJobByID.
+func (q *querier) GetWorkspaceResourceWithJobByIDNoAuth(ctx context.Context, id uuid.UUID) (database.GetWorkspaceResourceWithJobByIDRow, error) {
+	// Fetch the combined resource and job data.
+	row, err := q.db.GetWorkspaceResourceWithJobByID(ctx, id)
+	if err != nil {
+		return database.GetWorkspaceResourceWithJobByIDRow{}, err
+	}
+
+	// Still need to authorize via the job (which will cascade to workspace build).
+	// This triggers one GetWorkspaceBuildByJobID call instead of two.
+	_, err = q.GetProvisionerJobByID(ctx, row.JobID)
+	if err != nil {
+		return database.GetWorkspaceResourceWithJobByIDRow{}, err
+	}
+
+	return row, nil
+}
+
 // GetWorkspaceResourceMetadataByResourceIDs is only used for build data.
 // The workspace/job is already fetched.
 func (q *querier) GetWorkspaceResourceMetadataByResourceIDs(ctx context.Context, ids []uuid.UUID) ([]database.WorkspaceResourceMetadatum, error) {
diff --git a/coderd/httpmw/workspaceagentparam.go b/coderd/httpmw/workspaceagentparam.go
@@ -50,7 +50,9 @@ func ExtractWorkspaceAgentParam(db database.Store) func(http.Handler) http.Handl
 				return
 			}
 
-			resource, err := db.GetWorkspaceResourceByID(ctx, agent.ResourceID)
+			// Use the optimized combined query to fetch both resource and job info.
+			// This reduces GetWorkspaceBuildByJobID calls from 3 to 1.
+			resourceWithJob, err := db.GetWorkspaceResourceWithJobByIDNoAuth(ctx, agent.ResourceID)
 			if err != nil {
 				httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
 					Message: "Internal error fetching workspace resource.",
@@ -59,21 +61,15 @@ func ExtractWorkspaceAgentParam(db database.Store) func(http.Handler) http.Handl
 				return
 			}
 
-			job, err := db.GetProvisionerJobByID(ctx, resource.JobID)
-			if err != nil {
-				httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
-					Message: "Internal error fetching provisioner job.",
-					Detail:  err.Error(),
-				})
-				return
-			}
-			if job.Type != database.ProvisionerJobTypeWorkspaceBuild {
+			if resourceWithJob.JobType != database.ProvisionerJobTypeWorkspaceBuild {
 				httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
 					Message: "Workspace agents can only be fetched for builds.",
 				})
 				return
 			}
-			build, err := db.GetWorkspaceBuildByJobID(ctx, job.ID)
+
+			// We still need to get the build to extract the workspace ID.
+			build, err := db.GetWorkspaceBuildByJobID(ctx, resourceWithJob.JobID)
 			if err != nil {
 				httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
 					Message: "Internal error fetching workspace build.",
@@ -87,7 +83,7 @@ func ExtractWorkspaceAgentParam(db database.Store) func(http.Handler) http.Handl
 
 			if rlogger := loggermw.RequestLoggerFromContext(ctx); rlogger != nil {
 				rlogger.WithFields(
-					slog.F("workspace_name", resource.Name),
+					slog.F("workspace_name", resourceWithJob.Name),
 					slog.F("agent_name", agent.Name),
 				)
 			}
