mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
fix: make prebuild claiming durable and idempotent (#23108)
## Problem When a prebuilt workspace is claimed, the agent reinitializes via a single fire-and-forget pubsub event over SSE. If the agent's SSE connection is interrupted at claim time, the event is permanently lost — the workspace is stuck with no self-healing path. Additionally, regular (non-prebuild) workspaces had no way to opt out of the `/reinit` polling loop — agents would reconnect indefinitely to an endpoint that would never send them anything useful. ## Root Cause `workspaceAgentReinit` fetches the workspace (with its current `owner_id`) via `GetWorkspaceByAgentID`, but never checked whether a claim already happened. It only subscribed to pubsub for future events. The database already has durable claim state (`owner_id` changes from `PrebuildsSystemUserID` to the real user), but no layer ever consulted it on reconnection. ## Solution ### Server-side durable check with first-build-initiator gating **TOCTOU-safe ordering**: Subscribe to pubsub claim events *before* any durable checks, so a claim that fires during the check is buffered in the channel rather than lost. **First-build-initiator gating**: When `!workspace.IsPrebuild()` (owner is no longer the system user), look up the first build's `InitiatorID`. The prebuild reconciler always uses `PrebuildsSystemUserID` as the initiator. This distinguishes claimed prebuilds from regular workspaces without any SQL schema changes. - **Regular workspace** (first build initiator ≠ system user) → **409 Conflict**, agent stops reconnecting - **Claimed prebuild, build completed** → pre-seed channel with reinit event and close it, transmitter delivers one-shot then exits - **Claimed prebuild, build in-progress** → fall through to pubsub subscription, agent waits for completion event - **Unclaimed prebuild** → pubsub subscription (existing happy path) ### Declarative reinit events (defense-in-depth) - Added `UserID` field to `ReinitializationEvent` with JSON tags - Switched pubsub serialization from raw string to JSON (with backward-compat fallback for rolling upgrades) - Populated `UserID` at both the publish site and the durable check ### Agent SDK: 409 handling `WaitForReinitLoop` detects 409 Conflict from the server and closes the `reinitEvents` channel, cleanly exiting the retry goroutine. ### Agent CLI: fixed two bugs + added reinitCtx - **Closed channel (`!ok`)**: now blocks on `<-ctx.Done()` instead of `continue`, keeping the current agent running. Previously this would leak agents by skipping `agnt.Close()` and re-entering the loop. - **Duplicate owner reinit**: cancels `reinitCtx` (stops the reinit goroutine), then blocks on `<-ctx.Done()`. Previously `continue` would skip cleanup and create a new agent on the next loop iteration. - **`reinitCtx`**: a cancellable child of `ctx` passed to `WaitForReinitLoop`, allowing the agent to stop the reinit HTTP polling after reinit completes. ### Agent-side idempotency Tracks `lastOwnerID` in the agent reinit loop — duplicate events for the same owner are skipped. ## Testing - **"unclaimed prebuild receives reinit via pubsub"**: prebuild owned by system user, pubsub event triggers reinit - **"claimed prebuild receives one-shot reinit on reconnect"**: first build by system user, owner changed, build completed → immediate reinit (no pubsub needed) - **"claimed prebuild waits during in-progress claim build"**: claimed but build still running → no reinit until build completes - **"regular workspace gets 409"**: first build by real user → 409 Conflict, agent stops polling - Updated claim publisher/listener tests: verify `UserID` survives JSON round-trip + backward compat with raw string payloads - Updated SSE round-trip test: verify `UserID` survives transmit → receive cycle Fixes #22359 ## Rolling upgrade note During a rolling deploy where old coderd instances coexist with new ones, the pubsub `ReinitializationEvent` has a new `workspace_id` field (JSON key `workspace_id`). Old publishers send a raw reason string instead of JSON; the new listener gracefully falls back by treating the entire payload as the reason and filling in `WorkspaceID` from context. The only visible effect during the upgrade window is that `WorkspaceID` may be the zero UUID in agent-side logs — this is cosmetic and resolves once all instances are updated.
This commit is contained in:
+33
-6
@@ -17,6 +17,7 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
"golang.org/x/xerrors"
|
||||
"gopkg.in/natefinch/lumberjack.v2"
|
||||
@@ -272,11 +273,14 @@ func workspaceAgent() *serpent.Command {
|
||||
logger.Info(ctx, "agent devcontainer detection not enabled")
|
||||
}
|
||||
|
||||
reinitEvents := agentsdk.WaitForReinitLoop(ctx, logger, client)
|
||||
reinitCtx, reinitCancel := context.WithCancel(ctx)
|
||||
defer reinitCancel()
|
||||
reinitEvents := agentsdk.WaitForReinitLoop(reinitCtx, logger, client)
|
||||
|
||||
var (
|
||||
lastErr error
|
||||
mustExit bool
|
||||
lastOwnerID uuid.UUID
|
||||
lastErr error
|
||||
mustExit bool
|
||||
)
|
||||
for {
|
||||
prometheusRegistry := prometheus.NewRegistry()
|
||||
@@ -343,9 +347,32 @@ func workspaceAgent() *serpent.Command {
|
||||
case <-ctx.Done():
|
||||
logger.Info(ctx, "agent shutting down", slog.Error(context.Cause(ctx)))
|
||||
mustExit = true
|
||||
case event := <-reinitEvents:
|
||||
logger.Info(ctx, "agent received instruction to reinitialize",
|
||||
slog.F("workspace_id", event.WorkspaceID), slog.F("reason", event.Reason))
|
||||
case event, ok := <-reinitEvents:
|
||||
switch {
|
||||
case !ok:
|
||||
// Channel closed — the reinit loop exited
|
||||
// (terminal 409 or context expired). Keep
|
||||
// running the current agent until the parent
|
||||
// context is canceled.
|
||||
logger.Info(ctx, "reinit channel closed, running without reinit capability")
|
||||
reinitEvents = nil
|
||||
<-ctx.Done()
|
||||
mustExit = true
|
||||
case event.OwnerID != uuid.Nil && event.OwnerID == lastOwnerID:
|
||||
// Duplicate reinit for same owner — already
|
||||
// reinitialized. Cancel the reinit loop
|
||||
// goroutine and keep the current agent.
|
||||
logger.Info(ctx, "skipping redundant reinit, owner unchanged",
|
||||
slog.F("owner_id", event.OwnerID))
|
||||
reinitCancel()
|
||||
reinitEvents = nil
|
||||
<-ctx.Done()
|
||||
mustExit = true
|
||||
default:
|
||||
lastOwnerID = event.OwnerID
|
||||
logger.Info(ctx, "agent received instruction to reinitialize",
|
||||
slog.F("workspace_id", event.WorkspaceID), slog.F("reason", event.Reason))
|
||||
}
|
||||
}
|
||||
|
||||
lastErr = agnt.Close()
|
||||
|
||||
Reference in New Issue
Block a user