Skip to content

Commit a8f87c2

Browse files
mafredrimafredri
authored
feat(coderd): implement task to app linking (#20237)
This change adds workspace build/agent/app linking to tasks and wires it into `wsbuilder` and `provisionerdserver`. Closes coder/internal#948 Closes #20212 Closes #19773
1 parent 5dc57da commit a8f87c2

4 files changed

Lines changed: 225 additions & 3 deletions

File tree

‎coderd/provisionerdserver/provisionerdserver.go‎

Lines changed: 78 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1964,18 +1964,41 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
19641964
}
19651965

19661966
appIDs := make([]string, 0)
1967+
agentIDByAppID := make(map[string]uuid.UUID)
19671968
agentTimeouts := make(map[time.Duration]bool) // A set of agent timeouts.
19681969
// This could be a bulk insert to improve performance.
19691970
for _, protoResource := range jobType.WorkspaceBuild.Resources {
1970-
for _, protoAgent := range protoResource.Agents {
1971+
for _, protoAgent := range protoResource.GetAgents() {
1972+
if protoAgent == nil {
1973+
continue
1974+
}
1975+
// By default InsertWorkspaceResource ignores the protoAgent.Id
1976+
// and generates a new one, but we will insert these using the
1977+
// InsertWorkspaceResourceWithAgentIDsFromProto option so that
1978+
// we can properly map agent IDs to app IDs. This is needed for
1979+
// task linking.
1980+
agentID := uuid.New()
1981+
protoAgent.Id = agentID.String()
1982+
19711983
dur := time.Duration(protoAgent.GetConnectionTimeoutSeconds()) * time.Second
19721984
agentTimeouts[dur] = true
19731985
for _, app := range protoAgent.GetApps() {
19741986
appIDs = append(appIDs, app.GetId())
1987+
agentIDByAppID[app.GetId()] = agentID
19751988
}
19761989
}
19771990

1978-
err = InsertWorkspaceResource(ctx, db, job.ID, workspaceBuild.Transition, protoResource, telemetrySnapshot)
1991+
err = InsertWorkspaceResource(
1992+
ctx,
1993+
db,
1994+
job.ID,
1995+
workspaceBuild.Transition,
1996+
protoResource,
1997+
telemetrySnapshot,
1998+
// Ensure that the agent IDs we set previously
1999+
// are written to the database.
2000+
InsertWorkspaceResourceWithAgentIDsFromProto(),
2001+
)
19792002
if err != nil {
19802003
return xerrors.Errorf("insert provisioner job: %w", err)
19812004
}
@@ -1987,6 +2010,7 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
19872010
}
19882011

19892012
var taskAppID uuid.NullUUID
2013+
var taskAgentID uuid.NullUUID
19902014
var hasAITask bool
19912015
var warnUnknownTaskAppID bool
19922016
if tasks := jobType.WorkspaceBuild.GetAiTasks(); len(tasks) > 0 {
@@ -2014,6 +2038,9 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
20142038
}
20152039

20162040
taskAppID = uuid.NullUUID{UUID: id, Valid: true}
2041+
2042+
agentID, ok := agentIDByAppID[appID]
2043+
taskAgentID = uuid.NullUUID{UUID: agentID, Valid: ok}
20172044
}
20182045

20192046
// This is a hacky workaround for the issue with tasks 'disappearing' on stop:
@@ -2108,6 +2135,27 @@ func (s *server) completeWorkspaceBuildJob(ctx context.Context, job database.Pro
21082135
}
21092136
}
21102137

