diff --git a/coderd/coderd.go b/coderd/coderd.go index cb069fd6bf29d..4a9e3e61d9cf5 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -45,6 +45,7 @@ import ( "github.com/coder/coder/v2/coderd/entitlements" "github.com/coder/coder/v2/coderd/files" "github.com/coder/coder/v2/coderd/idpsync" + "github.com/coder/coder/v2/coderd/prebuilds" "github.com/coder/coder/v2/coderd/runtimeconfig" "github.com/coder/coder/v2/coderd/webpush" @@ -595,6 +596,7 @@ func New(options *Options) *API { f := appearance.NewDefaultFetcher(api.DeploymentValues.DocsURL.String()) api.AppearanceFetcher.Store(&f) api.PortSharer.Store(&portsharing.DefaultPortSharer) + api.PrebuildsClaimer.Store(&prebuilds.DefaultClaimer) buildInfo := codersdk.BuildInfoResponse{ ExternalURL: buildinfo.ExternalURL(), Version: buildinfo.Version(), @@ -1569,6 +1571,7 @@ type API struct { AccessControlStore *atomic.Pointer[dbauthz.AccessControlStore] PortSharer atomic.Pointer[portsharing.PortSharer] FileCache files.Cache + PrebuildsClaimer atomic.Pointer[prebuilds.Claimer] UpdatesProvider tailnet.WorkspaceUpdatesProvider diff --git a/coderd/prebuilds/api.go b/coderd/prebuilds/api.go index 6ebfb8acced44..ebc4c39c89b50 100644 --- a/coderd/prebuilds/api.go +++ b/coderd/prebuilds/api.go @@ -2,8 +2,13 @@ package prebuilds import ( "context" + + "github.com/google/uuid" + "golang.org/x/xerrors" ) +var ErrNoClaimablePrebuiltWorkspaces = xerrors.New("no claimable prebuilt workspaces found") + // ReconciliationOrchestrator manages the lifecycle of prebuild reconciliation. // It runs a continuous loop to check and reconcile prebuild states, and can be stopped gracefully. type ReconciliationOrchestrator interface { @@ -25,3 +30,8 @@ type Reconciler interface { // in parallel, creating or deleting prebuilds as needed to reach their desired states. ReconcileAll(ctx context.Context) error } + +type Claimer interface { + Claim(ctx context.Context, userID uuid.UUID, name string, presetID uuid.UUID) (*uuid.UUID, error) + Initiator() uuid.UUID +} diff --git a/coderd/prebuilds/noop.go b/coderd/prebuilds/noop.go index ffe4e7b442af9..d122a61ebb9c6 100644 --- a/coderd/prebuilds/noop.go +++ b/coderd/prebuilds/noop.go @@ -3,6 +3,8 @@ package prebuilds import ( "context" + "github.com/google/uuid" + "github.com/coder/coder/v2/coderd/database" ) @@ -33,3 +35,16 @@ func (NoopReconciler) CalculateActions(context.Context, PresetSnapshot) (*Reconc } var _ ReconciliationOrchestrator = NoopReconciler{} + +type AGPLPrebuildClaimer struct{} + +func (AGPLPrebuildClaimer) Claim(context.Context, uuid.UUID, string, uuid.UUID) (*uuid.UUID, error) { + // Not entitled to claim prebuilds in AGPL version. + return nil, ErrNoClaimablePrebuiltWorkspaces +} + +func (AGPLPrebuildClaimer) Initiator() uuid.UUID { + return uuid.Nil +} + +var DefaultClaimer Claimer = AGPLPrebuildClaimer{} diff --git a/coderd/provisionerdserver/provisionerdserver.go b/coderd/provisionerdserver/provisionerdserver.go index 22bc720736148..9362d2f3e5a85 100644 --- a/coderd/provisionerdserver/provisionerdserver.go +++ b/coderd/provisionerdserver/provisionerdserver.go @@ -2471,10 +2471,11 @@ type TemplateVersionImportJob struct { // WorkspaceProvisionJob is the payload for the "workspace_provision" job type. type WorkspaceProvisionJob struct { - WorkspaceBuildID uuid.UUID `json:"workspace_build_id"` - DryRun bool `json:"dry_run"` - IsPrebuild bool `json:"is_prebuild,omitempty"` - LogLevel string `json:"log_level,omitempty"` + WorkspaceBuildID uuid.UUID `json:"workspace_build_id"` + DryRun bool `json:"dry_run"` + IsPrebuild bool `json:"is_prebuild,omitempty"` + PrebuildClaimedByUser uuid.UUID `json:"prebuild_claimed_by,omitempty"` + LogLevel string `json:"log_level,omitempty"` } // TemplateVersionDryRunJob is the payload for the "template_version_dry_run" job type. diff --git a/coderd/workspaces.go b/coderd/workspaces.go index c1c8b1745c106..12b3787acf3d8 100644 --- a/coderd/workspaces.go +++ b/coderd/workspaces.go @@ -18,6 +18,7 @@ import ( "golang.org/x/xerrors" "cdr.dev/slog" + "github.com/coder/coder/v2/agent/proto" "github.com/coder/coder/v2/coderd/audit" "github.com/coder/coder/v2/coderd/database" @@ -28,6 +29,7 @@ import ( "github.com/coder/coder/v2/coderd/httpapi" "github.com/coder/coder/v2/coderd/httpmw" "github.com/coder/coder/v2/coderd/notifications" + "github.com/coder/coder/v2/coderd/prebuilds" "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/coderd/rbac/policy" "github.com/coder/coder/v2/coderd/schedule" @@ -636,33 +638,57 @@ func createWorkspace( workspaceBuild *database.WorkspaceBuild provisionerDaemons []database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow ) + err = api.Database.InTx(func(db database.Store) error { - now := dbtime.Now() - // Workspaces are created without any versions. - minimumWorkspace, err := db.InsertWorkspace(ctx, database.InsertWorkspaceParams{ - ID: uuid.New(), - CreatedAt: now, - UpdatedAt: now, - OwnerID: owner.ID, - OrganizationID: template.OrganizationID, - TemplateID: template.ID, - Name: req.Name, - AutostartSchedule: dbAutostartSchedule, - NextStartAt: nextStartAt, - Ttl: dbTTL, - // The workspaces page will sort by last used at, and it's useful to - // have the newly created workspace at the top of the list! - LastUsedAt: dbtime.Now(), - AutomaticUpdates: dbAU, - }) - if err != nil { - return xerrors.Errorf("insert workspace: %w", err) + var ( + workspaceID uuid.UUID + claimedWorkspace *database.Workspace + prebuildsClaimer = *api.PrebuildsClaimer.Load() + ) + + // If a template preset was chosen, try claim a prebuilt workspace. + if req.TemplateVersionPresetID != uuid.Nil { + // Try and claim an eligible prebuild, if available. + claimedWorkspace, err = claimPrebuild(ctx, prebuildsClaimer, db, api.Logger, req, owner) + if err != nil && !errors.Is(err, prebuilds.ErrNoClaimablePrebuiltWorkspaces) { + return xerrors.Errorf("claim prebuild: %w", err) + } + } + + // No prebuild found; regular flow. + if claimedWorkspace == nil { + now := dbtime.Now() + // Workspaces are created without any versions. + minimumWorkspace, err := db.InsertWorkspace(ctx, database.InsertWorkspaceParams{ + ID: uuid.New(), + CreatedAt: now, + UpdatedAt: now, + OwnerID: owner.ID, + OrganizationID: template.OrganizationID, + TemplateID: template.ID, + Name: req.Name, + AutostartSchedule: dbAutostartSchedule, + NextStartAt: nextStartAt, + Ttl: dbTTL, + // The workspaces page will sort by last used at, and it's useful to + // have the newly created workspace at the top of the list! + LastUsedAt: dbtime.Now(), + AutomaticUpdates: dbAU, + }) + if err != nil { + return xerrors.Errorf("insert workspace: %w", err) + } + workspaceID = minimumWorkspace.ID + } else { + // Prebuild found! + workspaceID = claimedWorkspace.ID + initiatorID = prebuildsClaimer.Initiator() } // We have to refetch the workspace for the joined in fields. // TODO: We can use WorkspaceTable for the builder to not require // this extra fetch. - workspace, err = db.GetWorkspaceByID(ctx, minimumWorkspace.ID) + workspace, err = db.GetWorkspaceByID(ctx, workspaceID) if err != nil { return xerrors.Errorf("get workspace by ID: %w", err) } @@ -676,6 +702,13 @@ func createWorkspace( if req.TemplateVersionID != uuid.Nil { builder = builder.VersionID(req.TemplateVersionID) } + if req.TemplateVersionPresetID != uuid.Nil { + builder = builder.TemplateVersionPresetID(req.TemplateVersionPresetID) + } + if claimedWorkspace != nil { + builder = builder.MarkPrebuildClaimedBy(owner.ID) + } + if req.EnableDynamicParameters && api.Experiments.Enabled(codersdk.ExperimentDynamicParameters) { builder = builder.UsingDynamicParameters() } @@ -842,6 +875,21 @@ func requestTemplate(ctx context.Context, rw http.ResponseWriter, req codersdk.C return template, true } +func claimPrebuild(ctx context.Context, claimer prebuilds.Claimer, db database.Store, logger slog.Logger, req codersdk.CreateWorkspaceRequest, owner workspaceOwner) (*database.Workspace, error) { + claimedID, err := claimer.Claim(ctx, owner.ID, req.Name, req.TemplateVersionPresetID) + if err != nil { + // TODO: enhance this by clarifying whether this *specific* prebuild failed or whether there are none to claim. + return nil, xerrors.Errorf("claim prebuild: %w", err) + } + + lookup, err := db.GetWorkspaceByID(ctx, *claimedID) + if err != nil { + logger.Error(ctx, "unable to find claimed workspace by ID", slog.Error(err), slog.F("claimed_prebuild_id", claimedID.String())) + return nil, xerrors.Errorf("find claimed workspace by ID %q: %w", claimedID.String(), err) + } + return &lookup, nil +} + func (api *API) notifyWorkspaceCreated( ctx context.Context, receiverID uuid.UUID, diff --git a/coderd/wsbuilder/wsbuilder.go b/coderd/wsbuilder/wsbuilder.go index 5ac0f54639a06..942829004309c 100644 --- a/coderd/wsbuilder/wsbuilder.go +++ b/coderd/wsbuilder/wsbuilder.go @@ -76,7 +76,8 @@ type Builder struct { parameterValues *[]string templateVersionPresetParameterValues []database.TemplateVersionPresetParameter - prebuild bool + prebuild bool + prebuildClaimedBy uuid.UUID verifyNoLegacyParametersOnce bool } @@ -179,6 +180,12 @@ func (b Builder) MarkPrebuild() Builder { return b } +func (b Builder) MarkPrebuildClaimedBy(userID uuid.UUID) Builder { + // nolint: revive + b.prebuildClaimedBy = userID + return b +} + func (b Builder) UsingDynamicParameters() Builder { b.dynamicParametersEnabled = true return b @@ -315,9 +322,10 @@ func (b *Builder) buildTx(authFunc func(action policy.Action, object rbac.Object workspaceBuildID := uuid.New() input, err := json.Marshal(provisionerdserver.WorkspaceProvisionJob{ - WorkspaceBuildID: workspaceBuildID, - LogLevel: b.logLevel, - IsPrebuild: b.prebuild, + WorkspaceBuildID: workspaceBuildID, + LogLevel: b.logLevel, + IsPrebuild: b.prebuild, + PrebuildClaimedByUser: b.prebuildClaimedBy, }) if err != nil { return nil, nil, nil, BuildError{ diff --git a/enterprise/coderd/prebuilds/claim.go b/enterprise/coderd/prebuilds/claim.go new file mode 100644 index 0000000000000..f040ee756e678 --- /dev/null +++ b/enterprise/coderd/prebuilds/claim.go @@ -0,0 +1,53 @@ +package prebuilds + +import ( + "context" + "database/sql" + "errors" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/prebuilds" +) + +type EnterpriseClaimer struct { + store database.Store +} + +func NewEnterpriseClaimer(store database.Store) *EnterpriseClaimer { + return &EnterpriseClaimer{ + store: store, + } +} + +func (c EnterpriseClaimer) Claim( + ctx context.Context, + userID uuid.UUID, + name string, + presetID uuid.UUID, +) (*uuid.UUID, error) { + result, err := c.store.ClaimPrebuiltWorkspace(ctx, database.ClaimPrebuiltWorkspaceParams{ + NewUserID: userID, + NewName: name, + PresetID: presetID, + }) + if err != nil { + switch { + // No eligible prebuilds found + case errors.Is(err, sql.ErrNoRows): + return nil, prebuilds.ErrNoClaimablePrebuiltWorkspaces + default: + return nil, xerrors.Errorf("claim prebuild for user %q: %w", userID.String(), err) + } + } + + return &result.ID, nil +} + +func (EnterpriseClaimer) Initiator() uuid.UUID { + return prebuilds.SystemUserID +} + +var _ prebuilds.Claimer = &EnterpriseClaimer{} diff --git a/enterprise/coderd/prebuilds/claim_test.go b/enterprise/coderd/prebuilds/claim_test.go new file mode 100644 index 0000000000000..4f398724b8265 --- /dev/null +++ b/enterprise/coderd/prebuilds/claim_test.go @@ -0,0 +1,564 @@ +package prebuilds_test + +import ( + "context" + "database/sql" + "slices" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + + "github.com/coder/quartz" + + "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbtestutil" + agplprebuilds "github.com/coder/coder/v2/coderd/prebuilds" + "github.com/coder/coder/v2/coderd/rbac" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/enterprise/coderd/coderdenttest" + "github.com/coder/coder/v2/enterprise/coderd/prebuilds" + "github.com/coder/coder/v2/provisioner/echo" + "github.com/coder/coder/v2/provisionersdk/proto" + "github.com/coder/coder/v2/testutil" +) + +type storeSpy struct { + database.Store + + claims *atomic.Int32 + claimParams *atomic.Pointer[database.ClaimPrebuiltWorkspaceParams] + claimedWorkspace *atomic.Pointer[database.ClaimPrebuiltWorkspaceRow] +} + +func newStoreSpy(db database.Store) *storeSpy { + return &storeSpy{ + Store: db, + claims: &atomic.Int32{}, + claimParams: &atomic.Pointer[database.ClaimPrebuiltWorkspaceParams]{}, + claimedWorkspace: &atomic.Pointer[database.ClaimPrebuiltWorkspaceRow]{}, + } +} + +func (m *storeSpy) InTx(fn func(store database.Store) error, opts *database.TxOptions) error { + // Pass spy down into transaction store. + return m.Store.InTx(func(store database.Store) error { + spy := newStoreSpy(store) + spy.claims = m.claims + spy.claimParams = m.claimParams + spy.claimedWorkspace = m.claimedWorkspace + + return fn(spy) + }, opts) +} + +func (m *storeSpy) ClaimPrebuiltWorkspace(ctx context.Context, arg database.ClaimPrebuiltWorkspaceParams) (database.ClaimPrebuiltWorkspaceRow, error) { + m.claims.Add(1) + m.claimParams.Store(&arg) + result, err := m.Store.ClaimPrebuiltWorkspace(ctx, arg) + if err == nil { + m.claimedWorkspace.Store(&result) + } + return result, err +} + +type errorStore struct { + claimingErr error + + database.Store +} + +func newErrorStore(db database.Store, claimingErr error) *errorStore { + return &errorStore{ + Store: db, + claimingErr: claimingErr, + } +} + +func (es *errorStore) InTx(fn func(store database.Store) error, opts *database.TxOptions) error { + // Pass failure store down into transaction store. + return es.Store.InTx(func(store database.Store) error { + newES := newErrorStore(store, es.claimingErr) + + return fn(newES) + }, opts) +} + +func (es *errorStore) ClaimPrebuiltWorkspace(ctx context.Context, arg database.ClaimPrebuiltWorkspaceParams) (database.ClaimPrebuiltWorkspaceRow, error) { + return database.ClaimPrebuiltWorkspaceRow{}, es.claimingErr +} + +func TestClaimPrebuild(t *testing.T) { + t.Parallel() + + if !dbtestutil.WillUsePostgres() { + t.Skip("This test requires postgres") + } + + const ( + desiredInstances = 1 + presetCount = 2 + ) + + cases := map[string]struct { + expectPrebuildClaimed bool + markPrebuildsClaimable bool + }{ + "no eligible prebuilds to claim": { + expectPrebuildClaimed: false, + markPrebuildsClaimable: false, + }, + "claiming an eligible prebuild should succeed": { + expectPrebuildClaimed: true, + markPrebuildsClaimable: true, + }, + } + + for name, tc := range cases { + tc := tc + + t.Run(name, func(t *testing.T) { + t.Parallel() + + // Setup. + ctx := testutil.Context(t, testutil.WaitSuperLong) + db, pubsub := dbtestutil.NewDB(t) + spy := newStoreSpy(db) + expectedPrebuildsCount := desiredInstances * presetCount + + logger := testutil.Logger(t) + client, _, api, owner := coderdenttest.NewWithAPI(t, &coderdenttest.Options{ + Options: &coderdtest.Options{ + IncludeProvisionerDaemon: true, + Database: spy, + Pubsub: pubsub, + }, + + EntitlementsUpdateInterval: time.Second, + }) + + reconciler := prebuilds.NewStoreReconciler(spy, pubsub, codersdk.PrebuildsConfig{}, logger, quartz.NewMock(t)) + var claimer agplprebuilds.Claimer = prebuilds.NewEnterpriseClaimer(spy) + api.AGPL.PrebuildsClaimer.Store(&claimer) + + version := coderdtest.CreateTemplateVersion(t, client, owner.OrganizationID, templateWithAgentAndPresetsWithPrebuilds(desiredInstances)) + _ = coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID) + coderdtest.CreateTemplate(t, client, owner.OrganizationID, version.ID) + presets, err := client.TemplateVersionPresets(ctx, version.ID) + require.NoError(t, err) + require.Len(t, presets, presetCount) + + userClient, user := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID, rbac.RoleMember()) + + // Given: the reconciliation state is snapshot. + state, err := reconciler.SnapshotState(ctx, spy) + require.NoError(t, err) + require.Len(t, state.Presets, presetCount) + + // When: a reconciliation is setup for each preset. + for _, preset := range presets { + ps, err := state.FilterByPreset(preset.ID) + require.NoError(t, err) + require.NotNil(t, ps) + actions, err := reconciler.CalculateActions(ctx, *ps) + require.NoError(t, err) + require.NotNil(t, actions) + + require.NoError(t, reconciler.ReconcilePreset(ctx, *ps)) + } + + // Given: a set of running, eligible prebuilds eventually starts up. + runningPrebuilds := make(map[uuid.UUID]database.GetRunningPrebuiltWorkspacesRow, desiredInstances*presetCount) + require.Eventually(t, func() bool { + rows, err := spy.GetRunningPrebuiltWorkspaces(ctx) + if err != nil { + return false + } + + for _, row := range rows { + runningPrebuilds[row.CurrentPresetID.UUID] = row + + if !tc.markPrebuildsClaimable { + continue + } + + agents, err := db.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, row.ID) + if err != nil { + return false + } + + // Workspaces are eligible once its agent is marked "ready". + for _, agent := range agents { + err = db.UpdateWorkspaceAgentLifecycleStateByID(ctx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{ + ID: agent.ID, + LifecycleState: database.WorkspaceAgentLifecycleStateReady, + StartedAt: sql.NullTime{Time: time.Now().Add(time.Hour), Valid: true}, + ReadyAt: sql.NullTime{Time: time.Now().Add(-1 * time.Hour), Valid: true}, + }) + if err != nil { + return false + } + } + } + + t.Logf("found %d running prebuilds so far, want %d", len(runningPrebuilds), expectedPrebuildsCount) + + return len(runningPrebuilds) == expectedPrebuildsCount + }, testutil.WaitSuperLong, testutil.IntervalSlow) + + // When: a user creates a new workspace with a preset for which prebuilds are configured. + workspaceName := strings.ReplaceAll(testutil.GetRandomName(t), "_", "-") + params := database.ClaimPrebuiltWorkspaceParams{ + NewUserID: user.ID, + NewName: workspaceName, + PresetID: presets[0].ID, + } + userWorkspace, err := userClient.CreateUserWorkspace(ctx, user.Username, codersdk.CreateWorkspaceRequest{ + TemplateVersionID: version.ID, + Name: workspaceName, + TemplateVersionPresetID: presets[0].ID, + }) + + require.NoError(t, err) + coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, userWorkspace.LatestBuild.ID) + + // Then: a prebuild should have been claimed. + require.EqualValues(t, spy.claims.Load(), 1) + require.EqualValues(t, *spy.claimParams.Load(), params) + + if !tc.expectPrebuildClaimed { + require.Nil(t, spy.claimedWorkspace.Load()) + return + } + + require.NotNil(t, spy.claimedWorkspace.Load()) + claimed := *spy.claimedWorkspace.Load() + require.NotEqual(t, claimed.ID, uuid.Nil) + + // Then: the claimed prebuild must now be owned by the requester. + workspace, err := spy.GetWorkspaceByID(ctx, claimed.ID) + require.NoError(t, err) + require.Equal(t, user.ID, workspace.OwnerID) + + // Then: the number of running prebuilds has changed since one was claimed. + currentPrebuilds, err := spy.GetRunningPrebuiltWorkspaces(ctx) + require.NoError(t, err) + require.Equal(t, expectedPrebuildsCount-1, len(currentPrebuilds)) + + // Then: the claimed prebuild is now missing from the running prebuilds set. + found := slices.ContainsFunc(currentPrebuilds, func(prebuild database.GetRunningPrebuiltWorkspacesRow) bool { + return prebuild.ID == claimed.ID + }) + require.False(t, found, "claimed prebuild should not still be considered a running prebuild") + + // Then: reconciling at this point will provision a new prebuild to replace the claimed one. + { + // Given: the reconciliation state is snapshot. + state, err = reconciler.SnapshotState(ctx, spy) + require.NoError(t, err) + + // When: a reconciliation is setup for each preset. + for _, preset := range presets { + ps, err := state.FilterByPreset(preset.ID) + require.NoError(t, err) + + // Then: the reconciliation takes place without error. + require.NoError(t, reconciler.ReconcilePreset(ctx, *ps)) + } + } + + require.Eventually(t, func() bool { + rows, err := spy.GetRunningPrebuiltWorkspaces(ctx) + if err != nil { + return false + } + + t.Logf("found %d running prebuilds so far, want %d", len(rows), expectedPrebuildsCount) + + return len(runningPrebuilds) == expectedPrebuildsCount + }, testutil.WaitSuperLong, testutil.IntervalSlow) + + // Then: when restarting the created workspace (which claimed a prebuild), it should not try and claim a new prebuild. + // Prebuilds should ONLY be used for net-new workspaces. + // This is expected by default anyway currently since new workspaces and operations on existing workspaces + // take different code paths, but it's worth validating. + + spy.claims.Store(0) // Reset counter because we need to check if any new claim requests happen. + + wp, err := userClient.WorkspaceBuildParameters(ctx, userWorkspace.LatestBuild.ID) + require.NoError(t, err) + + stopBuild, err := userClient.CreateWorkspaceBuild(ctx, workspace.ID, codersdk.CreateWorkspaceBuildRequest{ + TemplateVersionID: version.ID, + Transition: codersdk.WorkspaceTransitionStop, + RichParameterValues: wp, + }) + require.NoError(t, err) + coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, stopBuild.ID) + + startBuild, err := userClient.CreateWorkspaceBuild(ctx, workspace.ID, codersdk.CreateWorkspaceBuildRequest{ + TemplateVersionID: version.ID, + Transition: codersdk.WorkspaceTransitionStart, + RichParameterValues: wp, + }) + require.NoError(t, err) + coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, startBuild.ID) + + require.Zero(t, spy.claims.Load()) + }) + } +} + +func TestClaimPrebuild_CheckDifferentErrors(t *testing.T) { + t.Parallel() + + if !dbtestutil.WillUsePostgres() { + t.Skip("This test requires postgres") + } + + const ( + desiredInstances = 1 + presetCount = 2 + + expectedPrebuildsCount = desiredInstances * presetCount + ) + + cases := map[string]struct { + claimingErr error + checkFn func( + t *testing.T, + ctx context.Context, + store database.Store, + userClient *codersdk.Client, + user codersdk.User, + templateVersionID uuid.UUID, + presetID uuid.UUID, + ) + }{ + "ErrNoClaimablePrebuiltWorkspaces is returned": { + claimingErr: agplprebuilds.ErrNoClaimablePrebuiltWorkspaces, + checkFn: func( + t *testing.T, + ctx context.Context, + store database.Store, + userClient *codersdk.Client, + user codersdk.User, + templateVersionID uuid.UUID, + presetID uuid.UUID, + ) { + // When: a user creates a new workspace with a preset for which prebuilds are configured. + workspaceName := strings.ReplaceAll(testutil.GetRandomName(t), "_", "-") + userWorkspace, err := userClient.CreateUserWorkspace(ctx, user.Username, codersdk.CreateWorkspaceRequest{ + TemplateVersionID: templateVersionID, + Name: workspaceName, + TemplateVersionPresetID: presetID, + }) + + require.NoError(t, err) + coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, userWorkspace.LatestBuild.ID) + + // Then: the number of running prebuilds hasn't changed because claiming prebuild is failed and we fallback to creating new workspace. + currentPrebuilds, err := store.GetRunningPrebuiltWorkspaces(ctx) + require.NoError(t, err) + require.Equal(t, expectedPrebuildsCount, len(currentPrebuilds)) + }, + }, + "unexpected error during claim is returned": { + claimingErr: xerrors.New("unexpected error during claim"), + checkFn: func( + t *testing.T, + ctx context.Context, + store database.Store, + userClient *codersdk.Client, + user codersdk.User, + templateVersionID uuid.UUID, + presetID uuid.UUID, + ) { + // When: a user creates a new workspace with a preset for which prebuilds are configured. + workspaceName := strings.ReplaceAll(testutil.GetRandomName(t), "_", "-") + _, err := userClient.CreateUserWorkspace(ctx, user.Username, codersdk.CreateWorkspaceRequest{ + TemplateVersionID: templateVersionID, + Name: workspaceName, + TemplateVersionPresetID: presetID, + }) + + // Then: unexpected error happened and was propagated all the way to the caller + require.Error(t, err) + require.ErrorContains(t, err, "unexpected error during claim") + + // Then: the number of running prebuilds hasn't changed because claiming prebuild is failed. + currentPrebuilds, err := store.GetRunningPrebuiltWorkspaces(ctx) + require.NoError(t, err) + require.Equal(t, expectedPrebuildsCount, len(currentPrebuilds)) + }, + }, + } + + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + t.Parallel() + + // Setup. + ctx := testutil.Context(t, testutil.WaitSuperLong) + db, pubsub := dbtestutil.NewDB(t) + errorStore := newErrorStore(db, tc.claimingErr) + + logger := testutil.Logger(t) + client, _, api, owner := coderdenttest.NewWithAPI(t, &coderdenttest.Options{ + Options: &coderdtest.Options{ + IncludeProvisionerDaemon: true, + Database: errorStore, + Pubsub: pubsub, + }, + + EntitlementsUpdateInterval: time.Second, + }) + + reconciler := prebuilds.NewStoreReconciler(errorStore, pubsub, codersdk.PrebuildsConfig{}, logger, quartz.NewMock(t)) + var claimer agplprebuilds.Claimer = prebuilds.NewEnterpriseClaimer(errorStore) + api.AGPL.PrebuildsClaimer.Store(&claimer) + + version := coderdtest.CreateTemplateVersion(t, client, owner.OrganizationID, templateWithAgentAndPresetsWithPrebuilds(desiredInstances)) + _ = coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID) + coderdtest.CreateTemplate(t, client, owner.OrganizationID, version.ID) + presets, err := client.TemplateVersionPresets(ctx, version.ID) + require.NoError(t, err) + require.Len(t, presets, presetCount) + + userClient, user := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID, rbac.RoleMember()) + + // Given: the reconciliation state is snapshot. + state, err := reconciler.SnapshotState(ctx, errorStore) + require.NoError(t, err) + require.Len(t, state.Presets, presetCount) + + // When: a reconciliation is setup for each preset. + for _, preset := range presets { + ps, err := state.FilterByPreset(preset.ID) + require.NoError(t, err) + require.NotNil(t, ps) + actions, err := reconciler.CalculateActions(ctx, *ps) + require.NoError(t, err) + require.NotNil(t, actions) + + require.NoError(t, reconciler.ReconcilePreset(ctx, *ps)) + } + + // Given: a set of running, eligible prebuilds eventually starts up. + runningPrebuilds := make(map[uuid.UUID]database.GetRunningPrebuiltWorkspacesRow, desiredInstances*presetCount) + require.Eventually(t, func() bool { + rows, err := errorStore.GetRunningPrebuiltWorkspaces(ctx) + if err != nil { + return false + } + + for _, row := range rows { + runningPrebuilds[row.CurrentPresetID.UUID] = row + + agents, err := db.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, row.ID) + if err != nil { + return false + } + + // Workspaces are eligible once its agent is marked "ready". + for _, agent := range agents { + err = db.UpdateWorkspaceAgentLifecycleStateByID(ctx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{ + ID: agent.ID, + LifecycleState: database.WorkspaceAgentLifecycleStateReady, + StartedAt: sql.NullTime{Time: time.Now().Add(time.Hour), Valid: true}, + ReadyAt: sql.NullTime{Time: time.Now().Add(-1 * time.Hour), Valid: true}, + }) + if err != nil { + return false + } + } + } + + t.Logf("found %d running prebuilds so far, want %d", len(runningPrebuilds), expectedPrebuildsCount) + + return len(runningPrebuilds) == expectedPrebuildsCount + }, testutil.WaitSuperLong, testutil.IntervalSlow) + + tc.checkFn(t, ctx, errorStore, userClient, user, version.ID, presets[0].ID) + }) + } +} + +func templateWithAgentAndPresetsWithPrebuilds(desiredInstances int32) *echo.Responses { + return &echo.Responses{ + Parse: echo.ParseComplete, + ProvisionPlan: []*proto.Response{ + { + Type: &proto.Response_Plan{ + Plan: &proto.PlanComplete{ + Resources: []*proto.Resource{ + { + Type: "compute", + Name: "main", + Agents: []*proto.Agent{ + { + Name: "smith", + OperatingSystem: "linux", + Architecture: "i386", + }, + }, + }, + }, + Presets: []*proto.Preset{ + { + Name: "preset-a", + Parameters: []*proto.PresetParameter{ + { + Name: "k1", + Value: "v1", + }, + }, + Prebuild: &proto.Prebuild{ + Instances: desiredInstances, + }, + }, + { + Name: "preset-b", + Parameters: []*proto.PresetParameter{ + { + Name: "k1", + Value: "v2", + }, + }, + Prebuild: &proto.Prebuild{ + Instances: desiredInstances, + }, + }, + }, + }, + }, + }, + }, + ProvisionApply: []*proto.Response{ + { + Type: &proto.Response_Apply{ + Apply: &proto.ApplyComplete{ + Resources: []*proto.Resource{ + { + Type: "compute", + Name: "main", + Agents: []*proto.Agent{ + { + Name: "smith", + OperatingSystem: "linux", + Architecture: "i386", + }, + }, + }, + }, + }, + }, + }, + }, + } +}