Skip to content

Commit 8849ad3

Browse files
dannykoppingdannykopping
committed
Refactor
Signed-off-by: Danny Kopping <danny@coder.com>
1 parent b151e1b commit 8849ad3

File tree

3 files changed

+57
-56
lines changed

3 files changed

+57
-56
lines changed

‎coderd/provisionerdserver/provisionerdserver.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -1883,7 +1883,7 @@ func InsertWorkspacePresetAndParameters(ctx context.Context, db database.Store,
18831883
}
18841884

18851885
if protoPreset.Prebuild != nil {
1886-
_, err := db.InsertPresetPrebuild(ctx, database.InsertPresetPrebuildParams{
1886+
_, err := tx.InsertPresetPrebuild(ctx, database.InsertPresetPrebuildParams{
18871887
ID: uuid.New(),
18881888
PresetID: dbPreset.ID,
18891889
DesiredInstances: protoPreset.Prebuild.Instances,

‎coderd/workspaces.go

+56-53
Original file line numberDiff line numberDiff line change
@@ -633,62 +633,17 @@ func createWorkspace(
633633
runningWorkspaceAgentID uuid.UUID
634634
)
635635
err = api.Database.InTx(func(db database.Store) error {
636-
var claimedWorkspace *database.Workspace
637-
638-
// TODO: implement matching logic
639-
if true {
640-
//if req.ClaimPrebuildIfAvailable {
641-
// TODO: authz // Can't use existing profiles (i.e. AsSystemRestricted) because of dbauthz rules
642-
var ownerCtx = dbauthz.As(ctx, rbac.Subject{
643-
ID: "owner",
644-
Roles: rbac.RoleIdentifiers{rbac.RoleOwner()},
645-
Groups: []string{},
646-
Scope: rbac.ExpandableScope(rbac.ScopeAll),
647-
})
648-
649-
claimCtx, cancel := context.WithTimeout(ownerCtx, time.Second*10) // TODO: don't use elevated authz context
650-
defer cancel()
651-
652-
// TODO: pass down rich params for matching
653-
claimedID, err := prebuilds.Claim(claimCtx, db, owner.ID, req.Name)
654-
if err != nil {
655-
// TODO: enhance this by clarifying whether this *specific* prebuild failed or whether there are none to claim.
656-
api.Logger.Error(ctx, "failed to claim a prebuild", slog.Error(err))
657-
goto regularPath
658-
}
659-
660-
if claimedID == nil {
661-
api.Logger.Warn(ctx, "no claimable prebuild available", slog.Error(err))
662-
goto regularPath
663-
}
664-
665-
lookup, err := api.Database.GetWorkspaceByID(ownerCtx, *claimedID) // TODO: don't use elevated authz context
666-
if err != nil {
667-
api.Logger.Warn(ctx, "unable to find claimed workspace by ID", slog.Error(err), slog.F("claimed_prebuild_id", (*claimedID).String()))
668-
goto regularPath
669-
}
670-
claimedWorkspace = &lookup
636+
var workspaceID uuid.UUID
671637

672-
agents, err := api.Database.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ownerCtx, claimedWorkspace.ID)
673-
if err != nil {
674-
api.Logger.Error(ctx, "failed to retrieve running agents of claimed prebuilt workspace",
675-
slog.F("workspace_id", claimedWorkspace.ID), slog.Error(err))
676-
}
677-
if len(agents) >= 1 {
678-
// TODO: handle multiple agents
679-
runningWorkspaceAgentID = agents[0].ID
680-
}
638+
// Try and claim an eligible prebuild, if available.
639+
claimedWorkspace, err := claimPrebuild(ctx, db, api.Logger, req, owner)
640+
if err != nil {
641+
return xerrors.Errorf("claim prebuild: %w", err)
681642
}
682643

683-
regularPath:
684-
now := dbtime.Now()
685-
686-
var workspaceID uuid.UUID
687-
688-
if claimedWorkspace != nil {
689-
workspaceID = claimedWorkspace.ID
690-
initiatorID = prebuilds.PrebuildOwnerUUID
691-
} else {
644+
// No prebuild found; regular flow.
645+
if claimedWorkspace == nil {
646+
now := dbtime.Now()
692647
// Workspaces are created without any versions.
693648
minimumWorkspace, err := db.InsertWorkspace(ctx, database.InsertWorkspaceParams{
694649
ID: uuid.New(),
@@ -710,6 +665,20 @@ func createWorkspace(
710665
return xerrors.Errorf("insert workspace: %w", err)
711666
}
712667
workspaceID = minimumWorkspace.ID
668+
} else {
669+
// Prebuild found!
670+
workspaceID = claimedWorkspace.ID
671+
initiatorID = prebuilds.PrebuildOwnerUUID
672+
673+
agents, err := api.Database.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, claimedWorkspace.ID)
674+
if err != nil {
675+
api.Logger.Error(ctx, "failed to retrieve running agents of claimed prebuilt workspace",
676+
slog.F("workspace_id", claimedWorkspace.ID), slog.Error(err))
677+
}
678+
if len(agents) >= 1 {
679+
// TODO: handle multiple agents
680+
runningWorkspaceAgentID = agents[0].ID
681+
}
713682
}
714683

715684
// We have to refetch the workspace for the joined in fields.
@@ -828,6 +797,40 @@ func createWorkspace(
828797
httpapi.Write(ctx, rw, http.StatusCreated, w)
829798
}
830799

800+
func claimPrebuild(ctx context.Context, db database.Store, logger slog.Logger, req codersdk.CreateWorkspaceRequest, owner workspaceOwner) (*database.Workspace, error) {
801+
// TODO: authz // Can't use existing profiles (i.e. AsSystemRestricted) because of dbauthz rules
802+
var ownerCtx = dbauthz.As(ctx, rbac.Subject{
803+
ID: "owner",
804+
Roles: rbac.RoleIdentifiers{rbac.RoleOwner()},
805+
Groups: []string{},
806+
Scope: rbac.ExpandableScope(rbac.ScopeAll),
807+
})
808+
809+
// TODO: do we need a timeout here?
810+
claimCtx, cancel := context.WithTimeout(ownerCtx, time.Second*10) // TODO: don't use elevated authz context
811+
defer cancel()
812+
813+
// TODO: implement matching logic
814+
// TODO: pass down rich params for matching
815+
claimedID, err := prebuilds.Claim(claimCtx, db, owner.ID, req.Name)
816+
if err != nil {
817+
// TODO: enhance this by clarifying whether this *specific* prebuild failed or whether there are none to claim.
818+
return nil, xerrors.Errorf("claim prebuild: %w", err)
819+
}
820+
821+
// No prebuild available.
822+
if claimedID == nil {
823+
return nil, nil
824+
}
825+
826+
lookup, err := db.GetWorkspaceByID(ownerCtx, *claimedID) // TODO: don't use elevated authz context
827+
if err != nil {
828+
logger.Error(ctx, "unable to find claimed workspace by ID", slog.Error(err), slog.F("claimed_prebuild_id", (*claimedID).String()))
829+
return nil, xerrors.Errorf("find claimed workspace by ID %q: %w", (*claimedID).String(), err)
830+
}
831+
return &lookup, err
832+
}
833+
831834
func (api *API) notifyWorkspaceCreated(
832835
ctx context.Context,
833836
receiverID uuid.UUID,

‎codersdk/agentsdk/agentsdk.go

-2
Original file line numberDiff line numberDiff line change
@@ -698,8 +698,6 @@ func (c *Client) WaitForReinit(ctx context.Context, events chan<- Reinitializati
698698
if err != nil {
699699
return xerrors.Errorf("failed to read server-sent event: %w", err)
700700
}
701-
// TODO: remove
702-
fmt.Printf("RECEIVED SSE EVENT: %s\n", sse.Type)
703701
if sse.Type != codersdk.ServerSentEventTypeData {
704702
continue
705703
}

0 commit comments

Comments
 (0)