chore(coderd/database): optimize GetRunningPrebuiltWorkspaces (#18588)

johnstcn · web-flow · commit 258a839d27d9 · 2025-07-01T09:42:01.000+01:00
Fixes coder/internal#715 After this change, the only use of the `workspace_prebuilds` view is the `ClaimPrebuiltWorkspace` query. A subsequent PR will update the view. Before: ~44ms https://explain.dalibo.com/plan/76cbe21d1a4c9329#plan After: 7.3ms https://explain.dalibo.com/plan/5abbdf926315677e#plan
diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go
@@ -3758,9 +3758,9 @@ func createPrebuiltWorkspace(
 	job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
 		Type:           database.ProvisionerJobTypeWorkspaceBuild,
 		OrganizationID: orgID,
-
-		CreatedAt: now.Add(-1 * time.Minute),
-		Error:     jobError,
+		CreatedAt:      now.Add(-1 * time.Minute),
+		CompletedAt:    sql.NullTime{Time: now, Valid: true},
+		Error:          jobError,
 	})
 
 	// create ready agents
@@ -3930,6 +3930,95 @@ func TestWorkspacePrebuildsView(t *testing.T) {
 	}
 }
 
+func TestGetRunningPrebuiltWorkspaces(t *testing.T) {
+	t.Parallel()
+	if !dbtestutil.WillUsePostgres() {
+		t.SkipNow()
+	}
+
+	now := dbtime.Now()
+	orgID := uuid.New()
+	userID := uuid.New()
+
+	testCases := []struct {
+		name           string
+		readyAgents    int
+		notReadyAgents int
+		expectRows     int
+		expectReady    bool
+	}{
+		{
+			name:           "one ready agent",
+			readyAgents:    1,
+			notReadyAgents: 0,
+			expectRows:     1,
+			expectReady:    true,
+		},
+		{
+			name:           "one not ready agent",
+			readyAgents:    0,
+			notReadyAgents: 1,
+			expectRows:     1,
+			expectReady:    false,
+		},
+		{
+			name:           "one ready, one not ready",
+			readyAgents:    1,
+			notReadyAgents: 1,
+			expectRows:     1,
+			expectReady:    false,
+		},
+		{
+			name:           "both ready",
+			readyAgents:    2,
+			notReadyAgents: 0,
+			expectRows:     1,
+			expectReady:    true,
+		},
+		{
+			name:           "five ready, one not ready",
+			readyAgents:    5,
+			notReadyAgents: 1,
+			expectRows:     1,
+			expectReady:    false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			sqlDB := testSQLDB(t)
+			err := migrations.Up(sqlDB)
+			require.NoError(t, err)
+			db := database.New(sqlDB)
+
+			ctx := testutil.Context(t, testutil.WaitShort)
+
+			dbgen.Organization(t, db, database.Organization{
+				ID: orgID,
+			})
+			dbgen.User(t, db, database.User{
+				ID: userID,
+			})
+
+			tmpl := createTemplate(t, db, orgID, userID)
+			tmplV1 := createTmplVersionAndPreset(t, db, tmpl, tmpl.ActiveVersionID, now, nil)
+			createPrebuiltWorkspace(ctx, t, db, tmpl, tmplV1, orgID, now, &createPrebuiltWorkspaceOpts{
+				readyAgents:    tc.readyAgents,
+				notReadyAgents: tc.notReadyAgents,
+			})
+
+			workspacePrebuilds, err := db.GetRunningPrebuiltWorkspaces(ctx)
+			require.NoError(t, err)
+			require.Len(t, workspacePrebuilds, tc.expectRows)
+			if tc.expectRows > 0 {
+				require.Equal(t, tc.expectReady, workspacePrebuilds[0].Ready)
+			}
+		})
+	}
+}
+
 func TestGetPresetsBackoff(t *testing.T) {
 	t.Parallel()
 	if !dbtestutil.WillUsePostgres() {
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/prebuilds.sql b/coderd/database/queries/prebuilds.sql
@@ -49,18 +49,55 @@ WHERE tvp.desired_instances IS NOT NULL -- Consider only presets that have a pre
 	AND (t.id = sqlc.narg('template_id')::uuid OR sqlc.narg('template_id') IS NULL);
 
 -- name: GetRunningPrebuiltWorkspaces :many
+WITH latest_prebuilds AS (
+	SELECT
+		latest_build.workspace_id,
+		workspaces.name,
+		workspaces.template_id,
+		latest_build.template_version_id,
+		latest_build.template_version_preset_id,
+		latest_build.job_id,
+		workspaces.created_at
+	FROM workspaces
+	JOIN LATERAL (
+		SELECT
+			workspace_builds.workspace_id,
+			workspace_builds.template_version_id,
+			workspace_builds.job_id,
+			workspace_builds.template_version_preset_id
+		FROM workspace_builds
+		JOIN provisioner_jobs ON provisioner_jobs.id = workspace_builds.job_id
+		WHERE workspace_builds.workspace_id = workspaces.id
+		AND workspace_builds.transition = 'start'::workspace_transition
+		AND provisioner_jobs.job_status = 'succeeded'::provisioner_job_status
+		ORDER BY workspace_builds.build_number DESC
+		LIMIT 1
+	) AS latest_build ON true
+	WHERE workspaces.deleted = false
+	AND workspaces.owner_id = 'c42fdf75-3097-471c-8c33-fb52454d81c0'::UUID
+),
+ready_agents AS (
+	SELECT
+		latest_prebuilds.job_id,
+		BOOL_AND(workspace_agents.lifecycle_state = 'ready'::workspace_agent_lifecycle_state)::boolean AS ready
+	FROM latest_prebuilds
+	JOIN workspace_resources ON workspace_resources.job_id = latest_prebuilds.job_id
+	JOIN workspace_agents ON workspace_agents.resource_id = workspace_resources.id
+	WHERE workspace_agents.deleted = false
+	AND workspace_agents.parent_id IS NULL
+	GROUP BY latest_prebuilds.job_id
+)
 SELECT
-		p.id,
-		p.name,
-		p.template_id,
-		b.template_version_id,
-		p.current_preset_id AS current_preset_id,
-		p.ready,
-		p.created_at
-FROM workspace_prebuilds p
-		INNER JOIN workspace_latest_builds b ON b.workspace_id = p.id
-WHERE (b.transition = 'start'::workspace_transition
-	AND b.job_status = 'succeeded'::provisioner_job_status);
+	latest_prebuilds.workspace_id AS id,
+	latest_prebuilds.name,
+	latest_prebuilds.template_id,
+	latest_prebuilds.template_version_id,
+	latest_prebuilds.template_version_preset_id AS current_preset_id,
+	COALESCE(ready_agents.ready, false)::boolean AS ready,
+	latest_prebuilds.created_at
+FROM latest_prebuilds
+LEFT JOIN ready_agents ON ready_agents.job_id = latest_prebuilds.job_id
+ORDER BY latest_prebuilds.workspace_id ASC;
 
 -- name: CountInProgressPrebuilds :many
 -- CountInProgressPrebuilds returns the number of in-progress prebuilds, grouped by preset ID and transition.
