diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index edc175ada5..d779490785 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -1078,6 +1078,13 @@ func (q *querier) BulkMarkNotificationMessagesSent(ctx context.Context, arg data return q.db.BulkMarkNotificationMessagesSent(ctx, arg) } +func (q *querier) ClaimPrebuild(ctx context.Context, newOwnerID uuid.UUID) (uuid.UUID, error) { + if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceWorkspace); err != nil { + return uuid.Nil, err + } + return q.db.ClaimPrebuild(ctx, newOwnerID) +} + func (q *querier) CleanTailnetCoordinators(ctx context.Context) error { if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceTailnetCoordinator); err != nil { return err diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go index 95a8f312f8..d3a0be3fe0 100644 --- a/coderd/database/dbmem/dbmem.go +++ b/coderd/database/dbmem/dbmem.go @@ -1585,6 +1585,10 @@ func (*FakeQuerier) BulkMarkNotificationMessagesSent(_ context.Context, arg data return int64(len(arg.IDs)), nil } +func (q *FakeQuerier) ClaimPrebuild(ctx context.Context, newOwnerID uuid.UUID) (uuid.UUID, error) { + panic("not implemented") +} + func (*FakeQuerier) CleanTailnetCoordinators(_ context.Context) error { return ErrUnimplemented } diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 82c1f09456..fd6a069c88 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -147,6 +147,13 @@ func (m queryMetricsStore) BulkMarkNotificationMessagesSent(ctx context.Context, return r0, r1 } +func (m queryMetricsStore) ClaimPrebuild(ctx context.Context, newOwnerID uuid.UUID) (uuid.UUID, error) { + start := time.Now() + r0, r1 := m.s.ClaimPrebuild(ctx, newOwnerID) + m.queryLatencies.WithLabelValues("ClaimPrebuild").Observe(time.Since(start).Seconds()) + return r0, r1 +} + func (m queryMetricsStore) CleanTailnetCoordinators(ctx context.Context) error { start := time.Now() err := m.s.CleanTailnetCoordinators(ctx) diff --git a/coderd/database/querier.go b/coderd/database/querier.go index cbd242a223..bea602af7c 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -60,6 +60,7 @@ type sqlcQuerier interface { BatchUpdateWorkspaceNextStartAt(ctx context.Context, arg BatchUpdateWorkspaceNextStartAtParams) error BulkMarkNotificationMessagesFailed(ctx context.Context, arg BulkMarkNotificationMessagesFailedParams) (int64, error) BulkMarkNotificationMessagesSent(ctx context.Context, arg BulkMarkNotificationMessagesSentParams) (int64, error) + ClaimPrebuild(ctx context.Context, newOwnerID uuid.UUID) (uuid.UUID, error) CleanTailnetCoordinators(ctx context.Context) error CleanTailnetLostPeers(ctx context.Context) error CleanTailnetTunnels(ctx context.Context) error diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 0fb72e2610..d9d69cb6f0 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -5395,6 +5395,29 @@ func (q *sqlQuerier) GetParameterSchemasByJobID(ctx context.Context, jobID uuid. return items, nil } +const claimPrebuild = `-- name: ClaimPrebuild :one +UPDATE workspaces w +SET owner_id = $1::uuid, updated_at = NOW() -- TODO: annoying; having two input params breaks dbgen +WHERE w.id IN (SELECT p.id + FROM workspace_prebuilds p + INNER JOIN workspace_latest_build b ON b.workspace_id = p.id + INNER JOIN provisioner_jobs pj ON b.job_id = pj.id + INNER JOIN templates t ON p.template_id = t.id + WHERE (b.transition = 'start'::workspace_transition + AND pj.job_status IN ('succeeded'::provisioner_job_status)) + AND b.template_version_id = t.active_version_id + ORDER BY random() + LIMIT 1 FOR UPDATE OF p SKIP LOCKED) +RETURNING w.id +` + +func (q *sqlQuerier) ClaimPrebuild(ctx context.Context, newOwnerID uuid.UUID) (uuid.UUID, error) { + row := q.db.QueryRowContext(ctx, claimPrebuild, newOwnerID) + var id uuid.UUID + err := row.Scan(&id) + return id, err +} + const getTemplatePrebuildState = `-- name: GetTemplatePrebuildState :many WITH -- All prebuilds currently running diff --git a/coderd/database/queries/prebuilds.sql b/coderd/database/queries/prebuilds.sql index 28684bebae..ca06ccb3c7 100644 --- a/coderd/database/queries/prebuilds.sql +++ b/coderd/database/queries/prebuilds.sql @@ -62,3 +62,18 @@ FROM templates_with_prebuilds t LEFT JOIN prebuilds_in_progress pip ON pip.template_version_id = t.template_version_id GROUP BY t.using_active_version, t.template_id, t.template_version_id, p.count, p.ids, p.template_version_id, t.deleted, t.deprecated; + +-- name: ClaimPrebuild :one +UPDATE workspaces w +SET owner_id = @new_owner_id::uuid, updated_at = NOW() -- TODO: annoying; having two input params breaks dbgen +WHERE w.id IN (SELECT p.id + FROM workspace_prebuilds p + INNER JOIN workspace_latest_build b ON b.workspace_id = p.id + INNER JOIN provisioner_jobs pj ON b.job_id = pj.id + INNER JOIN templates t ON p.template_id = t.id + WHERE (b.transition = 'start'::workspace_transition + AND pj.job_status IN ('succeeded'::provisioner_job_status)) + AND b.template_version_id = t.active_version_id + ORDER BY random() + LIMIT 1 FOR UPDATE OF p SKIP LOCKED) +RETURNING w.id; diff --git a/coderd/prebuilds/claim.go b/coderd/prebuilds/claim.go new file mode 100644 index 0000000000..e1610d57b8 --- /dev/null +++ b/coderd/prebuilds/claim.go @@ -0,0 +1,35 @@ +package prebuilds + +import ( + "context" + "github.com/coder/coder/v2/coderd/database" + "github.com/google/uuid" + "golang.org/x/xerrors" +) + +func Claim(ctx context.Context, store database.Store, userID 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) + //} + + id, err := db.ClaimPrebuild(ctx, userID) + if err != nil { + return xerrors.Errorf("claim prebuild for user %q: %w", userID.String(), err) + } + + if id != uuid.Nil { + prebuildID = &id + } + + return nil + }, &database.TxOptions{ + TxIdentifier: "prebuild-claim", + }) + + return prebuildID, err +} diff --git a/coderd/prebuilds/controller.go b/coderd/prebuilds/controller.go index 03069cc7fd..eafdd21f58 100644 --- a/coderd/prebuilds/controller.go +++ b/coderd/prebuilds/controller.go @@ -376,5 +376,5 @@ func generateName() (string, error) { } // Encode the bytes to Base32 (A-Z2-7), strip any '=' padding - return fmt.Sprintf("prebuild-%s", base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(b)), nil + return fmt.Sprintf("prebuild-%s", strings.ToLower(base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(b))), nil } diff --git a/coderd/workspaces.go b/coderd/workspaces.go index 7a64648033..3851247099 100644 --- a/coderd/workspaces.go +++ b/coderd/workspaces.go @@ -6,6 +6,7 @@ import ( "encoding/json" "errors" "fmt" + "github.com/coder/coder/v2/coderd/prebuilds" "net/http" "slices" "strconv" @@ -628,32 +629,78 @@ func createWorkspace( provisionerDaemons []database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow ) err = api.Database.InTx(func(db database.Store) error { + var claimedWorkspace *database.Workspace + + // TODO: implement matching logic + if true { + //if req.ClaimPrebuildIfAvailable { + // TODO: authz // Can't use existing profiles (i.e. AsSystemRestricted) because of dbauthz rules + var ownerCtx = dbauthz.As(ctx, rbac.Subject{ + ID: "owner", + Roles: rbac.RoleIdentifiers{rbac.RoleOwner()}, + Groups: []string{}, + Scope: rbac.ExpandableScope(rbac.ScopeAll), + }) + + claimCtx, cancel := context.WithTimeout(ownerCtx, time.Second*10) // TODO: don't use elevated authz context + defer cancel() + + claimedID, err := prebuilds.Claim(claimCtx, db, owner.ID) + if err != nil { + // TODO: enhance this by clarifying whether this *specific* prebuild failed or whether there are none to claim. + api.Logger.Error(ctx, "failed to claim a prebuild", slog.Error(err)) + goto regularPath + } + + if claimedID == nil { + api.Logger.Warn(ctx, "no claimable prebuild available", slog.Error(err)) + goto regularPath + } + + lookup, err := api.Database.GetWorkspaceByID(ownerCtx, *claimedID) // TODO: don't use elevated authz context + if err != nil { + api.Logger.Warn(ctx, "unable to find claimed workspace by ID", slog.Error(err), slog.F("claimed_prebuild_id", (*claimedID).String())) + goto regularPath + } + + claimedWorkspace = &lookup + } + + regularPath: 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 + + if claimedWorkspace != nil { + workspaceID = claimedWorkspace.ID + } else { + // 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 } // 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) } diff --git a/codersdk/organizations.go b/codersdk/organizations.go index 98afd98fed..d9741e6e6b 100644 --- a/codersdk/organizations.go +++ b/codersdk/organizations.go @@ -207,8 +207,9 @@ type CreateWorkspaceRequest struct { TTLMillis *int64 `json:"ttl_ms,omitempty"` // RichParameterValues allows for additional parameters to be provided // during the initial provision. - RichParameterValues []WorkspaceBuildParameter `json:"rich_parameter_values,omitempty"` - AutomaticUpdates AutomaticUpdates `json:"automatic_updates,omitempty"` + RichParameterValues []WorkspaceBuildParameter `json:"rich_parameter_values,omitempty"` + AutomaticUpdates AutomaticUpdates `json:"automatic_updates,omitempty"` + ClaimPrebuildIfAvailable bool `json:"claim_prebuild_if_available,omitempty"` } func (c *Client) OrganizationByName(ctx context.Context, name string) (Organization, error) { diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 50b45ccd4d..ae49fc58a6 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -452,6 +452,7 @@ export interface CreateWorkspaceRequest { readonly ttl_ms?: number; readonly rich_parameter_values?: readonly WorkspaceBuildParameter[]; readonly automatic_updates?: AutomaticUpdates; + readonly claim_prebuild_if_available?: boolean; } // From codersdk/deployment.go