From 258a839d27d92dc78d4414fc67e07eaf23ae01e3 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Tue, 1 Jul 2025 09:42:01 +0100 Subject: [PATCH] chore(coderd/database): optimize GetRunningPrebuiltWorkspaces (#18588) Fixes https://github.com/coder/internal/issues/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 --- coderd/database/querier_test.go | 95 ++++++++++++++++++++++++++- coderd/database/queries.sql.go | 59 +++++++++++++---- coderd/database/queries/prebuilds.sql | 59 +++++++++++++---- 3 files changed, 188 insertions(+), 25 deletions(-) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index f80f68115a..08306c402d 100644 --- 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 index e3e893df0c..f07fddb68d 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -6357,18 +6357,55 @@ func (q *sqlQuerier) GetPresetsBackoff(ctx context.Context, lookback time.Time) } const getRunningPrebuiltWorkspaces = `-- 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 ` type GetRunningPrebuiltWorkspacesRow struct { diff --git a/coderd/database/queries/prebuilds.sql b/coderd/database/queries/prebuilds.sql index 2fc9f3f4a6..8f0e2981b1 100644 --- 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.