Skip to content

Commit 0367dba

Browse files
johnstcnjohnstcn
authored
chore: optimize GetPrebuiltWorkspaces query (#18717)
* Adds GetRunningPrebuiltWorkspacesOptimized query * Runs both original and updated query side-by-side and logs diffs
1 parent dc0919d commit 0367dba

File tree

12 files changed

+615
-2
lines changed

12 files changed

+615
-2
lines changed

‎coderd/database/dbauthz/dbauthz.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2596,6 +2596,14 @@ func (q *querier) GetRunningPrebuiltWorkspaces(ctx context.Context) ([]database.
25962596
return q.db.GetRunningPrebuiltWorkspaces(ctx)
25972597
}
25982598

2599+
func (q *querier) GetRunningPrebuiltWorkspacesOptimized(ctx context.Context) ([]database.GetRunningPrebuiltWorkspacesOptimizedRow, error) {
2600+
// This query returns only prebuilt workspaces, but we decided to require permissions for all workspaces.
2601+
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceWorkspace.All()); err != nil {
2602+
return nil, err
2603+
}
2604+
return q.db.GetRunningPrebuiltWorkspacesOptimized(ctx)
2605+
}
2606+
25992607
func (q *querier) GetRuntimeConfig(ctx context.Context, key string) (string, error) {
26002608
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
26012609
return "", err

‎coderd/database/dbauthz/dbauthz_test.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -178,7 +178,8 @@ func TestDBAuthzRecursive(t *testing.T) {
178178
if method.Name == "InTx" ||
179179
method.Name == "Ping" ||
180180
method.Name == "Wrappers" ||
181-
method.Name == "PGLocks" {
181+
method.Name == "PGLocks" ||
182+
method.Name == "GetRunningPrebuiltWorkspacesOptimized" {
182183
continue
183184
}
184185
// easy to know which method failed.

‎coderd/database/dbauthz/setup_test.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@ var skipMethods = map[string]string{
4141
"Wrappers": "Not relevant",
4242
"AcquireLock": "Not relevant",
4343
"TryAcquireLock": "Not relevant",
44+
// This method will be removed once we know this works correctly.
45+
"GetRunningPrebuiltWorkspacesOptimized": "Not relevant",
4446
}
4547

4648
// TestMethodTestSuite runs MethodTestSuite.

‎coderd/database/dbgen/dbgen.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -359,6 +359,14 @@ func Workspace(t testing.TB, db database.Store, orig database.WorkspaceTable) da
359359
NextStartAt: orig.NextStartAt,
360360
})
361361
require.NoError(t, err, "insert workspace")
362+
if orig.Deleted {
363+
err = db.UpdateWorkspaceDeletedByID(genCtx, database.UpdateWorkspaceDeletedByIDParams{
364+
ID: workspace.ID,
365+
Deleted: true,
366+
})
367+
require.NoError(t, err, "set workspace as deleted")
368+
workspace.Deleted = true
369+
}
362370
return workspace
363371
}
364372

‎coderd/database/dbmetrics/querymetrics.go

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dbmock/dbmock.go

