Skip to content

Commit 8b9e30d

Browse files
evgeniy-scherbinaevgeniy-scherbina
committed
feat: implement claiming of prebuilt workspaces
1 parent 183146e commit 8b9e30d

File tree

11 files changed

+769
-34
lines changed

11 files changed

+769
-34
lines changed

‎coderd/coderd.go

+3
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"expvar"
1010
"flag"
1111
"fmt"
12+
"github.com/coder/coder/v2/coderd/prebuilds"
1213
"io"
1314
"net/http"
1415
"net/url"
@@ -595,6 +596,7 @@ func New(options *Options) *API {
595596
f := appearance.NewDefaultFetcher(api.DeploymentValues.DocsURL.String())
596597
api.AppearanceFetcher.Store(&f)
597598
api.PortSharer.Store(&portsharing.DefaultPortSharer)
599+
api.PrebuildsClaimer.Store(&prebuilds.DefaultClaimer)
598600
buildInfo := codersdk.BuildInfoResponse{
599601
ExternalURL: buildinfo.ExternalURL(),
600602
Version: buildinfo.Version(),
@@ -1562,6 +1564,7 @@ type API struct {
15621564
AccessControlStore *atomic.Pointer[dbauthz.AccessControlStore]
15631565
PortSharer atomic.Pointer[portsharing.PortSharer]
15641566
FileCache files.Cache
1567+
PrebuildsClaimer atomic.Pointer[prebuilds.Claimer]
15651568

15661569
UpdatesProvider tailnet.WorkspaceUpdatesProvider
15671570

‎coderd/prebuilds/api.go

+33
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,10 @@ package prebuilds
22

33
import (
44
"context"
5+
6+
"github.com/google/uuid"
7+
8+
"github.com/coder/coder/v2/coderd/database"
59
)
610

711
// ReconciliationOrchestrator manages the lifecycle of prebuild reconciliation.
@@ -19,9 +23,38 @@ type ReconciliationOrchestrator interface {
1923
Stop(ctx context.Context, cause error)
2024
}
2125

26+
// Reconciler defines the core operations for managing prebuilds.
27+
// It provides both high-level orchestration (ReconcileAll) and lower-level operations
28+
// for more fine-grained control (SnapshotState, ReconcilePreset, CalculateActions).
29+
// All database operations must be performed within repeatable-read transactions
30+
// to ensure consistency.
2231
type Reconciler interface {
2332
// ReconcileAll orchestrates the reconciliation of all prebuilds across all templates.
2433
// It takes a global snapshot of the system state and then reconciles each preset
2534
// in parallel, creating or deleting prebuilds as needed to reach their desired states.
35+
// For more fine-grained control, you can use the lower-level methods SnapshotState
36+
// and ReconcilePreset directly.
2637
ReconcileAll(ctx context.Context) error
38+
39+
// SnapshotState captures the current state of all prebuilds across templates.
40+
// It creates a global database snapshot that can be viewed as a collection of PresetSnapshots,
41+
// each representing the state of prebuilds for a specific preset.
42+
// MUST be called inside a repeatable-read transaction.
43+
SnapshotState(ctx context.Context, store database.Store) (*GlobalSnapshot, error)
44+
45+
// ReconcilePreset handles a single PresetSnapshot, determining and executing
46+
// the required actions (creating or deleting prebuilds) based on the current state.
47+
// MUST be called inside a repeatable-read transaction.
48+
ReconcilePreset(ctx context.Context, snapshot PresetSnapshot) error
49+
50+
// CalculateActions determines what actions are needed to reconcile a preset's prebuilds
51+
// to their desired state. This includes creating new prebuilds, deleting excess ones,
52+
// or waiting due to backoff periods.
53+
// MUST be called inside a repeatable-read transaction.
54+
CalculateActions(ctx context.Context, state PresetSnapshot) (*ReconciliationActions, error)
55+
}
56+
57+
type Claimer interface {
58+
Claim(ctx context.Context, store database.Store, userID uuid.UUID, name string, presetID uuid.UUID) (*uuid.UUID, error)
59+
Initiator() uuid.UUID
2760
}

‎coderd/prebuilds/noop.go

+15
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ package prebuilds
33
import (
44
"context"
55

6+
"github.com/google/uuid"
7+
68
"github.com/coder/coder/v2/coderd/database"
79
)
810

@@ -33,3 +35,16 @@ func (NoopReconciler) CalculateActions(context.Context, PresetSnapshot) (*Reconc
3335
}
3436

3537
var _ ReconciliationOrchestrator = NoopReconciler{}
38+
39+
type AGPLPrebuildClaimer struct{}
40+
41+
func (c AGPLPrebuildClaimer) Claim(context.Context, database.Store, uuid.UUID, string, uuid.UUID) (*uuid.UUID, error) {
42+
// Not entitled to claim prebuilds in AGPL version.
43+
return nil, nil
44+
}
45+
46+
func (c AGPLPrebuildClaimer) Initiator() uuid.UUID {
47+
return uuid.Nil
48+
}
49+
50+
var DefaultClaimer Claimer = AGPLPrebuildClaimer{}

‎coderd/provisionerdserver/provisionerdserver.go

+12-4
Original file line numberDiff line numberDiff line change
@@ -2462,10 +2462,18 @@ type TemplateVersionImportJob struct {
24622462

24632463
// WorkspaceProvisionJob is the payload for the "workspace_provision" job type.
24642464
type WorkspaceProvisionJob struct {
2465-
WorkspaceBuildID uuid.UUID `json:"workspace_build_id"`
2466-
DryRun bool `json:"dry_run"`
2467-
IsPrebuild bool `json:"is_prebuild,omitempty"`
2468-
LogLevel string `json:"log_level,omitempty"`
2465+
WorkspaceBuildID uuid.UUID `json:"workspace_build_id"`
2466+
DryRun bool `json:"dry_run"`
2467+
IsPrebuild bool `json:"is_prebuild,omitempty"`
2468+
PrebuildClaimedByUser uuid.UUID `json:"prebuild_claimed_by,omitempty"`
2469+
// RunningWorkspaceAgentID is *only* used for prebuilds. We pass it down when we want to rebuild a prebuilt workspace
2470+
// but not generate a new agent token. The provisionerdserver will retrieve this token and push it down to
2471+
// the provisioner (and ultimately to the `coder_agent` resource in the Terraform provider) where it will be
2472+
// reused. Context: the agent token is often used in immutable attributes of workspace resource (e.g. VM/container)
2473+
// to initialize the agent, so if that value changes it will necessitate a replacement of that resource, thus
2474+
// obviating the whole point of the prebuild.
2475+
RunningWorkspaceAgentID uuid.UUID `json:"running_workspace_agent_id"`
2476+
LogLevel string `json:"log_level,omitempty"`
24692477
}
24702478

24712479
// TemplateVersionDryRunJob is the payload for the "template_version_dry_run" job type.

‎coderd/workspaces.go

+94-22
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"encoding/json"
77
"errors"
88
"fmt"
9+
"github.com/coder/coder/v2/coderd/prebuilds"
910
"net/http"
1011
"slices"
1112
"strconv"
@@ -635,34 +636,71 @@ func createWorkspace(
635636
provisionerJob *database.ProvisionerJob
636637
workspaceBuild *database.WorkspaceBuild
637638
provisionerDaemons []database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow
639+
640+
runningWorkspaceAgentID uuid.UUID
638641
)
642+
643+
prebuilds := (*api.PrebuildsClaimer.Load()).(prebuilds.Claimer)
644+
639645
err = api.Database.InTx(func(db database.Store) error {
640-
now := dbtime.Now()
641-
// Workspaces are created without any versions.
642-
minimumWorkspace, err := db.InsertWorkspace(ctx, database.InsertWorkspaceParams{
643-
ID: uuid.New(),
644-
CreatedAt: now,
645-
UpdatedAt: now,
646-
OwnerID: owner.ID,
647-
OrganizationID: template.OrganizationID,
648-
TemplateID: template.ID,
649-
Name: req.Name,
650-
AutostartSchedule: dbAutostartSchedule,
651-
NextStartAt: nextStartAt,
652-
Ttl: dbTTL,
653-
// The workspaces page will sort by last used at, and it's useful to
654-
// have the newly created workspace at the top of the list!
655-
LastUsedAt: dbtime.Now(),
656-
AutomaticUpdates: dbAU,
657-
})
658-
if err != nil {
659-
return xerrors.Errorf("insert workspace: %w", err)
646+
var (
647+
workspaceID uuid.UUID
648+
claimedWorkspace *database.Workspace
649+
)
650+
651+
// If a template preset was chosen, try claim a prebuild.
652+
if req.TemplateVersionPresetID != uuid.Nil {
653+
// Try and claim an eligible prebuild, if available.
654+
claimedWorkspace, err = claimPrebuild(ctx, prebuilds, db, api.Logger, req, owner)
655+
if err != nil {
656+
return xerrors.Errorf("claim prebuild: %w", err)
657+
}
658+
}
659+
660+
// No prebuild found; regular flow.
661+
if claimedWorkspace == nil {
662+
now := dbtime.Now()
663+
// Workspaces are created without any versions.
664+
minimumWorkspace, err := db.InsertWorkspace(ctx, database.InsertWorkspaceParams{
665+
ID: uuid.New(),
666+
CreatedAt: now,
667+
UpdatedAt: now,
668+
OwnerID: owner.ID,
669+
OrganizationID: template.OrganizationID,
670+
TemplateID: template.ID,
671+
Name: req.Name,
672+
AutostartSchedule: dbAutostartSchedule,
673+
NextStartAt: nextStartAt,
674+
Ttl: dbTTL,
675+
// The workspaces page will sort by last used at, and it's useful to
676+
// have the newly created workspace at the top of the list!
677+
LastUsedAt: dbtime.Now(),
678+
AutomaticUpdates: dbAU,
679+
})
680+
if err != nil {
681+
return xerrors.Errorf("insert workspace: %w", err)
682+
}
683+
workspaceID = minimumWorkspace.ID
684+
} else {
685+
// Prebuild found!
686+
workspaceID = claimedWorkspace.ID
687+
initiatorID = prebuilds.Initiator()
688+
agents, err := db.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, claimedWorkspace.ID)
689+
if err != nil {
690+
// TODO: comment about best-effort, workspace can be restarted if this fails...
691+
api.Logger.Error(ctx, "failed to retrieve running agents of claimed prebuilt workspace",
692+
slog.F("workspace_id", claimedWorkspace.ID), slog.Error(err))
693+
}
694+
if len(agents) >= 1 {
695+
// TODO: handle multiple agents
696+
runningWorkspaceAgentID = agents[0].ID
697+
}
660698
}
661699

662700
// We have to refetch the workspace for the joined in fields.
663701
// TODO: We can use WorkspaceTable for the builder to not require
664702
// this extra fetch.
665-
workspace, err = db.GetWorkspaceByID(ctx, minimumWorkspace.ID)
703+
workspace, err = db.GetWorkspaceByID(ctx, workspaceID)
666704
if err != nil {
667705
return xerrors.Errorf("get workspace by ID: %w", err)
668706
}
@@ -672,10 +710,18 @@ func createWorkspace(
672710
Initiator(initiatorID).
673711
ActiveVersion().
674712
RichParameterValues(req.RichParameterValues).
675-
TemplateVersionPresetID(req.TemplateVersionPresetID)
713+
TemplateVersionPresetID(req.TemplateVersionPresetID).
714+
RunningWorkspaceAgentID(runningWorkspaceAgentID)
676715
if req.TemplateVersionID != uuid.Nil {
677716
builder = builder.VersionID(req.TemplateVersionID)
678717
}
718+
if req.TemplateVersionPresetID != uuid.Nil {
719+
builder = builder.TemplateVersionPresetID(req.TemplateVersionPresetID)
720+
}
721+
722+
if claimedWorkspace != nil {
723+
builder = builder.MarkPrebuildClaimedBy(owner.ID)
724+
}
679725

680726
workspaceBuild, provisionerJob, provisionerDaemons, err = builder.Build(
681727
ctx,
@@ -839,6 +885,32 @@ func requestTemplate(ctx context.Context, rw http.ResponseWriter, req codersdk.C
839885
return template, true
840886
}
841887

888+
func claimPrebuild(ctx context.Context, claimer prebuilds.Claimer, db database.Store, logger slog.Logger, req codersdk.CreateWorkspaceRequest, owner workspaceOwner) (*database.Workspace, error) {
889+
prebuildsCtx := dbauthz.AsPrebuildsOrchestrator(ctx)
890+
891+
// TODO: do we need a timeout here?
892+
claimCtx, cancel := context.WithTimeout(prebuildsCtx, time.Second*10)
893+
defer cancel()
894+
895+
claimedID, err := claimer.Claim(claimCtx, db, owner.ID, req.Name, req.TemplateVersionPresetID)
896+
if err != nil {
897+
// TODO: enhance this by clarifying whether this *specific* prebuild failed or whether there are none to claim.
898+
return nil, xerrors.Errorf("claim prebuild: %w", err)
899+
}
900+
901+
// No prebuild available.
902+
if claimedID == nil {
903+
return nil, nil
904+
}
905+
906+
lookup, err := db.GetWorkspaceByID(prebuildsCtx, *claimedID)
907+
if err != nil {
908+
logger.Error(ctx, "unable to find claimed workspace by ID", slog.Error(err), slog.F("claimed_prebuild_id", (*claimedID).String()))
909+
return nil, xerrors.Errorf("find claimed workspace by ID %q: %w", (*claimedID).String(), err)
910+
}
911+
return &lookup, err
912+
}
913+
842914
func (api *API) notifyWorkspaceCreated(
843915
ctx context.Context,
844916
receiverID uuid.UUID,

‎coderd/wsbuilder/wsbuilder.go

+26-4
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,9 @@ type Builder struct {
7575
parameterValues *[]string
7676
templateVersionPresetParameterValues []database.TemplateVersionPresetParameter
7777

78-
prebuild bool
78+
prebuild bool
79+
prebuildClaimedBy uuid.UUID
80+
runningWorkspaceAgentID uuid.UUID
7981

8082
verifyNoLegacyParametersOnce bool
8183
}
@@ -178,6 +180,19 @@ func (b Builder) MarkPrebuild() Builder {
178180
return b
179181
}
180182

183+
func (b Builder) MarkPrebuildClaimedBy(userID uuid.UUID) Builder {
184+
// nolint: revive
185+
b.prebuildClaimedBy = userID
186+
return b
187+
}
188+
189+
// RunningWorkspaceAgentID is only used for prebuilds; see the associated field in `provisionerdserver.WorkspaceProvisionJob`.
190+
func (b Builder) RunningWorkspaceAgentID(id uuid.UUID) Builder {
191+
// nolint: revive
192+
b.runningWorkspaceAgentID = id
193+
return b
194+
}
195+
181196
// SetLastWorkspaceBuildInTx prepopulates the Builder's cache with the last workspace build. This allows us
182197
// to avoid a repeated database query when the Builder's caller also needs the workspace build, e.g. auto-start &
183198
// auto-stop.
@@ -309,9 +324,11 @@ func (b *Builder) buildTx(authFunc func(action policy.Action, object rbac.Object
309324

310325
workspaceBuildID := uuid.New()
311326
input, err := json.Marshal(provisionerdserver.WorkspaceProvisionJob{
312-
WorkspaceBuildID: workspaceBuildID,
313-
LogLevel: b.logLevel,
314-
IsPrebuild: b.prebuild,
327+
WorkspaceBuildID: workspaceBuildID,
328+
LogLevel: b.logLevel,
329+
IsPrebuild: b.prebuild,
330+
PrebuildClaimedByUser: b.prebuildClaimedBy,
331+
RunningWorkspaceAgentID: b.runningWorkspaceAgentID,
315332
})
316333
if err != nil {
317334
return nil, nil, nil, BuildError{
@@ -624,6 +641,11 @@ func (b *Builder) findNewBuildParameterValue(name string) *codersdk.WorkspaceBui
624641
}
625642

626643
func (b *Builder) getLastBuildParameters() ([]database.WorkspaceBuildParameter, error) {
644+
// TODO: exclude preset params from this list instead of returning nothing?
645+
if b.prebuildClaimedBy != uuid.Nil {
646+
return nil, nil
647+
}
648+
627649
if b.lastBuildParameters != nil {
628650
return *b.lastBuildParameters, nil
629651
}

0 commit comments

Comments
 (0)