chore: optimize GetPrebuiltWorkspaces query (#18717)

* Adds GetRunningPrebuiltWorkspacesOptimized query * Runs both original and updated query side-by-side and logs diffs
2026-09-24 15:04:27 +08:00 · 2025-07-09 11:30:42 +01:00
parent dc0919da33
commit 0367dbac43
12 changed files with 615 additions and 2 deletions
@@ -12,6 +12,7 @@ import (
 	"sync/atomic"
 	"time"

+	"github.com/google/go-cmp/cmp"
 	"github.com/hashicorp/go-multierror"
 	"github.com/prometheus/client_golang/prometheus"

@@ -398,11 +399,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 +933,30 @@ 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.
+// TODO(Cian): Remove this function once the optimized query is stable and correct.
+func CompareGetRunningPrebuiltWorkspacesResults(
+	ctx context.Context,
+	logger slog.Logger,
+	original []database.GetRunningPrebuiltWorkspacesRow,
+	optimized []database.GetRunningPrebuiltWorkspacesOptimizedRow,
+) {
+	if len(original) == 0 && len(optimized) == 0 {
+		return
+	}
+	// Convert optimized results to the same type as original for comparison
+	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))
+	}
+}
@@ -5,6 +5,7 @@ import (
 	"database/sql"
 	"fmt"
 	"sort"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -26,6 +27,7 @@ import (
 	"tailscale.com/types/ptr"

 	"cdr.dev/slog"
+	"cdr.dev/slog/sloggers/slogjson"
 	"cdr.dev/slog/sloggers/slogtest"
 	"github.com/coder/quartz"

@@ -370,6 +372,8 @@ func TestPrebuildReconciliation(t *testing.T) {
 									templateVersionID,
 								)

+								setupTestDBPrebuildAntagonists(t, db, pubSub, org)
+
 								if !templateVersionActive {
 									// Create a new template version and mark it as active
 									// This marks the template version that we care about as inactive
@@ -2116,6 +2120,115 @@ func setupTestDBWorkspaceAgent(t *testing.T, db database.Store, workspaceID uuid
 	return agent
 }

+// setupTestDBAntagonists creates test antagonists that should not influence running prebuild workspace tests.
+//  1. A stopped prebuilt workspace (STOP then START transitions, owned by
+//     prebuilds system user).
+//  2. A running regular workspace (not owned by the prebuilds system user).
+func setupTestDBPrebuildAntagonists(t *testing.T, db database.Store, ps pubsub.Pubsub, org database.Organization) {
+	t.Helper()
+
+	templateAdmin := dbgen.User(t, db, database.User{RBACRoles: []string{codersdk.RoleTemplateAdmin}})
+	_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
+		OrganizationID: org.ID,
+		UserID:         templateAdmin.ID,
+	})
+	member := dbgen.User(t, db, database.User{})
+	_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
+		OrganizationID: org.ID,
+		UserID:         member.ID,
+	})
+	tpl := dbgen.Template(t, db, database.Template{
+		OrganizationID: org.ID,
+		CreatedBy:      templateAdmin.ID,
+	})
+	tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{
+		TemplateID:     uuid.NullUUID{UUID: tpl.ID, Valid: true},
+		OrganizationID: org.ID,
+		CreatedBy:      templateAdmin.ID,
+	})
+
+	// 1) Stopped prebuilt workspace (owned by prebuilds system user)
+	stoppedPrebuild := dbgen.Workspace(t, db, database.WorkspaceTable{
+		OwnerID:    database.PrebuildsSystemUserID,
+		TemplateID: tpl.ID,
+		Name:       "prebuild-antagonist-stopped",
+		Deleted:    false,
+	})
+
+	// STOP build (build number 2, most recent)
+	stoppedJob2 := dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
+		OrganizationID: org.ID,
+		InitiatorID:    database.PrebuildsSystemUserID,
+		Provisioner:    database.ProvisionerTypeEcho,
+		Type:           database.ProvisionerJobTypeWorkspaceBuild,
+		StartedAt:      sql.NullTime{Time: dbtime.Now().Add(-30 * time.Second), Valid: true},
+		CompletedAt:    sql.NullTime{Time: dbtime.Now().Add(-20 * time.Second), Valid: true},
+		Error:          sql.NullString{},
+		ErrorCode:      sql.NullString{},
+	})
+	dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+		WorkspaceID:       stoppedPrebuild.ID,
+		TemplateVersionID: tv.ID,
+		JobID:             stoppedJob2.ID,
+		BuildNumber:       2,
+		Transition:        database.WorkspaceTransitionStop,
+		InitiatorID:       database.PrebuildsSystemUserID,
+		Reason:            database.BuildReasonInitiator,
+		// Explicitly not using a preset here. This shouldn't normally be possible,
+		// but without this the reconciler will try to create a new prebuild for
+		// this preset, which will affect the tests.
+		TemplateVersionPresetID: uuid.NullUUID{},
+	})
+
+	// START build (build number 1, older)
+	stoppedJob1 := dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
+		OrganizationID: org.ID,
+		InitiatorID:    database.PrebuildsSystemUserID,
+		Provisioner:    database.ProvisionerTypeEcho,
+		Type:           database.ProvisionerJobTypeWorkspaceBuild,
+		StartedAt:      sql.NullTime{Time: dbtime.Now().Add(-60 * time.Second), Valid: true},
+		CompletedAt:    sql.NullTime{Time: dbtime.Now().Add(-50 * time.Second), Valid: true},
+		Error:          sql.NullString{},
+		ErrorCode:      sql.NullString{},
+	})
+	dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+		WorkspaceID:       stoppedPrebuild.ID,
+		TemplateVersionID: tv.ID,
+		JobID:             stoppedJob1.ID,
+		BuildNumber:       1,
+		Transition:        database.WorkspaceTransitionStart,
+		InitiatorID:       database.PrebuildsSystemUserID,
+		Reason:            database.BuildReasonInitiator,
+	})
+
+	// 2) Running regular workspace (not owned by prebuilds system user)
+	regularWorkspace := dbgen.Workspace(t, db, database.WorkspaceTable{
+		OwnerID:    member.ID,
+		TemplateID: tpl.ID,
+		Name:       "antagonist-regular-workspace",
+		Deleted:    false,
+	})
+	regularJob := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
+		OrganizationID: org.ID,
+		InitiatorID:    member.ID,
+		Provisioner:    database.ProvisionerTypeEcho,
+		Type:           database.ProvisionerJobTypeWorkspaceBuild,
+		StartedAt:      sql.NullTime{Time: dbtime.Now().Add(-40 * time.Second), Valid: true},
+		CompletedAt:    sql.NullTime{Time: dbtime.Now().Add(-30 * time.Second), Valid: true},
+		Error:          sql.NullString{},
+		ErrorCode:      sql.NullString{},
+	})
+	dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+		WorkspaceID:       regularWorkspace.ID,
+		TemplateVersionID: tv.ID,
+		JobID:             regularJob.ID,
+		BuildNumber:       1,
+		Transition:        database.WorkspaceTransitionStart,
+		InitiatorID:       member.ID,
+		Reason:            database.BuildReasonInitiator,
+	})
+}
+
 var allTransitions = []database.WorkspaceTransition{
 	database.WorkspaceTransitionStart,
 	database.WorkspaceTransitionStop,
@@ -2220,3 +2333,164 @@ func TestReconciliationRespectsPauseSetting(t *testing.T) {
 	require.NoError(t, err)
 	require.Len(t, workspaces, 2, "should have recreated 2 prebuilds after resuming")
 }
+
+func TestCompareGetRunningPrebuiltWorkspacesResults(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+
+	// Helper to create test data
+	createWorkspaceRow := func(id string, name string, ready bool) database.GetRunningPrebuiltWorkspacesRow {
+		uid := uuid.MustParse(id)
+		return database.GetRunningPrebuiltWorkspacesRow{
+			ID:                uid,
+			Name:              name,
+			TemplateID:        uuid.New(),
+			TemplateVersionID: uuid.New(),
+			CurrentPresetID:   uuid.NullUUID{UUID: uuid.New(), Valid: true},
+			Ready:             ready,
+			CreatedAt:         time.Now(),
+		}
+	}
+
+	createOptimizedRow := func(row database.GetRunningPrebuiltWorkspacesRow) database.GetRunningPrebuiltWorkspacesOptimizedRow {
+		return database.GetRunningPrebuiltWorkspacesOptimizedRow(row)
+	}
+
+	t.Run("identical results - no logging", func(t *testing.T) {
+		t.Parallel()
+
+		var sb strings.Builder
+		logger := slog.Make(slogjson.Sink(&sb))
+
+		original := []database.GetRunningPrebuiltWorkspacesRow{
+			createWorkspaceRow("550e8400-e29b-41d4-a716-446655440000", "workspace1", true),
+			createWorkspaceRow("550e8400-e29b-41d4-a716-446655440001", "workspace2", false),
+		}
+
+		optimized := []database.GetRunningPrebuiltWorkspacesOptimizedRow{
+			createOptimizedRow(original[0]),
+			createOptimizedRow(original[1]),
+		}
+
+		prebuilds.CompareGetRunningPrebuiltWorkspacesResults(ctx, logger, original, optimized)
+
+		// Should not log any errors when results are identical
+		require.Empty(t, strings.TrimSpace(sb.String()))
+	})
+
+	t.Run("count mismatch - logs error", func(t *testing.T) {
+		t.Parallel()
+
+		var sb strings.Builder
+		logger := slog.Make(slogjson.Sink(&sb))
+
+		original := []database.GetRunningPrebuiltWorkspacesRow{
+			createWorkspaceRow("550e8400-e29b-41d4-a716-446655440000", "workspace1", true),
+		}
+
+		optimized := []database.GetRunningPrebuiltWorkspacesOptimizedRow{
+			createOptimizedRow(original[0]),
+			createOptimizedRow(createWorkspaceRow("550e8400-e29b-41d4-a716-446655440001", "workspace2", false)),
+		}
+
+		prebuilds.CompareGetRunningPrebuiltWorkspacesResults(ctx, logger, original, optimized)
+
+		// Should log exactly one error.
+		if lines := strings.Split(strings.TrimSpace(sb.String()), "\n"); assert.NotEmpty(t, lines) {
+			require.Len(t, lines, 1)
+			assert.Contains(t, lines[0], "ERROR")
+			assert.Contains(t, lines[0], "workspace2")
+			assert.Contains(t, lines[0], "CurrentPresetID")
+		}
+	})
+
+	t.Run("count mismatch - other direction", func(t *testing.T) {
+		t.Parallel()
+
+		var sb strings.Builder
+		logger := slog.Make(slogjson.Sink(&sb))
+
+		original := []database.GetRunningPrebuiltWorkspacesRow{}
+
+		optimized := []database.GetRunningPrebuiltWorkspacesOptimizedRow{
+			createOptimizedRow(createWorkspaceRow("550e8400-e29b-41d4-a716-446655440001", "workspace2", false)),
+		}
+
+		prebuilds.CompareGetRunningPrebuiltWorkspacesResults(ctx, logger, original, optimized)
+
+		if lines := strings.Split(strings.TrimSpace(sb.String()), "\n"); assert.NotEmpty(t, lines) {
+			require.Len(t, lines, 1)
+			assert.Contains(t, lines[0], "ERROR")
+			assert.Contains(t, lines[0], "workspace2")
+			assert.Contains(t, lines[0], "CurrentPresetID")
+		}
+	})
+
+	t.Run("field differences - logs errors", func(t *testing.T) {
+		t.Parallel()
+
+		var sb strings.Builder
+		logger := slog.Make(slogjson.Sink(&sb))
+
+		workspace1 := createWorkspaceRow("550e8400-e29b-41d4-a716-446655440000", "workspace1", true)
+		workspace2 := createWorkspaceRow("550e8400-e29b-41d4-a716-446655440001", "workspace2", false)
+
+		original := []database.GetRunningPrebuiltWorkspacesRow{workspace1, workspace2}
+
+		// Create optimized with different values
+		optimized1 := createOptimizedRow(workspace1)
+		optimized1.Name = "different-name" // Different name
+		optimized1.Ready = false           // Different ready status
+
+		optimized2 := createOptimizedRow(workspace2)
+		optimized2.CurrentPresetID = uuid.NullUUID{Valid: false} // Different preset ID (NULL)
+
+		optimized := []database.GetRunningPrebuiltWorkspacesOptimizedRow{optimized1, optimized2}
+
+		prebuilds.CompareGetRunningPrebuiltWorkspacesResults(ctx, logger, original, optimized)
+
+		// Should log exactly one error with a cmp.Diff output
+		if lines := strings.Split(strings.TrimSpace(sb.String()), "\n"); assert.NotEmpty(t, lines) {
+			require.Len(t, lines, 1)
+			assert.Contains(t, lines[0], "ERROR")
+			assert.Contains(t, lines[0], "different-name")
+			assert.Contains(t, lines[0], "workspace1")
+			assert.Contains(t, lines[0], "Ready")
+			assert.Contains(t, lines[0], "CurrentPresetID")
+		}
+	})
+
+	t.Run("empty results - no logging", func(t *testing.T) {
+		t.Parallel()
+
+		var sb strings.Builder
+		logger := slog.Make(slogjson.Sink(&sb))
+
+		original := []database.GetRunningPrebuiltWorkspacesRow{}
+		optimized := []database.GetRunningPrebuiltWorkspacesOptimizedRow{}
+
+		prebuilds.CompareGetRunningPrebuiltWorkspacesResults(ctx, logger, original, optimized)
+
+		// Should not log any errors when both results are empty
+		require.Empty(t, strings.TrimSpace(sb.String()))
+	})
+
+	t.Run("nil original", func(t *testing.T) {
+		t.Parallel()
+		var sb strings.Builder
+		logger := slog.Make(slogjson.Sink(&sb))
+		prebuilds.CompareGetRunningPrebuiltWorkspacesResults(ctx, logger, nil, []database.GetRunningPrebuiltWorkspacesOptimizedRow{})
+		// Should not log any errors when original is nil
+		require.Empty(t, strings.TrimSpace(sb.String()))
+	})
+
+	t.Run("nil optimized ", func(t *testing.T) {
+		t.Parallel()
+		var sb strings.Builder
+		logger := slog.Make(slogjson.Sink(&sb))
+		prebuilds.CompareGetRunningPrebuiltWorkspacesResults(ctx, logger, []database.GetRunningPrebuiltWorkspacesRow{}, nil)
+		// Should not log any errors when optimized is nil
+		require.Empty(t, strings.TrimSpace(sb.String()))
+	})
+}