feat: implement claiming of prebuilt workspaces (#17458)

Signed-off-by: Danny Kopping <dannykopping@gmail.com>
Co-authored-by: Danny Kopping <dannykopping@gmail.com>
Co-authored-by: Danny Kopping <danny@coder.com>
Co-authored-by: Edward Angert <EdwardAngert@users.noreply.github.com>
Co-authored-by: EdwardAngert <17991901+EdwardAngert@users.noreply.github.com>
Co-authored-by: Jaayden Halko <jaayden.halko@gmail.com>
Co-authored-by: Ethan <39577870+ethanndickson@users.noreply.github.com>
Co-authored-by: M Atif Ali <atif@coder.com>
Co-authored-by: Aericio <16523741+Aericio@users.noreply.github.com>
Co-authored-by: M Atif Ali <me@matifali.dev>
Co-authored-by: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com>

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
 

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
+}

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{}

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.

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,

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{

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{}

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",
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+}