chore: optimize GetPrebuiltWorkspaces query #18717

johnstcn · 2025-07-02T14:04:24Z

Hopefully fixes coder/internal#715 this time.

Second attempt at this.
Previous attempt incorrectly returned all rows for which there existed a prebuild that had previously had a successful start transition at some point: #18588

Explain before (71.1ms): https://explain.dalibo.com/plan/9be18ab833b7a000
~~Explain after (9.8ms): https://explain.dalibo.com/plan/b4a94742gaha229g~~ (EDIT: not selecting from correct CTE)
Explain after (11.2ms): https://explain.dalibo.com/plan/bea42b563ff7fbe7

Manually verified against dogfood db:

$ psql '<db url>' -f a.sql > a.txt
$ psql '<db url>' -f b.sql > b.txt
$ diff a.txt b.txt
28d27
<  18aa3c11-c1eb-4fc6-ae3b-8cbd01f81f21 | prebuild-hiyghekkb3gv2vq | 99064381-8750-407f-9d30-91f38b7911cc | 5d91f37e-5872-4821-b6ed-90574082628e | bdc27adb-7a8e-4a6f-9922-5a0f4a3885da | f     | 2025-07-02 13:34:58.537656+00
32a32
>  18aa3c11-c1eb-4fc6-ae3b-8cbd01f81f21 | prebuild-hiyghekkb3gv2vq | 99064381-8750-407f-9d30-91f38b7911cc | 5d91f37e-5872-4821-b6ed-90574082628e | bdc27adb-7a8e-4a6f-9922-5a0f4a3885da | f     | 2025-07-02 13:34:58.537656+00
$ diff <(sort a.txt) <(sort b.txt)
$

Diff is only due to unstable row order, also added additional testing in coderd/database/querier_test.go to validate changes.

~~If we want to be super careful about this, I can instead break out the updated query into a new function and diff the old versus the new on each reconcile call.~~
EDIT: decided to go ahead with this for safety.

johnstcn · 2025-07-03T13:01:28Z

coderd/database/queries/prebuilds.sql

@@ -60,7 +111,9 @@ SELECT
 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);
+	AND b.job_status = 'succeeded'::provisioner_job_status)
+ORDER BY p.id;


review: adding stable ordering for diffing

…Workspaces

Copilot

Pull Request Overview

This PR introduces and validates an optimized version of the GetRunningPrebuiltWorkspaces query, compares its output against the original for correctness, and wires it through the reconciler, mocks, metrics, and DB layers.

Add a new optimized SQL query and corresponding Go method (GetRunningPrebuiltWorkspacesOptimized)
Integrate a comparator in the reconciler to log differences between original and optimized results
Expand tests and mocks to cover the new optimized method

Reviewed Changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
enterprise/coderd/prebuilds/reconcile_test.go	Added antagonists setup and `CompareGetRunningPrebuiltWorkspacesResults` tests
enterprise/coderd/prebuilds/reconcile.go	Invoke optimized query, compare results, and log diffs
coderd/database/queries/prebuilds.sql	Defined `GetRunningPrebuiltWorkspacesOptimized` CTE and ordered original query
coderd/database/queries.sql.go	Added constant, row type, and method for optimized query
coderd/database/querier.go	Extended interface with optimized method
coderd/database/dbmock/dbmock.go	Mock recorder for optimized method
coderd/database/dbmetrics/querymetrics.go	Metrics wrapper for optimized method
coderd/database/dbmem/dbmem.go	Stub (panicking) implementation for optimized method
coderd/database/dbauthz/setup_test.go	Skipped optimized method in authz tests
coderd/database/dbauthz/dbauthz_test.go	Excluded optimized method from recursive authz test
coderd/database/dbauthz/dbauthz.go	Added authorization wrapper for optimized method

Comments suppressed due to low confidence (1)

coderd/database/querier_test.go:5025

Consider adding a parallel test for 'GetRunningPrebuiltWorkspacesOptimized' to verify the optimized SQL returns the same results as the original under PostgreSQL.

// TestGetRunningPrebuiltWorkspaces ensures the correct behavior of the

coderd/database/dbmem/dbmem.go

enterprise/coderd/prebuilds/reconcile.go

coderd/database/queries/prebuilds.sql

enterprise/coderd/prebuilds/reconcile.go

mafredri

Looks mostly alright, mainly just questioning the purpose/point of the workspace_latest_presets CTE.

mafredri · 2025-07-03T17:18:36Z

coderd/database/queries/prebuilds.sql

+	WHERE workspace_builds.transition = 'start'::workspace_transition
+	AND   workspace_builds.template_version_preset_id IS NOT NULL
+	ORDER BY latest_prebuilds.id, workspace_builds.build_number DESC
+),


I can't quite grasp the purpose of this CTE, it seems to pretty much do the same as latest_prebuilds? It's selecting distinct on workspace id from latest_prebuilds and joining workspace_builds on workspace id (we already have this data in the previous CTE, no?), and ordering by build number DESC, combined with distinct on that's == selecting directly from workspace_latest_builds, right? If we can't select the template_version_preset_id directly from the workspace_latest_builds then it'd still be more performant to join on build ID vs workspace ID and avoid t he distinct/order/where transition clause.

