diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index b4162f0db59d7..55665b4381862 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -2596,6 +2596,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/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 84c2067848ec4..fba199b637c06 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -178,7 +178,8 @@ func TestDBAuthzRecursive(t *testing.T) { if method.Name == "InTx" || method.Name == "Ping" || method.Name == "Wrappers" || - method.Name == "PGLocks" { + method.Name == "PGLocks" || + method.Name == "GetRunningPrebuiltWorkspacesOptimized" { continue } // easy to know which method failed. diff --git a/coderd/database/dbauthz/setup_test.go b/coderd/database/dbauthz/setup_test.go index 555a17fb2070f..23effafc632e0 100644 --- a/coderd/database/dbauthz/setup_test.go +++ b/coderd/database/dbauthz/setup_test.go @@ -41,6 +41,8 @@ var skipMethods = map[string]string{ "Wrappers": "Not relevant", "AcquireLock": "Not relevant", "TryAcquireLock": "Not relevant", + // This method will be removed once we know this works correctly. + "GetRunningPrebuiltWorkspacesOptimized": "Not relevant", } // TestMethodTestSuite runs MethodTestSuite. diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go index 00fc3aa006e70..fda7c6325899f 100644 --- a/coderd/database/dbgen/dbgen.go +++ b/coderd/database/dbgen/dbgen.go @@ -359,6 +359,14 @@ func Workspace(t testing.TB, db database.Store, orig database.WorkspaceTable) da NextStartAt: orig.NextStartAt, }) require.NoError(t, err, "insert workspace") + if orig.Deleted { + err = db.UpdateWorkspaceDeletedByID(genCtx, database.UpdateWorkspaceDeletedByIDParams{ + ID: workspace.ID, + Deleted: true, + }) + require.NoError(t, err, "set workspace as deleted") + workspace.Deleted = true + } return workspace } diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index debb8c2b89f56..b8ae92cd9f270 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -1335,6 +1335,13 @@ func (m queryMetricsStore) GetRunningPrebuiltWorkspaces(ctx context.Context) ([] return r0, r1 } +func (m queryMetricsStore) GetRunningPrebuiltWorkspacesOptimized(ctx context.Context) ([]database.GetRunningPrebuiltWorkspacesOptimizedRow, error) { + start := time.Now() + r0, r1 := m.s.GetRunningPrebuiltWorkspacesOptimized(ctx) + m.queryLatencies.WithLabelValues("GetRunningPrebuiltWorkspacesOptimized").Observe(time.Since(start).Seconds()) + return r0, r1 +} + func (m queryMetricsStore) GetRuntimeConfig(ctx context.Context, key string) (string, error) { start := time.Now() r0, r1 := m.s.GetRuntimeConfig(ctx, key) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 059f37f8852b9..ec9ca45b195e7 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -2778,6 +2778,21 @@ func (mr *MockStoreMockRecorder) GetRunningPrebuiltWorkspaces(ctx any) *gomock.C return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetRunningPrebuiltWorkspaces", reflect.TypeOf((*MockStore)(nil).GetRunningPrebuiltWorkspaces), ctx) } +// GetRunningPrebuiltWorkspacesOptimized mocks base method. +func (m *MockStore) GetRunningPrebuiltWorkspacesOptimized(ctx context.Context) ([]database.GetRunningPrebuiltWorkspacesOptimizedRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetRunningPrebuiltWorkspacesOptimized", ctx) + ret0, _ := ret[0].([]database.GetRunningPrebuiltWorkspacesOptimizedRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetRunningPrebuiltWorkspacesOptimized indicates an expected call of GetRunningPrebuiltWorkspacesOptimized. +func (mr *MockStoreMockRecorder) GetRunningPrebuiltWorkspacesOptimized(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetRunningPrebuiltWorkspacesOptimized", reflect.TypeOf((*MockStore)(nil).GetRunningPrebuiltWorkspacesOptimized), ctx) +} + // GetRuntimeConfig mocks base method. func (m *MockStore) GetRuntimeConfig(ctx context.Context, key string) (string, error) { m.ctrl.T.Helper() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index dcbac88611dd0..b83c7415a60c8 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -298,6 +298,7 @@ type sqlcQuerier interface { GetReplicaByID(ctx context.Context, id uuid.UUID) (Replica, error) GetReplicasUpdatedAfter(ctx context.Context, updatedAt time.Time) ([]Replica, error) GetRunningPrebuiltWorkspaces(ctx context.Context) ([]GetRunningPrebuiltWorkspacesRow, error) + GetRunningPrebuiltWorkspacesOptimized(ctx context.Context) ([]GetRunningPrebuiltWorkspacesOptimizedRow, error) GetRuntimeConfig(ctx context.Context, key string) (string, error) GetTailnetAgents(ctx context.Context, id uuid.UUID) ([]TailnetAgent, error) GetTailnetClientsForAgent(ctx context.Context, agentID uuid.UUID) ([]TailnetClient, error) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 4dc7bdd66d30a..789fc85655afb 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -5020,3 +5020,102 @@ func requireUsersMatch(t testing.TB, expected []database.User, found []database. t.Helper() require.ElementsMatch(t, expected, database.ConvertUserRows(found), msg) } + +// TestGetRunningPrebuiltWorkspaces ensures the correct behavior of the +// GetRunningPrebuiltWorkspaces query. +func TestGetRunningPrebuiltWorkspaces(t *testing.T) { + t.Parallel() + + if !dbtestutil.WillUsePostgres() { + t.Skip("Test requires PostgreSQL for complex queries") + } + + ctx := testutil.Context(t, testutil.WaitLong) + db, _ := dbtestutil.NewDB(t) + now := dbtime.Now() + + // Given: a prebuilt workspace with a successful start build and a stop build. + org := dbgen.Organization(t, db, database.Organization{}) + user := dbgen.User(t, db, database.User{}) + template := dbgen.Template(t, db, database.Template{ + CreatedBy: user.ID, + OrganizationID: org.ID, + }) + templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + TemplateID: uuid.NullUUID{UUID: template.ID, Valid: true}, + OrganizationID: org.ID, + CreatedBy: user.ID, + }) + preset := dbgen.Preset(t, db, database.InsertPresetParams{ + TemplateVersionID: templateVersion.ID, + DesiredInstances: sql.NullInt32{Int32: 1, Valid: true}, + }) + + setupFixture := func(t *testing.T, db database.Store, name string, deleted bool, transition database.WorkspaceTransition, jobStatus database.ProvisionerJobStatus) database.WorkspaceTable { + t.Helper() + ws := dbgen.Workspace(t, db, database.WorkspaceTable{ + OwnerID: database.PrebuildsSystemUserID, + TemplateID: template.ID, + Name: name, + Deleted: deleted, + }) + var canceledAt sql.NullTime + var jobError sql.NullString + switch jobStatus { + case database.ProvisionerJobStatusFailed: + jobError = sql.NullString{String: assert.AnError.Error(), Valid: true} + case database.ProvisionerJobStatusCanceled: + canceledAt = sql.NullTime{Time: now, Valid: true} + } + pj := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + OrganizationID: org.ID, + InitiatorID: database.PrebuildsSystemUserID, + Provisioner: database.ProvisionerTypeEcho, + Type: database.ProvisionerJobTypeWorkspaceBuild, + StartedAt: sql.NullTime{Time: now.Add(-time.Minute), Valid: true}, + CanceledAt: canceledAt, + CompletedAt: sql.NullTime{Time: now, Valid: true}, + Error: jobError, + }) + wb := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: ws.ID, + TemplateVersionID: templateVersion.ID, + TemplateVersionPresetID: uuid.NullUUID{UUID: preset.ID, Valid: true}, + JobID: pj.ID, + BuildNumber: 1, + Transition: transition, + InitiatorID: database.PrebuildsSystemUserID, + Reason: database.BuildReasonInitiator, + }) + // Ensure things are set up as expectd + require.Equal(t, transition, wb.Transition) + require.Equal(t, int32(1), wb.BuildNumber) + require.Equal(t, jobStatus, pj.JobStatus) + require.Equal(t, deleted, ws.Deleted) + + return ws + } + + // Given: a number of prebuild workspaces with different states exist. + runningPrebuild := setupFixture(t, db, "running-prebuild", false, database.WorkspaceTransitionStart, database.ProvisionerJobStatusSucceeded) + _ = 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) + + // Given: a regular workspace also exists. + _ = dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + OwnerID: user.ID, + TemplateID: template.ID, + Name: "test-running-regular-workspace", + Deleted: false, + }) + + // When: we query for running prebuild workspaces + runningPrebuilds, err := db.GetRunningPrebuiltWorkspaces(ctx) + require.NoError(t, err) + + // Then: only the running prebuild workspace should be returned. + require.Len(t, runningPrebuilds, 1, "expected only one running prebuilt workspace") + require.Equal(t, runningPrebuild.ID, runningPrebuilds[0].ID, "expected the running prebuilt workspace to be returned") +} diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 7b55c6a4db7b8..04ded71f1242a 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -6843,6 +6843,7 @@ 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 ` type GetRunningPrebuiltWorkspacesRow struct { @@ -6886,6 +6887,106 @@ func (q *sqlQuerier) GetRunningPrebuiltWorkspaces(ctx context.Context) ([]GetRun return items, nil } +const getRunningPrebuiltWorkspacesOptimized = `-- name: GetRunningPrebuiltWorkspacesOptimized :many +WITH latest_prebuilds AS ( + -- All workspaces that match the following criteria: + -- 1. Owned by prebuilds user + -- 2. Not deleted + -- 3. Latest build is a 'start' transition + -- 4. Latest build was successful + SELECT + workspaces.id, + workspaces.name, + workspaces.template_id, + workspace_latest_builds.template_version_id, + workspace_latest_builds.job_id, + workspaces.created_at + FROM workspace_latest_builds + JOIN workspaces ON workspaces.id = workspace_latest_builds.workspace_id + WHERE workspace_latest_builds.transition = 'start'::workspace_transition + AND workspace_latest_builds.job_status = 'succeeded'::provisioner_job_status + AND workspaces.owner_id = 'c42fdf75-3097-471c-8c33-fb52454d81c0'::UUID + AND NOT workspaces.deleted +), +workspace_latest_presets AS ( + -- For each of the above workspaces, the preset_id of the most recent + -- successful start transition. + SELECT DISTINCT ON (latest_prebuilds.id) + latest_prebuilds.id AS workspace_id, + workspace_builds.template_version_preset_id AS current_preset_id + FROM latest_prebuilds + JOIN workspace_builds ON workspace_builds.workspace_id = latest_prebuilds.id + 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 +), +ready_agents AS ( + -- For each of the above workspaces, check if all agents are ready. + SELECT + latest_prebuilds.job_id, + BOOL_AND(workspace_agents.lifecycle_state = 'ready'::workspace_agent_lifecycle_state)::boolean AS ready + FROM latest_prebuilds + JOIN workspace_resources ON workspace_resources.job_id = latest_prebuilds.job_id + JOIN workspace_agents ON workspace_agents.resource_id = workspace_resources.id + WHERE workspace_agents.deleted = false + AND workspace_agents.parent_id IS NULL + GROUP BY latest_prebuilds.job_id +) +SELECT + latest_prebuilds.id, + latest_prebuilds.name, + latest_prebuilds.template_id, + latest_prebuilds.template_version_id, + workspace_latest_presets.current_preset_id, + COALESCE(ready_agents.ready, false)::boolean AS ready, + latest_prebuilds.created_at +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 +` + +type GetRunningPrebuiltWorkspacesOptimizedRow struct { + ID uuid.UUID `db:"id" json:"id"` + Name string `db:"name" json:"name"` + TemplateID uuid.UUID `db:"template_id" json:"template_id"` + TemplateVersionID uuid.UUID `db:"template_version_id" json:"template_version_id"` + CurrentPresetID uuid.NullUUID `db:"current_preset_id" json:"current_preset_id"` + Ready bool `db:"ready" json:"ready"` + CreatedAt time.Time `db:"created_at" json:"created_at"` +} + +func (q *sqlQuerier) GetRunningPrebuiltWorkspacesOptimized(ctx context.Context) ([]GetRunningPrebuiltWorkspacesOptimizedRow, error) { + rows, err := q.db.QueryContext(ctx, getRunningPrebuiltWorkspacesOptimized) + if err != nil { + return nil, err + } + defer rows.Close() + var items []GetRunningPrebuiltWorkspacesOptimizedRow + for rows.Next() { + var i GetRunningPrebuiltWorkspacesOptimizedRow + if err := rows.Scan( + &i.ID, + &i.Name, + &i.TemplateID, + &i.TemplateVersionID, + &i.CurrentPresetID, + &i.Ready, + &i.CreatedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const getTemplatePresetsWithPrebuilds = `-- name: GetTemplatePresetsWithPrebuilds :many SELECT t.id AS template_id, diff --git a/coderd/database/queries/prebuilds.sql b/coderd/database/queries/prebuilds.sql index 2fc9f3f4a67f6..7e1dbc71f4a26 100644 --- a/coderd/database/queries/prebuilds.sql +++ b/coderd/database/queries/prebuilds.sql @@ -48,6 +48,64 @@ 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: GetRunningPrebuiltWorkspacesOptimized :many +WITH latest_prebuilds AS ( + -- All workspaces that match the following criteria: + -- 1. Owned by prebuilds user + -- 2. Not deleted + -- 3. Latest build is a 'start' transition + -- 4. Latest build was successful + SELECT + workspaces.id, + workspaces.name, + workspaces.template_id, + workspace_latest_builds.template_version_id, + workspace_latest_builds.job_id, + workspaces.created_at + FROM workspace_latest_builds + JOIN workspaces ON workspaces.id = workspace_latest_builds.workspace_id + WHERE workspace_latest_builds.transition = 'start'::workspace_transition + AND workspace_latest_builds.job_status = 'succeeded'::provisioner_job_status + AND workspaces.owner_id = 'c42fdf75-3097-471c-8c33-fb52454d81c0'::UUID + AND NOT workspaces.deleted +), +workspace_latest_presets AS ( + -- For each of the above workspaces, the preset_id of the most recent + -- successful start transition. + SELECT DISTINCT ON (latest_prebuilds.id) + latest_prebuilds.id AS workspace_id, + workspace_builds.template_version_preset_id AS current_preset_id + FROM latest_prebuilds + JOIN workspace_builds ON workspace_builds.workspace_id = latest_prebuilds.id + 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 +), +ready_agents AS ( + -- For each of the above workspaces, check if all agents are ready. + SELECT + latest_prebuilds.job_id, + BOOL_AND(workspace_agents.lifecycle_state = 'ready'::workspace_agent_lifecycle_state)::boolean AS ready + FROM latest_prebuilds + JOIN workspace_resources ON workspace_resources.job_id = latest_prebuilds.job_id + JOIN workspace_agents ON workspace_agents.resource_id = workspace_resources.id + WHERE workspace_agents.deleted = false + AND workspace_agents.parent_id IS NULL + GROUP BY latest_prebuilds.job_id +) +SELECT + latest_prebuilds.id, + latest_prebuilds.name, + latest_prebuilds.template_id, + latest_prebuilds.template_version_id, + workspace_latest_presets.current_preset_id, + COALESCE(ready_agents.ready, false)::boolean AS ready, + latest_prebuilds.created_at +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, @@ -60,7 +118,8 @@ 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; -- name: CountInProgressPrebuilds :many -- CountInProgressPrebuilds returns the number of in-progress prebuilds, grouped by preset ID and transition. diff --git a/enterprise/coderd/prebuilds/reconcile.go b/enterprise/coderd/prebuilds/reconcile.go index 44e0e82c8881a..cce39ea251323 100644 --- a/enterprise/coderd/prebuilds/reconcile.go +++ b/enterprise/coderd/prebuilds/reconcile.go @@ -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)) + } +} diff --git a/enterprise/coderd/prebuilds/reconcile_test.go b/enterprise/coderd/prebuilds/reconcile_test.go index fce5269214ed1..858b01abc00b9 100644 --- a/enterprise/coderd/prebuilds/reconcile_test.go +++ b/enterprise/coderd/prebuilds/reconcile_test.go @@ -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())) + }) +}