Lines changed: 15 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/querier.go

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/querier_test.go

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5020,3 +5020,102 @@ func requireUsersMatch(t testing.TB, expected []database.User, found []database.
50205020
t.Helper()
50215021
require.ElementsMatch(t, expected, database.ConvertUserRows(found), msg)
50225022
}
5023+
5024+
// TestGetRunningPrebuiltWorkspaces ensures the correct behavior of the
5025+
// GetRunningPrebuiltWorkspaces query.
5026+
func TestGetRunningPrebuiltWorkspaces(t *testing.T) {
5027+
t.Parallel()
5028+
5029+
if !dbtestutil.WillUsePostgres() {
5030+
t.Skip("Test requires PostgreSQL for complex queries")
5031+
}
5032+
5033+
ctx := testutil.Context(t, testutil.WaitLong)
5034+
db, _ := dbtestutil.NewDB(t)
5035+
now := dbtime.Now()
5036+
5037+
// Given: a prebuilt workspace with a successful start build and a stop build.
5038+
org := dbgen.Organization(t, db, database.Organization{})
5039+
user := dbgen.User(t, db, database.User{})
5040+
template := dbgen.Template(t, db, database.Template{
5041+
CreatedBy: user.ID,
5042+
OrganizationID: org.ID,
5043+
})
5044+
templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{
5045+
TemplateID: uuid.NullUUID{UUID: template.ID, Valid: true},
5046+
OrganizationID: org.ID,
5047+
CreatedBy: user.ID,
5048+
})
5049+
preset := dbgen.Preset(t, db, database.InsertPresetParams{
5050+
TemplateVersionID: templateVersion.ID,
5051+
DesiredInstances: sql.NullInt32{Int32: 1, Valid: true},
5052+
})
5053+
5054+
setupFixture := func(t *testing.T, db database.Store, name string, deleted bool, transition database.WorkspaceTransition, jobStatus database.ProvisionerJobStatus) database.WorkspaceTable {
5055+
t.Helper()
5056+
ws := dbgen.Workspace(t, db, database.WorkspaceTable{
5057+
OwnerID: database.PrebuildsSystemUserID,
5058+
TemplateID: template.ID,
5059+
Name: name,
5060+
Deleted: deleted,
5061+
})
5062+
var canceledAt sql.NullTime
5063+
var jobError sql.NullString
5064+
switch jobStatus {
5065+
case database.ProvisionerJobStatusFailed:
5066+
jobError = sql.NullString{String: assert.AnError.Error(), Valid: true}
5067+
case database.ProvisionerJobStatusCanceled:
5068+
canceledAt = sql.NullTime{Time: now, Valid: true}
5069+
}
5070+
pj := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
5071+
OrganizationID: org.ID,
5072+
InitiatorID: database.PrebuildsSystemUserID,
5073+
Provisioner: database.ProvisionerTypeEcho,
5074+
Type: database.ProvisionerJobTypeWorkspaceBuild,
5075+
StartedAt: sql.NullTime{Time: now.Add(-time.Minute), Valid: true},
5076+
CanceledAt: canceledAt,
5077+
CompletedAt: sql.NullTime{Time: now, Valid: true},
5078+
Error: jobError,
5079+
})
5080+
wb := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
5081+
WorkspaceID: ws.ID,
5082+
TemplateVersionID: templateVersion.ID,
5083+
TemplateVersionPresetID: uuid.NullUUID{UUID: preset.ID, Valid: true},
5084+
JobID: pj.ID,
5085+
BuildNumber: 1,
5086+
Transition: transition,
5087+
InitiatorID: database.PrebuildsSystemUserID,
5088+
Reason: database.BuildReasonInitiator,
5089+
})
5090+
// Ensure things are set up as expectd
5091+
require.Equal(t, transition, wb.Transition)
5092+
require.Equal(t, int32(1), wb.BuildNumber)
5093+
require.Equal(t, jobStatus, pj.JobStatus)
5094+
require.Equal(t, deleted, ws.Deleted)
5095+
5096+
return ws
5097+
}
5098+
5099+
// Given: a number of prebuild workspaces with different states exist.
5100+
runningPrebuild := setupFixture(t, db, "running-prebuild", false, database.WorkspaceTransitionStart, database.ProvisionerJobStatusSucceeded)
5101+
_ = setupFixture(t, db, "stopped-prebuild", false, database.WorkspaceTransitionStop, database.ProvisionerJobStatusSucceeded)
5102+
_ = setupFixture(t, db, "failed-prebuild", false, database.WorkspaceTransitionStart, database.ProvisionerJobStatusFailed)
5103+
_ = setupFixture(t, db, "canceled-prebuild", false, database.WorkspaceTransitionStart, database.ProvisionerJobStatusCanceled)
5104+
_ = setupFixture(t, db, "deleted-prebuild", true, database.WorkspaceTransitionStart, database.ProvisionerJobStatusSucceeded)
5105+
5106+
// Given: a regular workspace also exists.
5107+
_ = dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
5108+
OwnerID: user.ID,
5109+
TemplateID: template.ID,
5110+
Name: "test-running-regular-workspace",
5111+
Deleted: false,
5112+
})
5113+
5114+
// When: we query for running prebuild workspaces
5115+
runningPrebuilds, err := db.GetRunningPrebuiltWorkspaces(ctx)
5116+
require.NoError(t, err)
5117+
5118+
// Then: only the running prebuild workspace should be returned.
5119+
require.Len(t, runningPrebuilds, 1, "expected only one running prebuilt workspace")
5120+
require.Equal(t, runningPrebuild.ID, runningPrebuilds[0].ID, "expected the running prebuilt workspace to be returned")
5121+
}

