run both queries and diff results

coder · johnstcn · Jul 9, 2025 · Jul 2, 2025 · Jul 2, 2025 · Jul 2, 2025
commit 53c3ba5a03a9e6b859d9c4039672f05d65db318a
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -2560,6 +2560,14 @@ func (q *querier) GetRunningPrebuiltWorkspaces(ctx context.Context) ([]database.
 	return q.db.GetRunningPrebuiltWorkspaces(ctx)
 }
 
+func (q *querier) GetRunningPrebuiltWorkspacesOptimized(ctx context.Context) ([]database.GetRunningPrebuiltWorkspacesOptimizedRow, error) {
+	// This query returns only prebuilt workspaces, but we decided to require permissions for all workspaces.
+	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceWorkspace.All()); err != nil {
+		return nil, err
+	}
+	return q.db.GetRunningPrebuiltWorkspacesOptimized(ctx)
+}
+
 func (q *querier) GetRuntimeConfig(ctx context.Context, key string) (string, error) {
 	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
 		return "", err

diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -5101,6 +5101,10 @@ func (q *FakeQuerier) GetRunningPrebuiltWorkspaces(ctx context.Context) ([]datab
 	return nil, ErrUnimplemented
 }
 
+func (q *FakeQuerier) GetRunningPrebuiltWorkspacesOptimized(ctx context.Context) ([]database.GetRunningPrebuiltWorkspacesOptimizedRow, error) {
+	panic("not implemented")
+}
+
 func (q *FakeQuerier) GetRuntimeConfig(_ context.Context, key string) (string, error) {
 	q.mutex.Lock()
 	defer q.mutex.Unlock()

diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
diff --git a/coderd/database/querier.go b/coderd/database/querier.go
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
@@ -48,7 +48,7 @@ WHERE tvp.desired_instances IS NOT NULL -- Consider only presets that have a pre
   -- 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.
 	AND (t.id = sqlc.narg('template_id')::uuid OR sqlc.narg('template_id') IS NULL);
 
--- name: GetRunningPrebuiltWorkspaces :many
+-- name: GetRunningPrebuiltWorkspacesOptimized :many
 WITH latest_prebuilds AS (
 	SELECT
 		workspaces.id,
@@ -96,6 +96,23 @@ SELECT
 FROM latest_prebuilds
 LEFT JOIN ready_agents ON ready_agents.job_id = latest_prebuilds.job_id
 LEFT JOIN workspace_latest_presets ON workspace_latest_presets.workspace_id = latest_prebuilds.id
+ORDER BY latest_prebuilds.id
+;
+
+-- name: GetRunningPrebuiltWorkspaces :many
+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)
+ORDER BY p.id;
 ;
 
 -- name: CountInProgressPrebuilds :many

diff --git a/enterprise/coderd/prebuilds/reconcile.go b/enterprise/coderd/prebuilds/reconcile.go
@@ -12,6 +12,8 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/google/go-cmp/cmp"
+
 	"github.com/hashicorp/go-multierror"
 	"github.com/prometheus/client_golang/prometheus"
 
@@ -398,11 +400,21 @@ func (c *StoreReconciler) SnapshotState(ctx context.Context, store database.Stor
 			return xerrors.Errorf("failed to get preset prebuild schedules: %w", err)
 		}
 
+		// Get results from both original and optimized queries for comparison
 		allRunningPrebuilds, err := db.GetRunningPrebuiltWorkspaces(ctx)
 		if err != nil {
 			return xerrors.Errorf("failed to get running prebuilds: %w", err)
 		}
 
+		// Compare with optimized query to ensure behavioral correctness
+		optimized, err := db.GetRunningPrebuiltWorkspacesOptimized(ctx)
+		if err != nil {
+			// 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)
+		}
+
 		allPrebuildsInProgress, err := db.CountInProgressPrebuilds(ctx)
 		if err != nil {
 			return xerrors.Errorf("failed to get prebuilds in progress: %w", err)
@@ -922,3 +934,29 @@ func SetPrebuildsReconciliationPaused(ctx context.Context, db database.Store, pa
 	}
 	return db.UpsertPrebuildsSettings(ctx, string(settingsJSON))
 }
+
+// CompareGetRunningPrebuiltWorkspacesResults compares the original and optimized
+// query results and logs any differences found. This function can be easily
+// removed once we're confident the optimized query works correctly.
+func CompareGetRunningPrebuiltWorkspacesResults(
+	ctx context.Context,
+	logger slog.Logger,
+	original []database.GetRunningPrebuiltWorkspacesRow,
+	optimized []database.GetRunningPrebuiltWorkspacesOptimizedRow,
+) {
+	// Convert optimized results to the same type as original for comparison
+	var optimizedConverted []database.GetRunningPrebuiltWorkspacesRow
+	if original != nil {
+		optimizedConverted := make([]database.GetRunningPrebuiltWorkspacesRow, len(optimized))
+		for i, row := range optimized {
+			optimizedConverted[i] = database.GetRunningPrebuiltWorkspacesRow(row)
+		}
+	}
+
+	// Compare the results and log an error if they differ.
+	// NOTE: explicitly not sorting here as both query results are ordered by ID.
+	if diff := cmp.Diff(original, optimizedConverted); diff != "" {
+		logger.Error(ctx, "results differ for GetRunningPrebuiltWorkspacesOptimized",
+			slog.F("diff", diff))
+	}
+}