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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy