diff --git a/coderd/database/dbfake/dbfake.go b/coderd/database/dbfake/dbfake.go index 0b859a4fb1..82b66f504a 100644 --- a/coderd/database/dbfake/dbfake.go +++ b/coderd/database/dbfake/dbfake.go @@ -69,6 +69,8 @@ type WorkspaceBuildBuilder struct { jobErrorCode string // Error code for failed jobs provisionerState []byte + + prebuiltWorkspaceBuildStage sdkproto.PrebuiltWorkspaceBuildStage } // BuilderOption is a functional option for customizing job timestamps @@ -149,6 +151,14 @@ func (b WorkspaceBuildBuilder) ProvisionerState(state []byte) WorkspaceBuildBuil return b } +// MarkPrebuiltWorkspaceClaim marks the build's provisioner job as the claim +// of a prebuilt workspace, mirroring wsbuilder.MarkPrebuiltWorkspaceClaim. +func (b WorkspaceBuildBuilder) MarkPrebuiltWorkspaceClaim() WorkspaceBuildBuilder { + //nolint: revive // returns modified struct + b.prebuiltWorkspaceBuildStage = sdkproto.PrebuiltWorkspaceBuildStage_CLAIM + return b +} + func (b WorkspaceBuildBuilder) Resource(resource ...*sdkproto.Resource) WorkspaceBuildBuilder { //nolint: revive // returns modified struct b.resources = append(b.resources, resource...) @@ -368,7 +378,8 @@ func (b WorkspaceBuildBuilder) doInTX() WorkspaceResponse { // Create a provisioner job for the build! payload, err := json.Marshal(provisionerdserver.WorkspaceProvisionJob{ - WorkspaceBuildID: b.seed.ID, + WorkspaceBuildID: b.seed.ID, + PrebuiltWorkspaceBuildStage: b.prebuiltWorkspaceBuildStage, }) require.NoError(b.t, err) diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go index e424f9b536..915cd2ac90 100644 --- a/coderd/workspaceagents.go +++ b/coderd/workspaceagents.go @@ -37,6 +37,7 @@ import ( "github.com/coder/coder/v2/coderd/httpmw/loggermw" "github.com/coder/coder/v2/coderd/jwtutils" "github.com/coder/coder/v2/coderd/prebuilds" + "github.com/coder/coder/v2/coderd/provisionerdserver" "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/coderd/rbac/policy" "github.com/coder/coder/v2/coderd/telemetry" @@ -1537,14 +1538,13 @@ func (api *API) workspaceAgentReinit(rw http.ResponseWriter, r *http.Request) { return } - // This workspace was a prebuild that got claimed. Check if - // the claim build completed successfully before sending - // reinit. We assume the latest build is the claim build - // (build 2). If a third build (e.g. a restart) starts - // between the claim and the agent's reconnection, this - // would check that build instead. The window is extremely - // small in practice, and a restart would trigger its own - // reinit path. + // This workspace was a prebuild that got claimed. The seeded + // reinit below recovers a claim event that was missed while + // the agent's /reinit connection was down. It only applies + // while the latest build is the claim build itself, which the + // build's provisioner job input records, mirroring the check + // the provisioner server uses when publishing the claim + // event. latestBuild, err := api.Database.GetLatestWorkspaceBuildByWorkspaceID(ctx, workspace.ID) if err != nil { log.Error(ctx, "failed to get latest workspace build", slog.Error(err)) @@ -1557,43 +1557,73 @@ func (api *API) workspaceAgentReinit(rw http.ResponseWriter, r *http.Request) { httpapi.InternalServerError(rw, xerrors.New("failed to get provisioner job")) return } - - if job.CompletedAt.Valid && !job.Error.Valid { - // Claim build succeeded — cancel the pubsub - // subscription (no longer needed) and swap in a - // pre-seeded channel so the transmitter delivers - // exactly one reinit event. - cancelSub() - seeded := make(chan agentsdk.ReinitializationEvent, 1) - seeded <- agentsdk.ReinitializationEvent{ - WorkspaceID: workspace.ID, - Reason: agentsdk.ReinitializeReasonPrebuildClaimed, - OwnerID: workspace.OwnerID, - } - reinitEvents = seeded - } else if job.CompletedAt.Valid && job.Error.Valid { - // Claim build failed permanently. Return 409 so the - // agent treats this as terminal and stops retrying - // (WaitForReinitLoop exits on any 409). - cancelSub() - log.Warn(ctx, "claim build failed", - slog.F("job_id", job.ID), - slog.F("error", job.Error.String)) - httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ - Message: "Claim build failed permanently.", - Detail: job.Error.String, - }) + var jobInput provisionerdserver.WorkspaceProvisionJob + if err := json.Unmarshal(job.Input, &jobInput); err != nil { + log.Error(ctx, "failed to unmarshal provisioner job input", slog.Error(err)) + httpapi.InternalServerError(rw, xerrors.New("failed to unmarshal provisioner job input")) return } - // Claim build still in progress — fall through to the - // transmitter. The pubsub subscription (set up above) - // will deliver the event when the build completes - // successfully. Note: FailJob does not publish a claim - // event, so a failed in-progress build will leave the - // agent blocking here until it disconnects and - // reconnects (at which point the durable check above - // handles it). + switch { + case jobInput.PrebuiltWorkspaceBuildStage.IsPrebuiltWorkspaceClaim(): + if job.CompletedAt.Valid && !job.Error.Valid { + // Claim build succeeded: cancel the pubsub + // subscription (no longer needed) and swap in a + // pre-seeded channel so the transmitter delivers + // exactly one reinit event. + cancelSub() + seeded := make(chan agentsdk.ReinitializationEvent, 1) + seeded <- agentsdk.ReinitializationEvent{ + WorkspaceID: workspace.ID, + Reason: agentsdk.ReinitializeReasonPrebuildClaimed, + OwnerID: workspace.OwnerID, + } + reinitEvents = seeded + } else if job.CompletedAt.Valid && job.Error.Valid { + // Claim build failed permanently. Return 409 so the + // agent treats this as terminal and stops retrying + // (WaitForReinitLoop exits on any 409). + cancelSub() + log.Warn(ctx, "claim build failed", + slog.F("job_id", job.ID), + slog.F("error", job.Error.String)) + httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ + Message: "Claim build failed permanently.", + Detail: job.Error.String, + }) + return + } + // Claim build still in progress: proceed to the + // transmitter below. The pubsub subscription (set up + // above) will deliver the event when the build completes + // successfully. Note: FailJob does not publish a claim + // event, so a failed in-progress build will leave the + // agent blocking here until it disconnects and + // reconnects (at which point the durable check above + // handles it). + case latestBuild.InitiatorID == database.PrebuildsSystemUserID: + // The workspace owner has changed but the claim build has + // not been created yet. Proceed to the transmitter below; + // the pubsub subscription set up above delivers the claim + // event once the claim build completes. + default: + // The latest build is a user-initiated build other than + // the claim build, so the claim has already been handled. + // Re-sending the reinit event here would needlessly + // restart the agent of a long-claimed workspace on every + // reconnection. Return 409 so the agent stops polling, + // the same as a regular workspace. + log.Debug(ctx, "prebuild claim already handled, stopping reinit polling", + slog.F("job_id", job.ID), + slog.F("latest_build_id", latestBuild.ID), + slog.F("latest_build_number", latestBuild.BuildNumber)) + cancelSub() + httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ + Message: "Workspace is not a prebuilt workspace waiting to be claimed.", + Detail: "The prebuild claim for this workspace has already been handled by an earlier build.", + }) + return + } } transmitter := agentsdk.NewSSEAgentReinitTransmitter(log, rw, r) diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go index 583332ebba..1e52d1e35c 100644 --- a/coderd/workspaceagents_test.go +++ b/coderd/workspaceagents_test.go @@ -3464,6 +3464,7 @@ func TestReinit(t *testing.T) { InitiatorID: claimerID, Transition: database.WorkspaceTransitionStart, }). + MarkPrebuiltWorkspaceClaim(). WithAgent() if !complete { builder = builder.Starting() @@ -3562,6 +3563,52 @@ func TestReinit(t *testing.T) { require.Equal(t, user.UserID, reinitEvent.OwnerID) }) + // Verifies that the durable claim check only applies while the + // latest build is the claim build. A workspace that was claimed + // in the past and has since had user-initiated builds must get a + // 409 instead of another reinit, otherwise its agent would be + // restarted on every /reinit reconnection for the rest of the + // workspace's life. + t.Run("workspace claimed in the past gets 409", func(t *testing.T) { + t.Parallel() + + db, ps, sqlDB := dbtestutil.NewDBWithSQLDB(t) + client := coderdtest.New(t, &coderdtest.Options{ + Database: db, + Pubsub: ps, + }) + user := coderdtest.CreateFirstUser(t, client) + + // Create an unclaimed prebuild (build 1, completed) and claim + // it (build 2, completed). + r := setupPrebuildWorkspace(t, db, user.OrganizationID) + claimPrebuild(t, db, sqlDB, r.Workspace, user.UserID, r.TemplateVersion.ID, true) + + // A later build initiated by the owner (e.g. a restart) means + // the claim has already been handled. + ws := r.Workspace + ws.OwnerID = user.UserID + laterR := dbfake.WorkspaceBuild(t, db, ws). + Seed(database.WorkspaceBuild{ + TemplateVersionID: r.TemplateVersion.ID, + BuildNumber: 3, + InitiatorID: user.UserID, + Transition: database.WorkspaceTransitionStart, + }). + WithAgent(). + Do() + + agentCtx := testutil.Context(t, testutil.WaitShort) + agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(laterR.AgentToken)) + + // WaitForReinit should return an error wrapping a 409. + _, err := agentClient.WaitForReinit(agentCtx) + require.Error(t, err) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusConflict, sdkErr.StatusCode()) + }) + // Verifies that when the claim build completed with an error, // the handler returns 409 so the agent treats it as terminal // and stops retrying (WaitForReinitLoop exits on any 409).