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",
+ },
+ },
+ },
+ },
+ },
+ },
+ },
+ },
+ }
+}
pFad - Phonifier reborn
Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.
Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.
Alternative Proxies:
Alternative Proxy
pFad Proxy
pFad v3 Proxy
pFad v4 Proxy