Might make more sense to just join workspace_builds in latest_prebuilds instead to get template_version_preset_id if required and remove this CTE.

When I omitted this CTE it caused different results for one specific prebuild instance, but I'm now not seeing the same discrepancy when I run the queries side-by-side. I'll try the modifications you suggest and see if it still yields the same results.

The more I look at that CTE, the more I question it too. On the dogfood database, all of the workspace builds owned by the prebuild user have a non-null template_version_preset_id, whereas almost all subsequent ones have a null preset ID.

AFAIU, whenever a template version is added, new presets are inserted into the database. I don't think we ever delete presets from the database. Except maybe when a template is deleted?
A prebuild should always correlate to a non-null preset. Could this be an issue when the active template version changes? 🤔

After spending some more time looking at this, it appears that the current correct approach is to query the most recent non-null preset ID for a successful start transition. Simply selecting from workspace_latest_builds may not result in the correct preset_id. I discussed with Sas and a future change may result in preset_id not being a thing any more.

Prebuilt workspaces should ideally never change their preset. This means that the first build should contain a preset id and all subsequent builds should have a null preset ID. We could optimise this by dropping the workspace_latest_presets CTE from the query and instead rely on the preset ID for build 0 in all cases.

This would probably be much more performant.

The issue with this optimisation is the potentially negligent edge case where a human being restarts a prebuilt workspace with a new preset. In the optimized case such a restarted prebuilt workspace would be claimed with the wrong preset.

If we're willing to tolerate this edge case, we can optimize by dropping the CTE. Otherwise, the CTE remains necessary in order to find the last set preset id for a prebuilt workspace.

I'll leave the CTE for now. There is an argument to be made that once a human touches a prebuild that it's no longer valid to be claimed and as such only prebuilds with build_number = 1 should be claimable.

But what do we say to the God of Arguments? Not today.

Prebuilt workspaces should ideally never change their preset. This means that the first build should contain a preset id and all subsequent builds should have a null preset ID.

Why do we nullify the preset_id in subsequent builds? Conceptually, a prebuilt workspace is always associated with a preset, and clearing the field might imply the association no longer exists, which isn't quite accurate. Wouldn’t it make more sense to retain the preset_id to reflect this persistent relationship?

The issue with this optimisation is the potentially negligent edge case where a human being restarts a prebuilt workspace with a new preset.

How would this actually happen? Is it really possible for a prebuilt workspace to be restarted with a different preset? If we consider the association to the preset immutable once the prebuilt workspace is created, this edge case shouldn’t occur, right? Or is there a specific workflow that allows this to happen?

Either way, this is not a blocker for the PR, we can definitely move forward and revisit this discussion later if needed.

nullify the preset_id in subsequent builds

We don't explicitly nullify them. The fact that they are null on subsequent build is rather because there's nothing to tell a subsequent build what the preset_id should be. We could transfer the preset id from one build to the next, but that would roughly look like running the CTE that we're discussing here on every new build. Presets are set initially and when they change by specification of the user. For normal workspaces this user is a human being who chooses a preset. For prebuilt workspaces this user is the prebuilds system user and they specify the preset in the reconciliation loop.

How would this actually happen? Is it really possible for a prebuilt workspace to be restarted with a different preset?

Any workspace, prebuilt or not can be restarted with new parameters as long as those parameters are mutable. The same goes for presets. It is only possible via API right now. Neither the UI nor the CLI support changing a preset.

This means it is theoretically possible for a user with sufficient privileges to send an API call to restart with a new preset, but it is definitely not supported.

mafredri · 2025-07-03T17:22:10Z

enterprise/coderd/prebuilds/reconcile.go

+			// Log the error but continue with original results
+			c.logger.Error(ctx, "optimized GetRunningPrebuiltWorkspacesOptimized failed", slog.Error(err))
+		} else {
+			CompareGetRunningPrebuiltWorkspacesResults(ctx, c.logger, allRunningPrebuilds, optimized)


This is a fun way to test it out 😄

enterprise/coderd/prebuilds/reconcile_test.go

coderd/database/querier_test.go

coderd/database/queries/prebuilds.sql

ssncferreira · 2025-07-07T18:41:28Z

coderd/database/queries/prebuilds.sql

+	WHERE workspace_builds.transition = 'start'::workspace_transition
+	AND   workspace_builds.template_version_preset_id IS NOT NULL
+	ORDER BY latest_prebuilds.id, workspace_builds.build_number DESC
+),


AFAIU, whenever a template version is added, new presets are inserted into the database. I don't think we ever delete presets from the database. Except maybe when a template is deleted?
A prebuild should always correlate to a non-null preset. Could this be an issue when the active template version changes? 🤔

enterprise/coderd/prebuilds/reconcile.go

