Skip to content

Commit fdabb8c

Browse files
dannykoppingdannykopping
committed
Very basic prebuild reassignment
Signed-off-by: Danny Kopping <danny@coder.com>
1 parent 9d5c663 commit fdabb8c

File tree

11 files changed

+164
-23
lines changed

11 files changed

+164
-23
lines changed

‎coderd/database/dbauthz/dbauthz.go

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1078,6 +1078,13 @@ func (q *querier) BulkMarkNotificationMessagesSent(ctx context.Context, arg data
10781078
return q.db.BulkMarkNotificationMessagesSent(ctx, arg)
10791079
}
10801080

1081+
func (q *querier) ClaimPrebuild(ctx context.Context, newOwnerID uuid.UUID) (uuid.UUID, error) {
1082+
if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceWorkspace); err != nil {
1083+
return uuid.Nil, err
1084+
}
1085+
return q.db.ClaimPrebuild(ctx, newOwnerID)
1086+
}
1087+
10811088
func (q *querier) CleanTailnetCoordinators(ctx context.Context) error {
10821089
if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceTailnetCoordinator); err != nil {
10831090
return err

‎coderd/database/dbmem/dbmem.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1585,6 +1585,10 @@ func (*FakeQuerier) BulkMarkNotificationMessagesSent(_ context.Context, arg data
15851585
return int64(len(arg.IDs)), nil
15861586
}
15871587

1588+
func (q *FakeQuerier) ClaimPrebuild(ctx context.Context, newOwnerID uuid.UUID) (uuid.UUID, error) {
1589+
panic("not implemented")
1590+
}
1591+
15881592
func (*FakeQuerier) CleanTailnetCoordinators(_ context.Context) error {
15891593
return ErrUnimplemented
15901594
}

‎coderd/database/dbmetrics/querymetrics.go

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/querier.go

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries.sql.go

Lines changed: 23 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/prebuilds.sql

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,3 +62,18 @@ FROM templates_with_prebuilds t
6262
LEFT JOIN prebuilds_in_progress pip ON pip.template_version_id = t.template_version_id
6363
GROUP BY t.using_active_version, t.template_id, t.template_version_id, p.count, p.ids,
6464
p.template_version_id, t.deleted, t.deprecated;
65+
66+
-- name: ClaimPrebuild :one
67+
UPDATE workspaces w
68+
SET owner_id = @new_owner_id::uuid, updated_at = NOW() -- TODO: annoying; having two input params breaks dbgen
69+
WHERE w.id IN (SELECT p.id
70+
FROM workspace_prebuilds p
71+
INNER JOIN workspace_latest_build b ON b.workspace_id = p.id
72+
INNER JOIN provisioner_jobs pj ON b.job_id = pj.id
73+
INNER JOIN templates t ON p.template_id = t.id
74+
WHERE (b.transition = 'start'::workspace_transition
75+
AND pj.job_status IN ('succeeded'::provisioner_job_status))
76+
AND b.template_version_id = t.active_version_id
77+
ORDER BY random()
78+
LIMIT 1 FOR UPDATE OF p SKIP LOCKED)
79+
RETURNING w.id;

‎coderd/prebuilds/claim.go

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
package prebuilds
2+
3+
import (
4+
"context"
5+
"github.com/coder/coder/v2/coderd/database"
6+
"github.com/google/uuid"
7+
"golang.org/x/xerrors"
8+
)
9+
10+
func Claim(ctx context.Context, store database.Store, userID uuid.UUID) (*uuid.UUID, error) {
11+
var prebuildID *uuid.UUID
12+
err := store.InTx(func(db database.Store) error {
13+
// TODO: do we need this?
14+
//// Ensure no other replica can claim a prebuild for this user simultaneously.
15+
//err := store.AcquireLock(ctx, database.GenLockID(fmt.Sprintf("prebuild-user-claim-%s", userID.String())))
16+
//if err != nil {
17+
// return xerrors.Errorf("acquire claim lock for user %q: %w", userID.String(), err)
18+
//}
19+
20+
id, err := db.ClaimPrebuild(ctx, userID)
21+
if err != nil {
22+
return xerrors.Errorf("claim prebuild for user %q: %w", userID.String(), err)
23+
}
24+
25+
if id != uuid.Nil {
26+
prebuildID = &id
27+
}
28+
29+
return nil
30+
}, &database.TxOptions{
31+
TxIdentifier: "prebuild-claim",
32+
})
33+
34+
return prebuildID, err
35+
}

‎coderd/prebuilds/controller.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -376,5 +376,5 @@ func generateName() (string, error) {
376376
}
377377

378378
// Encode the bytes to Base32 (A-Z2-7), strip any '=' padding
379-
return fmt.Sprintf("prebuild-%s", base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(b)), nil
379+
return fmt.Sprintf("prebuild-%s", strings.ToLower(base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(b))), nil
380380
}

‎coderd/workspaces.go