2138+
if task, err := db.GetTaskByWorkspaceID(ctx, workspace.ID); err == nil {
2139+
// Irrespective of whether the agent or sidebar app is present,
2140+
// perform the upsert to ensure a link between the task and
2141+
// workspace build. Linking the task to the build is typically
2142+
// already established by wsbuilder.
2143+
_, err = db.UpsertTaskWorkspaceApp(
2144+
ctx,
2145+
database.UpsertTaskWorkspaceAppParams{
2146+
TaskID: task.ID,
2147+
WorkspaceBuildNumber: workspaceBuild.BuildNumber,
2148+
WorkspaceAgentID: taskAgentID,
2149+
WorkspaceAppID: taskAppID,
2150+
},
2151+
)
2152+
if err != nil {
2153+
return xerrors.Errorf("upsert task workspace app: %w", err)
2154+
}
2155+
} else if !errors.Is(err, sql.ErrNoRows) {
2156+
return xerrors.Errorf("get task by workspace id: %w", err)
2157+
}
2158+
21112159
// Regardless of whether there is an AI task or not, update the field to indicate one way or the other since it
21122160
// always defaults to nil. ONLY if has_ai_task=true MUST ai_task_sidebar_app_id be set.
21132161
if err := db.UpdateWorkspaceBuildFlagsByID(ctx, database.UpdateWorkspaceBuildFlagsByIDParams{
@@ -2578,7 +2626,28 @@ func InsertWorkspacePresetAndParameters(ctx context.Context, db database.Store,
25782626
return nil
25792627
}
25802628

2581-
func InsertWorkspaceResource(ctx context.Context, db database.Store, jobID uuid.UUID, transition database.WorkspaceTransition, protoResource *sdkproto.Resource, snapshot *telemetry.Snapshot) error {
2629+
type insertWorkspaceResourceOptions struct {
2630+
useAgentIDsFromProto bool
2631+
}
2632+
2633+
// InsertWorkspaceResourceOption represents a functional option for
2634+
// InsertWorkspaceResource.
2635+
type InsertWorkspaceResourceOption func(*insertWorkspaceResourceOptions)
2636+
2637+
// InsertWorkspaceResourceWithAgentIDsFromProto allows inserting agents into the
2638+
// database using the agent IDs defined in the proto resource.
2639+
func InsertWorkspaceResourceWithAgentIDsFromProto() InsertWorkspaceResourceOption {
2640+
return func(opts *insertWorkspaceResourceOptions) {
2641+
opts.useAgentIDsFromProto = true
2642+
}
2643+
}
2644+
2645+
func InsertWorkspaceResource(ctx context.Context, db database.Store, jobID uuid.UUID, transition database.WorkspaceTransition, protoResource *sdkproto.Resource, snapshot *telemetry.Snapshot, opt ...InsertWorkspaceResourceOption) error {
2646+
opts := &insertWorkspaceResourceOptions{}
2647+
for _, o := range opt {
2648+
o(opts)
2649+
}
2650+
25822651
resource, err := db.InsertWorkspaceResource(ctx, database.InsertWorkspaceResourceParams{
25832652
ID: uuid.New(),
25842653
CreatedAt: dbtime.Now(),
@@ -2675,6 +2744,12 @@ func InsertWorkspaceResource(ctx context.Context, db database.Store, jobID uuid.
26752744
}
26762745

26772746
agentID := uuid.New()
2747+
if opts.useAgentIDsFromProto {
2748+
agentID, err = uuid.Parse(prAgent.Id)
2749+
if err != nil {
2750+
return xerrors.Errorf("invalid agent ID format; must be uuid: %w", err)
2751+
}
2752+
}
26782753
dbAgent, err := db.InsertWorkspaceAgent(ctx, database.InsertWorkspaceAgentParams{
26792754
ID: agentID,
26802755
ParentID: uuid.NullUUID{},

‎coderd/provisionerdserver/provisionerdserver_test.go‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2850,6 +2850,8 @@ func TestCompleteJob(t *testing.T) {
28502850
seedFunc func(context.Context, testing.TB, database.Store) error // If you need to insert other resources
28512851
transition database.WorkspaceTransition
28522852
input *proto.CompletedJob_WorkspaceBuild
2853+
isTask bool
2854+
expectTaskStatus database.TaskStatus
28532855
expectHasAiTask bool
28542856
expectUsageEvent bool
28552857
}
@@ -2862,6 +2864,7 @@ func TestCompleteJob(t *testing.T) {
28622864
input: &proto.CompletedJob_WorkspaceBuild{
28632865
// No AiTasks defined.
28642866
},
2867+
isTask: false,
28652868
expectHasAiTask: false,
28662869
expectUsageEvent: false,
28672870
},
@@ -2894,6 +2897,8 @@ func TestCompleteJob(t *testing.T) {
28942897
},
28952898
},
28962899
},
2900+
isTask: true,
2901+
expectTaskStatus: database.TaskStatusInitializing,
28972902
expectHasAiTask: true,
28982903
expectUsageEvent: true,
28992904
},
@@ -2912,6 +2917,8 @@ func TestCompleteJob(t *testing.T) {
29122917
},
29132918
},
29142919
},
2920+
isTask: true,
2921+
expectTaskStatus: database.TaskStatusInitializing,
29152922
expectHasAiTask: false,
29162923
expectUsageEvent: false,
29172924
},
@@ -2944,6 +2951,8 @@ func TestCompleteJob(t *testing.T) {
29442951
},
29452952
},
29462953
},
2954+
isTask: true,
2955+
expectTaskStatus: database.TaskStatusPaused,
29472956
expectHasAiTask: true,
29482957
expectUsageEvent: false,
29492958
},
@@ -2955,6 +2964,8 @@ func TestCompleteJob(t *testing.T) {
29552964
AiTasks: []*sdkproto.AITask{},
29562965
Resources: []*sdkproto.Resource{},
29572966
},
2967+
isTask: true,
2968+
expectTaskStatus: database.TaskStatusPaused,
29582969
expectHasAiTask: true,
29592970
expectUsageEvent: false,
29602971
},
@@ -2992,6 +3003,15 @@ func TestCompleteJob(t *testing.T) {
29923003
OwnerID: user.ID,
29933004
OrganizationID: pd.OrganizationID,
29943005
})
3006+
var taskTable database.TaskTable
3007+
if tc.isTask {
3008+
taskTable = dbgen.Task(t, db, database.TaskTable{
3009+
OwnerID: user.ID,
3010+
OrganizationID: pd.OrganizationID,
3011+
WorkspaceID: uuid.NullUUID{UUID: workspaceTable.ID, Valid: true},
3012+
TemplateVersionID: version.ID,
3013+
})
3014+
}
29953015