SasSwart

You're going to have to remove the dbmem bit (🥳), but otherwise this looks good. See my comment about potentially further optimising your query at the cost of some robustness. That's non-blocking though. Since we're just logging the diffs for now.

…ebuiltWorkspace-Optimize-2

ssncferreira

LGTM, let's try it 🤞

ssncferreira · 2025-07-09T10:14:53Z

coderd/database/querier_test.go

+	_ = setupFixture(t, db, "stopped-prebuild", false, database.WorkspaceTransitionStop, database.ProvisionerJobStatusSucceeded)
+	_ = setupFixture(t, db, "failed-prebuild", false, database.WorkspaceTransitionStart, database.ProvisionerJobStatusFailed)
+	_ = setupFixture(t, db, "canceled-prebuild", false, database.WorkspaceTransitionStart, database.ProvisionerJobStatusCanceled)
+	_ = setupFixture(t, db, "deleted-prebuild", true, database.WorkspaceTransitionStart, database.ProvisionerJobStatusSucceeded)


ssncferreira · 2025-07-09T10:28:25Z

coderd/database/queries/prebuilds.sql

+	WHERE workspace_builds.transition = 'start'::workspace_transition
+	AND   workspace_builds.template_version_preset_id IS NOT NULL
+	ORDER BY latest_prebuilds.id, workspace_builds.build_number DESC
+),


Prebuilt workspaces should ideally never change their preset. This means that the first build should contain a preset id and all subsequent builds should have a null preset ID.

Why do we nullify the preset_id in subsequent builds? Conceptually, a prebuilt workspace is always associated with a preset, and clearing the field might imply the association no longer exists, which isn't quite accurate. Wouldn’t it make more sense to retain the preset_id to reflect this persistent relationship?

The issue with this optimisation is the potentially negligent edge case where a human being restarts a prebuilt workspace with a new preset.

How would this actually happen? Is it really possible for a prebuilt workspace to be restarted with a different preset? If we consider the association to the preset immutable once the prebuilt workspace is created, this edge case shouldn’t occur, right? Or is there a specific workflow that allows this to happen?

Either way, this is not a blocker for the PR, we can definitely move forward and revisit this discussion later if needed.

johnstcn self-assigned this Jul 2, 2025

johnstcn marked this pull request as ready for review July 2, 2025 14:20

johnstcn requested review from mafredri, ssncferreira and SasSwart July 2, 2025 14:21

johnstcn commented Jul 3, 2025

View reviewed changes

johnstcn added 5 commits July 3, 2025 14:05

chore: add test to ensure correct functionality of GetRunningPrebuilt…

2c48c95

…Workspaces

chore: add test antagonists to enterprise/prebuild/reconcile tests

68a4304

chore: alternate optimization for GetRunningPrebuiltWorkspaces

c95b6d0

chore: get latest preset id

6b4d6dc

run both queries and diff results

53c3ba5

johnstcn force-pushed the cj/chore/GetRunningPrebuiltWorkspace-Optimize-2 branch from da75520 to 53c3ba5 Compare July 3, 2025 13:08

johnstcn added 4 commits July 3, 2025 14:31

fixup! run both queries and diff results

b237251

fixup! run both queries and diff results

5dc69d9

fixup! run both queries and diff results

5158046

fixup! run both queries and diff results

3e69a1b

johnstcn requested a review from Copilot July 3, 2025 16:06

Copilot AI reviewed Jul 3, 2025

View reviewed changes

address copilot comments

2e760c7

mafredri reviewed Jul 3, 2025

View reviewed changes

use sink instead of spy

a0f4f36

ssncferreira reviewed Jul 7, 2025

View reviewed changes

johnstcn added 3 commits July 8, 2025 12:28

add comments to new query CTEs

30c86cc

improve querier test

1d6b4c0

add TODO

6a1c635

SasSwart approved these changes Jul 9, 2025

View reviewed changes

Merge remote-tracking branch 'origin/main' into cj/chore/GetRunningPr…

e43b5ba

…ebuiltWorkspace-Optimize-2

ssncferreira approved these changes Jul 9, 2025

View reviewed changes

johnstcn merged commit 0367dba into main Jul 9, 2025
30 checks passed

johnstcn deleted the cj/chore/GetRunningPrebuiltWorkspace-Optimize-2 branch July 9, 2025 10:30

github-actions bot locked and limited conversation to collaborators Jul 9, 2025

chore: optimize GetPrebuiltWorkspaces query #18717

chore: optimize GetPrebuiltWorkspaces query #18717

Uh oh!

Conversation

johnstcn commented Jul 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mafredri left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

johnstcn Jul 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SasSwart Jul 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

SasSwart left a comment

Choose a reason for hiding this comment

Uh oh!

ssncferreira left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

johnstcn commented Jul 2, 2025 •

edited

Loading

mafredri left a comment •

edited

Loading

johnstcn Jul 8, 2025 •

edited

Loading

SasSwart Jul 9, 2025 •

edited

Loading