coder · evgeniy-scherbina · Apr 24, 2025 · Apr 17, 2025 · Apr 22, 2025 · Apr 21, 2025
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -19,6 +19,8 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/coder/coder/v2/coderd/prebuilds"
+
 	"github.com/andybalholm/brotli"
 	"github.com/go-chi/chi/v5"
 	"github.com/go-chi/chi/v5/middleware"
@@ -595,6 +597,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 +1572,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
@@ -2,6 +2,10 @@ package prebuilds
 
 import (
 	"context"
+
+	"github.com/google/uuid"
+
+	"github.com/coder/coder/v2/coderd/database"
 )
 
 // ReconciliationOrchestrator manages the lifecycle of prebuild reconciliation.
@@ -19,9 +23,38 @@ type ReconciliationOrchestrator interface {
 	Stop(ctx context.Context, cause error)
 }
 
+// Reconciler defines the core operations for managing prebuilds.
+// It provides both high-level orchestration (ReconcileAll) and lower-level operations
+// for more fine-grained control (SnapshotState, ReconcilePreset, CalculateActions).
+// All database operations must be performed within repeatable-read transactions
+// to ensure consistency.
 type Reconciler interface {
 	// ReconcileAll orchestrates the reconciliation of all prebuilds across all templates.
 	// It takes a global snapshot of the system state and then reconciles each preset
 	// in parallel, creating or deleting prebuilds as needed to reach their desired states.
+	// For more fine-grained control, you can use the lower-level methods SnapshotState
+	// and ReconcilePreset directly.
 	ReconcileAll(ctx context.Context) error
+
+	// SnapshotState captures the current state of all prebuilds across templates.
+	// It creates a global database snapshot that can be viewed as a collection of PresetSnapshots,
+	// each representing the state of prebuilds for a specific preset.
+	// MUST be called inside a repeatable-read transaction.
+	SnapshotState(ctx context.Context, store database.Store) (*GlobalSnapshot, error)
+
+	// ReconcilePreset handles a single PresetSnapshot, determining and executing
+	// the required actions (creating or deleting prebuilds) based on the current state.
+	// MUST be called inside a repeatable-read transaction.
+	ReconcilePreset(ctx context.Context, snapshot PresetSnapshot) error
+
+	// CalculateActions determines what actions are needed to reconcile a preset's prebuilds
+	// to their desired state. This includes creating new prebuilds, deleting excess ones,
+	// or waiting due to backoff periods.
+	// MUST be called inside a repeatable-read transaction.
+	CalculateActions(ctx context.Context, state PresetSnapshot) (*ReconciliationActions, error)
+}
+
+type Claimer interface {
+	Claim(ctx context.Context, store database.Store, 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
@@ -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 (c AGPLPrebuildClaimer) Claim(context.Context, database.Store, uuid.UUID, string, uuid.UUID) (*uuid.UUID, error) {
+	// Not entitled to claim prebuilds in AGPL version.
+	return nil, nil
+}
+
+func (c AGPLPrebuildClaimer) Initiator() uuid.UUID {
+	return uuid.Nil
+}
+
+var DefaultClaimer Claimer = AGPLPrebuildClaimer{}
diff --git a/coderd/provisionerdserver/provisionerdserver.go b/coderd/provisionerdserver/provisionerdserver.go
@@ -2469,10 +2469,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
@@ -11,13 +11,16 @@ import (
 	"strconv"
 	"time"
 
+	"github.com/coder/coder/v2/coderd/prebuilds"
+
 	"github.com/dustin/go-humanize"
 	"github.com/go-chi/chi/v5"
 	"github.com/google/uuid"
 	"golang.org/x/sync/errgroup"
 	"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"
@@ -636,33 +639,58 @@ func createWorkspace(
 		workspaceBuild     *database.WorkspaceBuild
 		provisionerDaemons []database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow
 	)
+
+	prebuilds := *api.PrebuildsClaimer.Load()
+
 	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
+		)
+
+		// If a template preset was chosen, try claim a prebuild.
+		if req.TemplateVersionPresetID != uuid.Nil {
+			// Try and claim an eligible prebuild, if available.
+			claimedWorkspace, err = claimPrebuild(ctx, prebuilds, db, api.Logger, req, owner)
+			if err != nil {
+				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 = prebuilds.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 +704,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)
+		}
 
 		workspaceBuild, provisionerJob, provisionerDaemons, err = builder.Build(
 			ctx,
@@ -839,6 +874,32 @@ 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) {
+	prebuildsCtx := dbauthz.AsPrebuildsOrchestrator(ctx)
+
+	// TODO: do we need a timeout here?
+	claimCtx, cancel := context.WithTimeout(prebuildsCtx, time.Second*10)
+	defer cancel()
+
+	claimedID, err := claimer.Claim(claimCtx, db, 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)
+	}
+
+	// No prebuild available.
+	if claimedID == nil {
+		return nil, nil
+	}
+
+	lookup, err := db.GetWorkspaceByID(prebuildsCtx, *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, err
+}
+
 func (api *API) notifyWorkspaceCreated(
 	ctx context.Context,
 	receiverID uuid.UUID,

diff --git a/coderd/wsbuilder/wsbuilder.go b/coderd/wsbuilder/wsbuilder.go
@@ -75,7 +75,8 @@ type Builder struct {
 	parameterValues                      *[]string
 	templateVersionPresetParameterValues []database.TemplateVersionPresetParameter
 
-	prebuild bool
+	prebuild          bool
+	prebuildClaimedBy uuid.UUID
 
 	verifyNoLegacyParametersOnce bool
 }
@@ -178,6 +179,12 @@ func (b Builder) MarkPrebuild() Builder {
 	return b
 }
 
+func (b Builder) MarkPrebuildClaimedBy(userID uuid.UUID) Builder {
+	// nolint: revive
+	b.prebuildClaimedBy = userID
+	return b
+}
+
 // SetLastWorkspaceBuildInTx prepopulates the Builder's cache with the last workspace build.  This allows us
 // to avoid a repeated database query when the Builder's caller also needs the workspace build, e.g. auto-start &
 // auto-stop.
@@ -309,9 +316,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{
@@ -624,6 +632,11 @@ func (b *Builder) findNewBuildParameterValue(name string) *codersdk.WorkspaceBui
 }
 
 func (b *Builder) getLastBuildParameters() ([]database.WorkspaceBuildParameter, error) {
+	// TODO: exclude preset params from this list instead of returning nothing?
+	if b.prebuildClaimedBy != uuid.Nil {
+		return nil, nil
+	}
+
 	if b.lastBuildParameters != nil {
 		return *b.lastBuildParameters, nil
 	}

diff --git a/enterprise/coderd/prebuilds/claim.go b/enterprise/coderd/prebuilds/claim.go
@@ -0,0 +1,64 @@
+package prebuilds
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+
+	"github.com/coder/coder/v2/coderd/prebuilds"
+
+	"github.com/google/uuid"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/v2/coderd/database"
+)
+
+type EnterpriseClaimer struct{}
+
+func (_ EnterpriseClaimer) Claim(
+	ctx context.Context,
+	store database.Store,
+	userID uuid.UUID,
+	name string,
+	presetID uuid.UUID,
+) (*uuid.UUID, error) {
+	var prebuildID *uuid.UUID
+	err := store.InTx(func(db database.Store) error {
+		// TODO: do we need this?
+		//// Ensure no other replica can claim a prebuild for this user simultaneously.
+		// err := store.AcquireLock(ctx, database.GenLockID(fmt.Sprintf("prebuild-user-claim-%s", userID.String())))
+		// if err != nil {
+		//	return xerrors.Errorf("acquire claim lock for user %q: %w", userID.String(), err)
+		//}
+
+		result, err := db.ClaimPrebuiltWorkspace(ctx, database.ClaimPrebuiltWorkspaceParams{
+			NewUserID: userID,
+			NewName:   name,
+			PresetID:  presetID,
+		})
+		if err != nil {
+			switch {
+			// No eligible prebuilds found
+			case errors.Is(err, sql.ErrNoRows):
+				// Exit, this will result in a nil prebuildID being returned, which is fine
+				return nil
+			default:
+				return xerrors.Errorf("claim prebuild for user %q: %w", userID.String(), err)
+			}
+		}
+
+		prebuildID = &result.ID
+
+		return nil
+	}, &database.TxOptions{
+		TxIdentifier: "prebuild-claim",
+	})
+
+	return prebuildID, err
+}
+
+func (_ EnterpriseClaimer) Initiator() uuid.UUID {
+	return prebuilds.SystemUserID
+}
+
+var _ prebuilds.Claimer = &EnterpriseClaimer{}