Skip to content

Commit c741a08

Browse files
dannykoppingdannykopping
committed
WIP CODER_AGENT_TOKEN reuse
Signed-off-by: Danny Kopping <danny@coder.com>
1 parent 7e1abd1 commit c741a08

File tree

7 files changed

+62
-7
lines changed

7 files changed

+62
-7
lines changed

‎coderd/provisionerdserver/provisionerdserver.go

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -500,6 +500,17 @@ func (s *server) acquireProtoJob(ctx context.Context, job database.ProvisionerJo
500500
for _, group := range ownerGroups {
501501
ownerGroupNames = append(ownerGroupNames, group.Group.Name)
502502
}
503+
var runningWorkspaceAgentToken string
504+
if input.RunningWorkspaceAgentID != uuid.Nil {
505+
agent, err := s.Database.GetWorkspaceAgentByID(ctx, input.RunningWorkspaceAgentID)
506+
if err != nil {
507+
s.Logger.Warn(ctx, "failed to retrieve running workspace agent by ID; this may affect prebuilds",
508+
slog.F("workspace_agent_id", input.RunningWorkspaceAgentID),
509+
slog.F("job_id", job.ID))
510+
} else {
511+
runningWorkspaceAgentToken = agent.AuthToken.String()
512+
}
513+
}
503514

504515
msg, err := json.Marshal(wspubsub.WorkspaceEvent{
505516
Kind: wspubsub.WorkspaceEventKindStateChange,
@@ -622,6 +633,7 @@ func (s *server) acquireProtoJob(ctx context.Context, job database.ProvisionerJo
622633
WorkspaceBuildId: workspaceBuild.ID.String(),
623634
WorkspaceOwnerLoginType: string(owner.LoginType),
624635
IsPrebuild: input.IsPrebuild,
636+
RunningWorkspaceAgentToken: runningWorkspaceAgentToken,
625637
},
626638
LogLevel: input.LogLevel,
627639
},
@@ -2347,7 +2359,14 @@ type WorkspaceProvisionJob struct {
23472359
WorkspaceBuildID uuid.UUID `json:"workspace_build_id"`
23482360
DryRun bool `json:"dry_run"`
23492361
IsPrebuild bool `json:"is_prebuild,omitempty"`
2350-
LogLevel string `json:"log_level,omitempty"`
2362+
// RunningWorkspaceAgentID is *only* used for prebuilds. We pass it down when we want to rebuild a prebuilt workspace
2363+
// but not generate a new agent token. The provisionerdserver will retrieve this token and push it down to
2364+
// the provisioner (and ultimately to the `coder_agent` resource in the Terraform provider) where it will be
2365+
// reused. Context: the agent token is often used in immutable attributes of workspace resource (e.g. VM/container)
2366+
// to initialize the agent, so if that value changes it will necessitate a replacement of that resource, thus
2367+
// obviating the whole point of the prebuild.
2368+
RunningWorkspaceAgentID uuid.UUID `json:"running_workspace_agent_id"`
2369+
LogLevel string `json:"log_level,omitempty"`
23512370
}
23522371

23532372
// TemplateVersionDryRunJob is the payload for the "template_version_dry_run" job type.

‎coderd/workspaces.go

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -627,6 +627,8 @@ func createWorkspace(
627627
provisionerJob *database.ProvisionerJob
628628
workspaceBuild *database.WorkspaceBuild
629629
provisionerDaemons []database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow
630+
631+
runningWorkspaceAgentID uuid.UUID
630632
)
631633
err = api.Database.InTx(func(db database.Store) error {
632634
var claimedWorkspace *database.Workspace
@@ -663,8 +665,17 @@ func createWorkspace(
663665
api.Logger.Warn(ctx, "unable to find claimed workspace by ID", slog.Error(err), slog.F("claimed_prebuild_id", (*claimedID).String()))
664666
goto regularPath
665667
}
666-
667668
claimedWorkspace = &lookup
669+
670+
agents, err := api.Database.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ownerCtx, claimedWorkspace.ID)
671+
if err != nil {
672+
api.Logger.Error(ctx, "failed to retrieve running agents of claimed prebuilt workspace",
673+
slog.F("workspace_id", claimedWorkspace.ID), slog.Error(err))
674+
}
675+
if len(agents) >= 1 {
676+
// TODO: handle multiple agents
677+
runningWorkspaceAgentID = agents[0].ID
678+
}
668679
}
669680

670681
regularPath:
@@ -710,7 +721,8 @@ func createWorkspace(
710721
Reason(database.BuildReasonInitiator).
711722
Initiator(initiatorID).
712723
ActiveVersion().
713-
RichParameterValues(req.RichParameterValues)
724+
RichParameterValues(req.RichParameterValues).
725+
RunningWorkspaceAgentID(runningWorkspaceAgentID)
714726
if req.TemplateVersionID != uuid.Nil {
715727
builder = builder.VersionID(req.TemplateVersionID)
716728
}

‎coderd/wsbuilder/wsbuilder.go

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ type Builder struct {
7373
parameterNames *[]string
7474
parameterValues *[]string
7575
prebuild bool
76+
runningWorkspaceAgentID uuid.UUID
7677

7778
verifyNoLegacyParametersOnce bool
7879
}
@@ -175,6 +176,13 @@ func (b Builder) MarkPrebuild() Builder {
175176
return b
176177
}
177178

179+
// RunningWorkspaceAgentID is only used for prebuilds; see the associated field in `provisionerdserver.WorkspaceProvisionJob`.
180+
func (b Builder) RunningWorkspaceAgentID(id uuid.UUID) Builder {
181+
// nolint: revive
182+
b.runningWorkspaceAgentID = id
183+
return b
184+
}
185+
178186
// SetLastWorkspaceBuildInTx prepopulates the Builder's cache with the last workspace build. This allows us
179187
// to avoid a repeated database query when the Builder's caller also needs the workspace build, e.g. auto-start &
180188
// auto-stop.
@@ -300,9 +308,10 @@ func (b *Builder) buildTx(authFunc func(action policy.Action, object rbac.Object
300308

301309
workspaceBuildID := uuid.New()
302310
input, err := json.Marshal(provisionerdserver.WorkspaceProvisionJob{
303-
WorkspaceBuildID: workspaceBuildID,
304-
LogLevel: b.logLevel,
305-
IsPrebuild: b.prebuild,
311+
WorkspaceBuildID: workspaceBuildID,
312+
LogLevel: b.logLevel,
313+
IsPrebuild: b.prebuild,
314+
RunningWorkspaceAgentID: b.runningWorkspaceAgentID,
306315
})
307316
if err != nil {
308317
return nil, nil, nil, BuildError{

‎provisioner/terraform/provision.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -263,7 +263,9 @@ func provisionEnv(
263263
"CODER_WORKSPACE_BUILD_ID="+metadata.GetWorkspaceBuildId(),
264264
)
265265
if metadata.GetIsPrebuild() {
266-
env = append(env, "CODER_WORKSPACE_IS_PREBUILD=true")
266+
env = append(env, provider.IsPrebuildEnvironmentVariable()+"=true")
267+
} else {
268+
env = append(env, provider.RunningAgentTokenEnvironmentVariable()+"="+metadata.GetRunningWorkspaceAgentToken())
267269
}
268270
for key, value := range provisionersdk.AgentScriptEnv() {
269271
env = append(env, key+"="+value)

‎provisioner/terraform/safeenv.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,14 @@ func envName(env string) string {
3030
return ""
3131
}
3232

33+
func envVar(env string) string {
34+
parts := strings.SplitN(env, "=", 1)
35+
if len(parts) > 0 {
36+
return parts[1]
37+
}
38+
return ""
39+
}
40+
3341
func isCanarySet(env []string) bool {
3442
for _, e := range env {
3543
if envName(e) == unsafeEnvCanary {

‎provisionersdk/proto/provisioner.proto

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -276,6 +276,7 @@ message Metadata {
276276
string workspace_build_id = 17;
277277
string workspace_owner_login_type = 18;
278278
bool is_prebuild = 19;
279+
string running_workspace_agent_token = 20;
279280
}
280281

281282
// Config represents execution configuration shared by all subsequent requests in the Session

‎site/e2e/provisionerGenerated.ts

Lines changed: 4 additions & 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)