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