29963016
ctx := testutil.Context(t, testutil.WaitShort)
29973017
if tc.seedFunc != nil {
@@ -3060,6 +3080,12 @@ func TestCompleteJob(t *testing.T) {
30603080
require.True(t, build.HasAITask.Valid) // We ALWAYS expect a value to be set, therefore not nil, i.e. valid = true.
30613081
require.Equal(t, tc.expectHasAiTask, build.HasAITask.Bool)
30623082

3083+
if tc.isTask {
3084+
task, err := db.GetTaskByID(ctx, taskTable.ID)
3085+
require.NoError(t, err)
3086+
require.Equal(t, tc.expectTaskStatus, task.Status)
3087+
}
3088+
30633089
if tc.expectHasAiTask && build.Transition != database.WorkspaceTransitionStop {
30643090
require.Equal(t, sidebarAppID, build.AITaskSidebarAppID.UUID.String())
30653091
}

‎coderd/wsbuilder/wsbuilder.go‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"context"
77
"database/sql"
88
"encoding/json"
9+
"errors"
910
"fmt"
1011
"net/http"
1112
"time"
@@ -488,6 +489,21 @@ func (b *Builder) buildTx(authFunc func(action policy.Action, object rbac.Object
488489
return BuildError{code, "insert workspace build", err}
489490
}
490491

492+
// If this is a task workspace, link it to the latest workspace build.
493+
if task, err := store.GetTaskByWorkspaceID(b.ctx, b.workspace.ID); err == nil {
494+
_, err = store.UpsertTaskWorkspaceApp(b.ctx, database.UpsertTaskWorkspaceAppParams{
495+
TaskID: task.ID,
496+
WorkspaceBuildNumber: buildNum,
497+
WorkspaceAgentID: uuid.NullUUID{}, // Updated by the provisioner upon job completion.
498+
WorkspaceAppID: uuid.NullUUID{}, // Updated by the provisioner upon job completion.
499+
})
500+
if err != nil {
501+
return BuildError{http.StatusInternalServerError, "upsert task workspace app", err}
502+
}
503+
} else if !errors.Is(err, sql.ErrNoRows) {
504+
return BuildError{http.StatusInternalServerError, "get task by workspace id", err}
505+
}
506+
491507
err = store.InsertWorkspaceBuildParameters(b.ctx, database.InsertWorkspaceBuildParametersParams{
492508
WorkspaceBuildID: workspaceBuildID,
493509
Name: names,

0 commit comments

Comments
 (0)