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

Lines changed: 3 additions & 0 deletions
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

Lines changed: 33 additions & 0 deletions
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

Lines changed: 15 additions & 0 deletions
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

Lines changed: 12 additions & 4 deletions
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

Lines changed: 94 additions & 22 deletions
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

Lines changed: 26 additions & 4 deletions
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)
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