synced/coder
1
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2025-07-06 15:41:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

Very basic prebuild reassignment

Browse Source 
Signed-off-by: Danny Kopping <danny@coder.com>
This commit is contained in:
Danny Kopping
2025-01-29 08:15:07 +00:00
parent 9d5c6633de
commit fdabb8cf07
11 changed files with 164 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
coderd/database/dbauthz/dbauthz.go
View File

@ -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

4
coderd/database/dbmem/dbmem.go
View File

@ -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
}

7
coderd/database/dbmetrics/querymetrics.go
View File

@ -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)

1
coderd/database/querier.go
View File

@ -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

23
coderd/database/queries.sql.go
View File

@ -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

15
coderd/database/queries/prebuilds.sql
View File

@ -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;

35
coderd/prebuilds/claim.go Normal file
View File

@ -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
}

2
coderd/prebuilds/controller.go
View File

@ -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
}

87
coderd/workspaces.go
View File

@ -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)
}