Lines changed: 67 additions & 20 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"
@@ -628,32 +629,78 @@ func createWorkspace(
628629
provisionerDaemons []database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow
629630
)
630631
err = api.Database.InTx(func(db database.Store) error {
632+
var claimedWorkspace *database.Workspace
633+
634+
// TODO: implement matching logic
635+
if true {
636+
//if req.ClaimPrebuildIfAvailable {
637+
// TODO: authz // Can't use existing profiles (i.e. AsSystemRestricted) because of dbauthz rules
638+
var ownerCtx = dbauthz.As(ctx, rbac.Subject{
639+
ID: "owner",
640+
Roles: rbac.RoleIdentifiers{rbac.RoleOwner()},
641+
Groups: []string{},
642+
Scope: rbac.ExpandableScope(rbac.ScopeAll),
643+
})
644+
645+
claimCtx, cancel := context.WithTimeout(ownerCtx, time.Second*10) // TODO: don't use elevated authz context
646+
defer cancel()
647+
648+
claimedID, err := prebuilds.Claim(claimCtx, db, owner.ID)
649+
if err != nil {
650+
// TODO: enhance this by clarifying whether this *specific* prebuild failed or whether there are none to claim.
651+
api.Logger.Error(ctx, "failed to claim a prebuild", slog.Error(err))
652+
goto regularPath
653+
}
654+
655+
if claimedID == nil {
656+
api.Logger.Warn(ctx, "no claimable prebuild available", slog.Error(err))
657+
goto regularPath
658+
}
659+
660+
lookup, err := api.Database.GetWorkspaceByID(ownerCtx, *claimedID) // TODO: don't use elevated authz context
661+
if err != nil {
662+
api.Logger.Warn(ctx, "unable to find claimed workspace by ID", slog.Error(err), slog.F("claimed_prebuild_id", (*claimedID).String()))
663+
goto regularPath
664+
}
665+
666+
claimedWorkspace = &lookup
667+
}
668+
669+
regularPath:
631670
now := dbtime.Now()
632-
// Workspaces are created without any versions.
633-
minimumWorkspace, err := db.InsertWorkspace(ctx, database.InsertWorkspaceParams{
634-
ID: uuid.New(),
635-
CreatedAt: now,
636-
UpdatedAt: now,
637-
OwnerID: owner.ID,
638-
OrganizationID: template.OrganizationID,
639-
TemplateID: template.ID,
640-
Name: req.Name,
641-
AutostartSchedule: dbAutostartSchedule,
642-
NextStartAt: nextStartAt,
643-
Ttl: dbTTL,
644-
// The workspaces page will sort by last used at, and it's useful to
645-
// have the newly created workspace at the top of the list!
646-
LastUsedAt: dbtime.Now(),
647-
AutomaticUpdates: dbAU,
648-
})
649-
if err != nil {
650-
return xerrors.Errorf("insert workspace: %w", err)
671+
672+
var workspaceID uuid.UUID
673+
674+
if claimedWorkspace != nil {
675+
workspaceID = claimedWorkspace.ID
676+
} else {
677+
// Workspaces are created without any versions.
678+
minimumWorkspace, err := db.InsertWorkspace(ctx, database.InsertWorkspaceParams{
679+
ID: uuid.New(),
680+
CreatedAt: now,
681+
UpdatedAt: now,
682+
OwnerID: owner.ID,
683+
OrganizationID: template.OrganizationID,
684+
TemplateID: template.ID,
685+
Name: req.Name,
686+
AutostartSchedule: dbAutostartSchedule,
687+
NextStartAt: nextStartAt,
688+
Ttl: dbTTL,
689+
// The workspaces page will sort by last used at, and it's useful to
690+
// have the newly created workspace at the top of the list!
691+
LastUsedAt: dbtime.Now(),
692+
AutomaticUpdates: dbAU,
693+
})
694+
if err != nil {
695+
return xerrors.Errorf("insert workspace: %w", err)
696+
}
697+
workspaceID = minimumWorkspace.ID
651698
}
652699

653700
// We have to refetch the workspace for the joined in fields.
654701
// TODO: We can use WorkspaceTable for the builder to not require
655702
// this extra fetch.
656-
workspace, err = db.GetWorkspaceByID(ctx, minimumWorkspace.ID)
703+
workspace, err = db.GetWorkspaceByID(ctx, workspaceID)
657704
if err != nil {
658705
return xerrors.Errorf("get workspace by ID: %w", err)
659706
}

‎codersdk/organizations.go

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -207,8 +207,9 @@ type CreateWorkspaceRequest struct {
207207
TTLMillis *int64 `json:"ttl_ms,omitempty"`
208208
// RichParameterValues allows for additional parameters to be provided
209209
// during the initial provision.
210-
RichParameterValues []WorkspaceBuildParameter `json:"rich_parameter_values,omitempty"`
211-
AutomaticUpdates AutomaticUpdates `json:"automatic_updates,omitempty"`
210+
RichParameterValues []WorkspaceBuildParameter `json:"rich_parameter_values,omitempty"`
211+
AutomaticUpdates AutomaticUpdates `json:"automatic_updates,omitempty"`
212+
ClaimPrebuildIfAvailable bool `json:"claim_prebuild_if_available,omitempty"`
212213
}
213214

214215
func (c *Client) OrganizationByName(ctx context.Context, name string) (Organization, error) {

‎site/src/api/typesGenerated.ts

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)