Skip to content

Commit 118f12a

Browse files
evgeniy-scherbinadannykoppingEdwardAngertjaaydenhethanndickson
authored
feat: implement claiming of prebuilt workspaces (coder#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>
1 parent 25dacd3 commit 118f12a

File tree

8 files changed

+731
-29
lines changed

8 files changed

+731
-29
lines changed

‎coderd/coderd.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ import (
4545
"github.com/coder/coder/v2/coderd/entitlements"
4646
"github.com/coder/coder/v2/coderd/files"
4747
"github.com/coder/coder/v2/coderd/idpsync"
48+
"github.com/coder/coder/v2/coderd/prebuilds"
4849
"github.com/coder/coder/v2/coderd/runtimeconfig"
4950
"github.com/coder/coder/v2/coderd/webpush"
5051

@@ -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(),
@@ -1569,6 +1571,7 @@ type API struct {
15691571
AccessControlStore *atomic.Pointer[dbauthz.AccessControlStore]
15701572
PortSharer atomic.Pointer[portsharing.PortSharer]
15711573
FileCache files.Cache
1574+
PrebuildsClaimer atomic.Pointer[prebuilds.Claimer]
15721575

15731576
UpdatesProvider tailnet.WorkspaceUpdatesProvider
15741577

‎coderd/prebuilds/api.go

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,13 @@ package prebuilds
22

33
import (
44
"context"
5+
6+
"github.com/google/uuid"
7+
"golang.org/x/xerrors"
58
)
69

10+
var ErrNoClaimablePrebuiltWorkspaces = xerrors.New("no claimable prebuilt workspaces found")
11+
712
// ReconciliationOrchestrator manages the lifecycle of prebuild reconciliation.
813
// It runs a continuous loop to check and reconcile prebuild states, and can be stopped gracefully.
914
type ReconciliationOrchestrator interface {
@@ -25,3 +30,8 @@ type Reconciler interface {
2530
// in parallel, creating or deleting prebuilds as needed to reach their desired states.
2631
ReconcileAll(ctx context.Context) error
2732
}
33+
34+
type Claimer interface {
35+
Claim(ctx context.Context, userID uuid.UUID, name string, presetID uuid.UUID) (*uuid.UUID, error)
36+
Initiator() uuid.UUID
37+
}

‎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 (AGPLPrebuildClaimer) Claim(context.Context, uuid.UUID, string, uuid.UUID) (*uuid.UUID, error) {
42+
// Not entitled to claim prebuilds in AGPL version.
43+
return nil, ErrNoClaimablePrebuiltWorkspaces
44+
}
45+
46+
func (AGPLPrebuildClaimer) Initiator() uuid.UUID {
47+
return uuid.Nil
48+
}
49+
50+
var DefaultClaimer Claimer = AGPLPrebuildClaimer{}

‎coderd/provisionerdserver/provisionerdserver.go

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2471,10 +2471,11 @@ type TemplateVersionImportJob struct {
24712471

24722472
// WorkspaceProvisionJob is the payload for the "workspace_provision" job type.
24732473
type WorkspaceProvisionJob struct {
2474-
WorkspaceBuildID uuid.UUID `json:"workspace_build_id"`
2475-
DryRun bool `json:"dry_run"`
2476-
IsPrebuild bool `json:"is_prebuild,omitempty"`
2477-
LogLevel string `json:"log_level,omitempty"`
2474+
WorkspaceBuildID uuid.UUID `json:"workspace_build_id"`
2475+
DryRun bool `json:"dry_run"`
2476+
IsPrebuild bool `json:"is_prebuild,omitempty"`
2477+
PrebuildClaimedByUser uuid.UUID `json:"prebuild_claimed_by,omitempty"`
2478+
LogLevel string `json:"log_level,omitempty"`
24782479
}
24792480

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

‎coderd/workspaces.go

Lines changed: 69 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import (
1818
"golang.org/x/xerrors"
1919

2020
"cdr.dev/slog"
21+
2122
"github.com/coder/coder/v2/agent/proto"
2223
"github.com/coder/coder/v2/coderd/audit"
2324
"github.com/coder/coder/v2/coderd/database"
@@ -28,6 +29,7 @@ import (
2829
"github.com/coder/coder/v2/coderd/httpapi"
2930
"github.com/coder/coder/v2/coderd/httpmw"
3031
"github.com/coder/coder/v2/coderd/notifications"
32+
"github.com/coder/coder/v2/coderd/prebuilds"
3133
"github.com/coder/coder/v2/coderd/rbac"
3234
"github.com/coder/coder/v2/coderd/rbac/policy"
3335
"github.com/coder/coder/v2/coderd/schedule"
@@ -636,33 +638,57 @@ func createWorkspace(
636638
workspaceBuild *database.WorkspaceBuild
637639
provisionerDaemons []database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow
638640
)
641+
639642
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)
643+
var (
644+
workspaceID uuid.UUID
645+
claimedWorkspace *database.Workspace
646+
prebuildsClaimer = *api.PrebuildsClaimer.Load()
647+
)
648+
649+
// If a template preset was chosen, try claim a prebuilt workspace.
650+
if req.TemplateVersionPresetID != uuid.Nil {
651+
// Try and claim an eligible prebuild, if available.
652+
claimedWorkspace, err = claimPrebuild(ctx, prebuildsClaimer, db, api.Logger, req, owner)
653+
if err != nil && !errors.Is(err, prebuilds.ErrNoClaimablePrebuiltWorkspaces) {
654+
return xerrors.Errorf("claim prebuild: %w", err)
655+
}
656+
}
657+
658+
// No prebuild found; regular flow.
659+
if claimedWorkspace == nil {
660+
now := dbtime.Now()
661+
// Workspaces are created without any versions.
662+
minimumWorkspace, err := db.InsertWorkspace(ctx, database.InsertWorkspaceParams{
663+
ID: uuid.New(),
664+
CreatedAt: now,
665+
UpdatedAt: now,
666+
OwnerID: owner.ID,
667+
OrganizationID: template.OrganizationID,
668+
TemplateID: template.ID,
669+
Name: req.Name,
670+
AutostartSchedule: dbAutostartSchedule,
671+
NextStartAt: nextStartAt,
672+
Ttl: dbTTL,
673+
// The workspaces page will sort by last used at, and it's useful to
674+
// have the newly created workspace at the top of the list!
675+
LastUsedAt: dbtime.Now(),
676+
AutomaticUpdates: dbAU,
677+
})
678+
if err != nil {
679+
return xerrors.Errorf("insert workspace: %w", err)
680+
}
681+
workspaceID = minimumWorkspace.ID
682+
} else {
683+
// Prebuild found!
684+
workspaceID = claimedWorkspace.ID
685+
initiatorID = prebuildsClaimer.Initiator()
660686
}
661687

662688
// We have to refetch the workspace for the joined in fields.
663689
// TODO: We can use WorkspaceTable for the builder to not require
664690
// this extra fetch.
665-
workspace, err = db.GetWorkspaceByID(ctx, minimumWorkspace.ID)
691+
workspace, err = db.GetWorkspaceByID(ctx, workspaceID)
666692
if err != nil {
667693
return xerrors.Errorf("get workspace by ID: %w", err)
668694
}
@@ -676,6 +702,13 @@ func createWorkspace(
676702
if req.TemplateVersionID != uuid.Nil {
677703
builder = builder.VersionID(req.TemplateVersionID)
678704
}
705+
if req.TemplateVersionPresetID != uuid.Nil {
706+
builder = builder.TemplateVersionPresetID(req.TemplateVersionPresetID)
707+
}
708+
if claimedWorkspace != nil {
709+
builder = builder.MarkPrebuildClaimedBy(owner.ID)
710+
}
711+
679712
if req.EnableDynamicParameters && api.Experiments.Enabled(codersdk.ExperimentDynamicParameters) {
680713
builder = builder.UsingDynamicParameters()
681714
}
@@ -842,6 +875,21 @@ func requestTemplate(ctx context.Context, rw http.ResponseWriter, req codersdk.C
842875
return template, true
843876
}
844877

878+
func claimPrebuild(ctx context.Context, claimer prebuilds.Claimer, db database.Store, logger slog.Logger, req codersdk.CreateWorkspaceRequest, owner workspaceOwner) (*database.Workspace, error) {
879+
claimedID, err := claimer.Claim(ctx, owner.ID, req.Name, req.TemplateVersionPresetID)
880+
if err != nil {
881+
// TODO: enhance this by clarifying whether this *specific* prebuild failed or whether there are none to claim.
882+
return nil, xerrors.Errorf("claim prebuild: %w", err)
883+
}
884+
885+
lookup, err := db.GetWorkspaceByID(ctx, *claimedID)
886+
if err != nil {
887+
logger.Error(ctx, "unable to find claimed workspace by ID", slog.Error(err), slog.F("claimed_prebuild_id", claimedID.String()))
888+
return nil, xerrors.Errorf("find claimed workspace by ID %q: %w", claimedID.String(), err)
889+
}
890+
return &lookup, nil
891+
}
892+
845893
func (api *API) notifyWorkspaceCreated(
846894
ctx context.Context,
847895
receiverID uuid.UUID,

‎coderd/wsbuilder/wsbuilder.go

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,8 @@ type Builder struct {
7676
parameterValues *[]string
7777
templateVersionPresetParameterValues []database.TemplateVersionPresetParameter
7878

79-
prebuild bool
79+
prebuild bool
80+
prebuildClaimedBy uuid.UUID
8081

8182
verifyNoLegacyParametersOnce bool
8283
}
@@ -179,6 +180,12 @@ func (b Builder) MarkPrebuild() Builder {
179180
return b
180181
}
181182

183+
func (b Builder) MarkPrebuildClaimedBy(userID uuid.UUID) Builder {
184+
// nolint: revive
185+
b.prebuildClaimedBy = userID
186+
return b
187+
}
188+
182189
func (b Builder) UsingDynamicParameters() Builder {
183190
b.dynamicParametersEnabled = true
184191
return b
@@ -315,9 +322,10 @@ func (b *Builder) buildTx(authFunc func(action policy.Action, object rbac.Object
315322

316323
workspaceBuildID := uuid.New()
317324
input, err := json.Marshal(provisionerdserver.WorkspaceProvisionJob{
318-
WorkspaceBuildID: workspaceBuildID,
319-
LogLevel: b.logLevel,
320-
IsPrebuild: b.prebuild,
325+
WorkspaceBuildID: workspaceBuildID,
326+
LogLevel: b.logLevel,
327+
IsPrebuild: b.prebuild,
328+
PrebuildClaimedByUser: b.prebuildClaimedBy,
321329
})
322330
if err != nil {
323331
return nil, nil, nil, BuildError{

‎enterprise/coderd/prebuilds/claim.go

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
package prebuilds
2+
3+
import (
4+
"context"
5+
"database/sql"
6+
"errors"
7+
8+
"github.com/google/uuid"
9+
"golang.org/x/xerrors"
10+
11+
"github.com/coder/coder/v2/coderd/database"
12+
"github.com/coder/coder/v2/coderd/prebuilds"
13+
)
14+
15+
type EnterpriseClaimer struct {
16+
store database.Store
17+
}
18+
19+
func NewEnterpriseClaimer(store database.Store) *EnterpriseClaimer {
20+
return &EnterpriseClaimer{
21+
store: store,
22+
}
23+
}
24+
25+
func (c EnterpriseClaimer) Claim(
26+
ctx context.Context,
27+
userID uuid.UUID,
28+
name string,
29+
presetID uuid.UUID,
30+
) (*uuid.UUID, error) {
31+
result, err := c.store.ClaimPrebuiltWorkspace(ctx, database.ClaimPrebuiltWorkspaceParams{
32+
NewUserID: userID,
33+
NewName: name,
34+
PresetID: presetID,
35+
})
36+
if err != nil {
37+
switch {
38+
// No eligible prebuilds found
39+
case errors.Is(err, sql.ErrNoRows):
40+
return nil, prebuilds.ErrNoClaimablePrebuiltWorkspaces
41+
default:
42+
return nil, xerrors.Errorf("claim prebuild for user %q: %w", userID.String(), err)
43+
}
44+
}
45+
46+
return &result.ID, nil
47+
}
48+
49+
func (EnterpriseClaimer) Initiator() uuid.UUID {
50+
return prebuilds.SystemUserID
51+
}
52+
53+
var _ prebuilds.Claimer = &EnterpriseClaimer{}

0 commit comments

Comments
 (0)