‎coderd/database/queries.sql.go

Lines changed: 101 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/prebuilds.sql

Lines changed: 60 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,64 @@ WHERE tvp.desired_instances IS NOT NULL -- Consider only presets that have a pre
4848
-- AND NOT t.deleted -- We don't exclude deleted templates because there's no constraint in the DB preventing a soft deletion on a template while workspaces are running.
4949
AND (t.id = sqlc.narg('template_id')::uuid OR sqlc.narg('template_id') IS NULL);
5050

51+
-- name: GetRunningPrebuiltWorkspacesOptimized :many
52+
WITH latest_prebuilds AS (
53+
-- All workspaces that match the following criteria:
54+
-- 1. Owned by prebuilds user
55+
-- 2. Not deleted
56+
-- 3. Latest build is a 'start' transition
57+
-- 4. Latest build was successful
58+
SELECT
59+
workspaces.id,
60+
workspaces.name,
61+
workspaces.template_id,
62+
workspace_latest_builds.template_version_id,
63+
workspace_latest_builds.job_id,
64+
workspaces.created_at
65+
FROM workspace_latest_builds
66+
JOIN workspaces ON workspaces.id = workspace_latest_builds.workspace_id
67+
WHERE workspace_latest_builds.transition = 'start'::workspace_transition
68+
AND workspace_latest_builds.job_status = 'succeeded'::provisioner_job_status
69+
AND workspaces.owner_id = 'c42fdf75-3097-471c-8c33-fb52454d81c0'::UUID
70+
AND NOT workspaces.deleted
71+
),
72+
workspace_latest_presets AS (
73+
-- For each of the above workspaces, the preset_id of the most recent
74+
-- successful start transition.
75+
SELECT DISTINCT ON (latest_prebuilds.id)
76+
latest_prebuilds.id AS workspace_id,
77+
workspace_builds.template_version_preset_id AS current_preset_id
78+
FROM latest_prebuilds
79+
JOIN workspace_builds ON workspace_builds.workspace_id = latest_prebuilds.id
80+
WHERE workspace_builds.transition = 'start'::workspace_transition
81+
AND workspace_builds.template_version_preset_id IS NOT NULL
82+
ORDER BY latest_prebuilds.id, workspace_builds.build_number DESC
83+
),
84+
ready_agents AS (
85+
-- For each of the above workspaces, check if all agents are ready.
86+
SELECT
87+
latest_prebuilds.job_id,
88+
BOOL_AND(workspace_agents.lifecycle_state = 'ready'::workspace_agent_lifecycle_state)::boolean AS ready
89+
FROM latest_prebuilds
90+
JOIN workspace_resources ON workspace_resources.job_id = latest_prebuilds.job_id
91+
JOIN workspace_agents ON workspace_agents.resource_id = workspace_resources.id
92+
WHERE workspace_agents.deleted = false
93+
AND workspace_agents.parent_id IS NULL
94+
GROUP BY latest_prebuilds.job_id
95+
)
96+
SELECT
97+
latest_prebuilds.id,
98+
latest_prebuilds.name,
99+
latest_prebuilds.template_id,
100+
latest_prebuilds.template_version_id,
101+
workspace_latest_presets.current_preset_id,
102+
COALESCE(ready_agents.ready, false)::boolean AS ready,
103+
latest_prebuilds.created_at
104+
FROM latest_prebuilds
105+
LEFT JOIN ready_agents ON ready_agents.job_id = latest_prebuilds.job_id
106+
LEFT JOIN workspace_latest_presets ON workspace_latest_presets.workspace_id = latest_prebuilds.id
107+
ORDER BY latest_prebuilds.id;
108+
51109
-- name: GetRunningPrebuiltWorkspaces :many
52110
SELECT
53111
p.id,
@@ -60,7 +118,8 @@ SELECT
60118
FROM workspace_prebuilds p
61119
INNER JOIN workspace_latest_builds b ON b.workspace_id = p.id
62120
WHERE (b.transition = 'start'::workspace_transition
63-
AND b.job_status = 'succeeded'::provisioner_job_status);
121+
AND b.job_status = 'succeeded'::provisioner_job_status)
122+
ORDER BY p.id;
64123

65124
-- name: CountInProgressPrebuilds :many
66125
-- CountInProgressPrebuilds returns the number of in-progress prebuilds, grouped by preset ID and transition.

0 commit comments

Comments
 (0)