diff --git a/apps/docs/content/docs/daemon-runtimes.mdx b/apps/docs/content/docs/daemon-runtimes.mdx index 63ea6b60e31..e3fa2de7ce2 100644 --- a/apps/docs/content/docs/daemon-runtimes.mdx +++ b/apps/docs/content/docs/daemon-runtimes.mdx @@ -31,7 +31,7 @@ Common commands: |---|---| | `multica daemon start` | Start (background by default; add `--foreground` to run in the foreground) | | `multica daemon stop` | Stop | -| `multica daemon restart` | Restart | +| `multica daemon restart` | Restart immediately (add `--drain` to let active tasks finish first) | | `multica daemon status` | Show status | | `multica daemon logs` | Show logs (add `-f` to follow) | diff --git a/docs/plans/2026-07-19-001-daemon-drain-restart-plan.md b/docs/plans/2026-07-19-001-daemon-drain-restart-plan.md new file mode 100644 index 00000000000..03ef670e5a0 --- /dev/null +++ b/docs/plans/2026-07-19-001-daemon-drain-restart-plan.md @@ -0,0 +1,65 @@ +--- +title: Drain-Aware Daemon Restart - Plan +type: feat +date: 2026-07-19 +topic: daemon-drain-restart +artifact_contract: ce-unified-plan/v1 +artifact_readiness: implementation-ready +execution: code +--- + +# Drain-Aware Daemon Restart - Plan + +## Goal + +Add `multica daemon restart --drain`: stop claiming, let accepted work finish, then restart through the invoking CLI. Plain restart remains immediate for compatibility. + +## Problem + +`daemon restart` currently calls the local `/shutdown` endpoint. That cancels the daemon root context, which is also passed to active agent tasks. The subsequent 30-second wait only waits for cancelled task goroutines to clean up; it does not let agent work finish. + +`runtime_recovery` can retry the interrupted row, but still records a failure, redispatches work, and can spend more model tokens. It is crash recovery, not safe restart. + +## Design + +1. Add an opt-in `--drain` flag to `daemon restart`. +2. Use a dedicated local endpoint so an older daemon returns 404 rather than ignoring a query parameter and stopping immediately. +3. Under `claimMu`, represent the shared claim barrier with one explicit owner: `none`, `drain`, or `update`. Any non-`none` owner pauses new claims. This single state replaces independent `pauseClaims`, `draining`, and `updating` flags so acquisition and release cannot disagree about which lifecycle operation is active. +4. Manual drain acquires the `drain` owner even while claims or tasks are active, then waits for `claimsInFlight == 0 && activeTasks == 0` without cancelling active contexts. Preserve the existing handoff invariant: `activeTasks` increments before `claimsInFlight` decrements, so the drain cannot observe false idle. +5. Periodic auto-update acquires the `update` owner only while fully idle. Heartbeat-triggered update acquires the same `update` owner without requiring idle, preserving its existing immediate-update behavior while preventing it from bypassing an active drain. +6. A heartbeat update that cannot acquire ownership remains pending and is retried by a later heartbeat; claim-barrier contention is not reported as an update failure. A drain that finds `update` ownership already held returns `409 Conflict`. +7. Release is owner-specific. A cancelled drain releases only `drain`; a failed update releases only `update`. A successful drain or update keeps ownership through root-context cancellation so no new claim can enter during shutdown. +8. After drain reaches idle, cancel the daemon and use the existing CLI start path, preserving binary, profile, foreground mode, and overrides. + +## Ownership Transitions + +| Current owner | Request | Result | +| --- | --- | --- | +| `none` | new task claim | Claim enters and increments `claimsInFlight` | +| `none` | manual drain | Acquire `drain`, pause new claims, wait for accepted work | +| `none`, idle | periodic auto-update | Acquire `update` and run the upgrade | +| `none` | heartbeat update | Acquire `update` and run the existing immediate upgrade path | +| `drain` | task claim | Reject the claim attempt | +| `drain` | periodic or heartbeat update | Defer without running or reporting failure | +| `update` | manual drain | Return `409 Conflict` | +| `update` | another update | Defer without starting a second upgrade | +| `drain` | requester cancellation | Release `drain` and resume claims | +| `update` | upgrade failure | Release `update` and resume claims | +| `drain` or `update` | successful shutdown/restart | Retain ownership until process exit | + +Plain `multica daemon restart` remains an intentionally immediate shutdown and does not participate in this opt-in drain protocol. + +## Concurrency Verification + +Add deterministic tests for both reviewer-reported orderings using channel-gated operations rather than timing-only assertions: + +1. **Drain then heartbeat update:** keep one task active, start `/shutdown/drain`, wait until `drain` owns the barrier, then invoke `handleUpdate`. Assert that the update function and restart are not called, the drain retains ownership, and the update is not reported failed. +2. **Heartbeat update then drain:** start `handleUpdate` with an update function blocked after `update` ownership is acquired, then call `/shutdown/drain`. Assert `409 Conflict`; release the update and verify its existing completion/restart path. + +Keep the existing auto-update, claim-handoff, drain cancellation, concurrent-drain, endpoint compatibility, and immediate-shutdown regression tests. Run the focused package tests under the race detector because the contract is specifically about cross-goroutine ownership and ordering. + +## Scope and Verification + +Change only daemon lifecycle/health code, focused tests, and the daemon command doc. Do not change default restart, server recovery, executable trust, or PR #5494 ownership semantics. + +Cover active work, claim handoff, client cancellation, all owner transition conflicts, dedicated endpoint selection, immediate-shutdown regression, focused package tests, race tests, vet, and formatting. diff --git a/server/cmd/multica/cmd_daemon.go b/server/cmd/multica/cmd_daemon.go index 724b033f4ae..5a24b93c8d7 100644 --- a/server/cmd/multica/cmd_daemon.go +++ b/server/cmd/multica/cmd_daemon.go @@ -119,6 +119,7 @@ func init() { rf.Int("max-concurrent-tasks", 0, "Max tasks running in parallel (env: MULTICA_DAEMON_MAX_CONCURRENT_TASKS)") rf.Bool("no-auto-update", false, "Disable periodic CLI self-update (env: MULTICA_DAEMON_AUTO_UPDATE=false)") rf.Duration("auto-update-interval", 0, "How often to poll GitHub for a newer release (env: MULTICA_DAEMON_AUTO_UPDATE_INTERVAL)") + rf.Bool("drain", false, "Stop claiming new tasks and wait for active tasks to finish before restarting") df := daemonDiskUsageCmd.Flags() df.Bool("by-workspace", false, "Aggregate output by workspace instead of by task") @@ -932,6 +933,7 @@ func requireDaemonRestartPreflight(cmd *cobra.Command, profile string) error { func runDaemonRestart(cmd *cobra.Command, args []string) error { profile := resolveProfile(cmd) healthPort := healthPortForProfile(profile) + drain, _ := cmd.Flags().GetBool("drain") ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() @@ -949,10 +951,18 @@ func runDaemonRestart(cmd *cobra.Command, args []string) error { } pid, _ := health["pid"].(float64) if pid > 0 { - fmt.Fprintf(os.Stderr, "Stopping daemon (pid %d)...\n", int(pid)) - if err := requestDaemonShutdown(healthPort); err != nil { - if p, perr := os.FindProcess(int(pid)); perr == nil { - _ = p.Kill() + if drain { + active, _ := health["active_task_count"].(float64) + fmt.Fprintf(os.Stderr, "Draining daemon (pid %d); waiting for %d active task(s)...\n", int(pid), int(active)) + if err := requestDaemonDrainShutdown(cmd.Context(), healthPort); err != nil { + return fmt.Errorf("drain daemon before restart: %w", err) + } + } else { + fmt.Fprintf(os.Stderr, "Stopping daemon (pid %d)...\n", int(pid)) + if err := requestDaemonShutdown(healthPort); err != nil { + if p, perr := os.FindProcess(int(pid)); perr == nil { + _ = p.Kill() + } } } // Wait until the port is fully released (not merely past "running"), @@ -1055,6 +1065,29 @@ func requestDaemonShutdown(healthPort int) error { return nil } +// requestDaemonDrainShutdown asks the daemon to stop claiming new tasks, wait +// for every in-flight claim and active task to finish naturally, and then shut +// down. It deliberately has no fixed client timeout: a legitimate agent task +// can run for hours. Cancelling ctx aborts the request and makes the daemon +// release its claim barrier instead of leaving the runtime paused. +func requestDaemonDrainShutdown(ctx context.Context, healthPort int) error { + url := fmt.Sprintf("http://127.0.0.1:%d/shutdown/drain", healthPort) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, nil) + if err != nil { + return err + } + resp, err := (&http.Client{}).Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + body, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<10)) + return fmt.Errorf("unexpected status %d: %s", resp.StatusCode, strings.TrimSpace(string(body))) + } + return nil +} + // --- daemon status --- func runDaemonStatus(cmd *cobra.Command, _ []string) error { diff --git a/server/cmd/multica/cmd_daemon_test.go b/server/cmd/multica/cmd_daemon_test.go index 136cfc44834..828e42ac333 100644 --- a/server/cmd/multica/cmd_daemon_test.go +++ b/server/cmd/multica/cmd_daemon_test.go @@ -2,6 +2,7 @@ package main import ( "bytes" + "context" "fmt" "net" "net/http" @@ -384,6 +385,57 @@ func TestDaemonRestartUnauthenticatedFailsBeforeStopping(t *testing.T) { } } +func TestDaemonRestartCommandExposesDrainFlag(t *testing.T) { + t.Parallel() + + flag := daemonRestartCmd.Flags().Lookup("drain") + if flag == nil { + t.Fatal("daemon restart is missing the --drain flag") + } + if flag.DefValue != "false" { + t.Fatalf("--drain default = %q, want false for backwards compatibility", flag.DefValue) + } +} + +func TestRequestDaemonDrainShutdownUsesDedicatedEndpointAndCallerContext(t *testing.T) { + t.Parallel() + + requestSeen := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + t.Errorf("method = %s, want POST", r.Method) + } + if got := r.URL.Path; got != "/shutdown/drain" { + t.Errorf("path = %q, want /shutdown/drain", got) + } + close(requestSeen) + <-r.Context().Done() + })) + defer srv.Close() + + port := srv.Listener.Addr().(*net.TCPAddr).Port + ctx, cancel := context.WithCancel(context.Background()) + errCh := make(chan error, 1) + go func() { + errCh <- requestDaemonDrainShutdown(ctx, port) + }() + + select { + case <-requestSeen: + case <-time.After(time.Second): + t.Fatal("drain shutdown request was not delivered") + } + cancel() + select { + case err := <-errCh: + if err == nil || !strings.Contains(err.Error(), "context canceled") { + t.Fatalf("request error = %v, want caller context cancellation", err) + } + case <-time.After(time.Second): + t.Fatal("drain shutdown request ignored caller cancellation") + } +} + // fakeRunningDaemon serves a fake healthy daemon on the given profile's health // port and reports any /shutdown request on the returned channel. The PID in // /health is our own so a kill-fallback would be visible as a test crash too. diff --git a/server/internal/daemon/auto_update.go b/server/internal/daemon/auto_update.go index 2d1673524cc..d4b7f4a35f8 100644 --- a/server/internal/daemon/auto_update.go +++ b/server/internal/daemon/auto_update.go @@ -83,7 +83,7 @@ func (d *Daemon) autoUpdateLoop(ctx context.Context) { } // tryAutoUpdate runs one check-and-maybe-upgrade cycle. Bails early on any of: -// already updating (server-triggered upgrade in flight), active tasks (defer +// update ownership already held, active tasks (defer // to next tick — we never interrupt running agents), version fetch failure, // or no newer release. The function never returns an error: a check that // fails today will be retried at the next tick, and we don't want a transient @@ -96,7 +96,7 @@ func (d *Daemon) tryAutoUpdate(ctx context.Context) { // the Runtimes page is already in flight, let it finish and re-check next // tick (by which time we'll either be on the new binary or it failed and // we can retry). - if d.updating.Load() { + if d.isUpdating() { d.logger.Debug("auto-update: skip — update already in progress") return } @@ -104,7 +104,7 @@ func (d *Daemon) tryAutoUpdate(ctx context.Context) { // HTTPS call to GitHub, and there is no point paying that cost (or the // rate-limit budget) when we already know we are going to defer. A task // that starts between this load and the barrier check below is caught - // by the strict re-check under claimMu inside trySetClaimBarrier. + // by the strict re-check under claimMu inside tryBeginUpdate. if running := d.activeTasks.Load(); running > 0 { d.logger.Debug("auto-update: skip — tasks running", "active", running) return @@ -122,36 +122,19 @@ func (d *Daemon) tryAutoUpdate(ctx context.Context) { return } - // CAS the updating flag so a concurrent server-triggered handleUpdate - // dropped onto a heartbeat tick can't double-fire. Release on every exit - // path before triggerRestart — once that lands, the daemon ctx is - // cancelled and the flag dies with the process. - if !d.updating.CompareAndSwap(false, true) { - d.logger.Debug("auto-update: skip — update already in progress (raced)") - return - } - released := false - defer func() { - if !released { - d.updating.Store(false) - } - }() - // Strict barrier: between the cheap pre-fetch idle check and now the // release fetch took anywhere from tens of milliseconds (typical) to - // seconds (slow link, GitHub hiccup), plenty of time for a poller to - // claim a fresh task. trySetClaimBarrier checks claimsInFlight + - // activeTasks under claimMu and only flips pauseClaims to true if both - // are zero, so once it returns true we can run the upgrade knowing that - // no in-flight task will be cancelled by triggerRestart. - if !d.trySetClaimBarrier() { - d.logger.Info("auto-update: deferring — task or claim in flight at barrier check") + // seconds (slow link, GitHub hiccup), plenty of time for a task, drain, or + // heartbeat-triggered update to acquire lifecycle ownership. The update + // owner atomically covers all three cases and requires full idle here. + if !d.tryBeginUpdate(true) { + d.logger.Info("auto-update: deferring — lifecycle barrier unavailable or daemon not idle") return } - barrierReleased := false + keepBarrier := false defer func() { - if !barrierReleased { - d.releaseClaimBarrier() + if !keepBarrier { + d.releaseUpdate() } }() @@ -166,11 +149,8 @@ func (d *Daemon) tryAutoUpdate(ctx context.Context) { d.logger.Info("auto-update: upgrade completed, restarting", "target", release.TagName, "output", output) // triggerRestart cancels the root context, which causes Run() to return - // and the parent (cmd_daemon.go) to re-exec the new binary. Leave both - // the updating flag and the claim barrier held — process exit is - // imminent and clearing either would open a window for new claims / a - // second auto-update tick to fire mid-shutdown. - released = true - barrierReleased = true + // and the parent (cmd_daemon.go) to re-exec the new binary. Keep ownership + // only if a restart was actually scheduled; otherwise resume claims. d.triggerRestart() + keepBarrier = d.RestartBinary() != "" } diff --git a/server/internal/daemon/auto_update_test.go b/server/internal/daemon/auto_update_test.go index ad7f54f63af..a7f3defdc68 100644 --- a/server/internal/daemon/auto_update_test.go +++ b/server/internal/daemon/auto_update_test.go @@ -41,7 +41,9 @@ func withStubRelease(t *testing.T, release *cli.GitHubRelease, err error) { func TestTryAutoUpdate_SkipsWhenUpdating(t *testing.T) { d, restartCalls := newAutoUpdateTestDaemon(t, "v0.1.13") - d.updating.Store(true) + if !d.tryBeginUpdate(false) { + t.Fatal("failed to acquire update owner for test setup") + } withStubRelease(t, &cli.GitHubRelease{TagName: "v0.1.14"}, nil) d.tryAutoUpdate(context.Background()) @@ -61,15 +63,15 @@ func TestTryAutoUpdate_SkipsWhenTasksRunning(t *testing.T) { if restartCalls.Load() != 0 { t.Fatalf("triggerRestart fired with active tasks; auto-update must defer") } - if d.updating.Load() { - t.Fatalf("updating flag should not have been claimed while tasks were running") + if d.isUpdating() { + t.Fatalf("update owner should not have been acquired while tasks were running") } } // TestTryAutoUpdate_DefersWhenClaimInFlightAtBarrier covers the race the // review flagged: cheap pre-fetch idle check passes (activeTasks == 0), then // during the release fetch a poller decides to claim and bumps -// claimsInFlight. trySetClaimBarrier must observe that and defer rather than +// claimsInFlight. tryBeginUpdate must observe that and defer rather than // proceed into runUpdate (which would lead to a triggerRestart cancelling // the just-claimed task mid-run). func TestTryAutoUpdate_DefersWhenClaimInFlightAtBarrier(t *testing.T) { @@ -83,16 +85,17 @@ func TestTryAutoUpdate_DefersWhenClaimInFlightAtBarrier(t *testing.T) { if restartCalls.Load() != 0 { t.Fatalf("triggerRestart fired despite a claim being in flight at the barrier") } - if d.updating.Load() { - t.Fatalf("updating flag must be released after a deferred upgrade so the next tick can retry") + if d.isUpdating() { + t.Fatalf("update owner must be released after a deferred upgrade so the next tick can retry") } - if d.pauseClaims { - t.Fatalf("pauseClaims must be cleared after a deferred upgrade") + if !d.tryEnterClaim() { + t.Fatal("claims must remain enabled after a deferred upgrade") } + d.exitClaim() } // TestTryAutoUpdate_HoldsBarrierAcrossRestart asserts the success path leaves -// pauseClaims set: process exit is imminent and clearing the barrier would +// update ownership held: process exit is imminent and clearing the barrier would // open a window for a poller to claim a task that the imminent restart is // about to cancel. func TestTryAutoUpdate_HoldsBarrierAcrossRestart(t *testing.T) { @@ -105,13 +108,13 @@ func TestTryAutoUpdate_HoldsBarrierAcrossRestart(t *testing.T) { if restartCalls.Load() != 1 { t.Fatalf("triggerRestart fired %d times, want 1", restartCalls.Load()) } - if !d.pauseClaims { - t.Fatalf("pauseClaims must remain set across the restart kick; got cleared") + if !d.isUpdating() { + t.Fatalf("update ownership must remain held across the restart kick; got cleared") } } // TestTryAutoUpdate_ReleasesBarrierOnUpgradeFailure asserts the failure path -// clears pauseClaims so the daemon can keep claiming tasks normally and +// clears update ownership so the daemon can keep claiming tasks normally and // retry the upgrade on the next tick. func TestTryAutoUpdate_ReleasesBarrierOnUpgradeFailure(t *testing.T) { d, restartCalls := newAutoUpdateTestDaemon(t, "v0.1.13") @@ -125,13 +128,13 @@ func TestTryAutoUpdate_ReleasesBarrierOnUpgradeFailure(t *testing.T) { if restartCalls.Load() != 0 { t.Fatalf("triggerRestart fired despite upgrade failure") } - if d.pauseClaims { - t.Fatalf("pauseClaims must be cleared after a failed upgrade so pollers resume claiming") + if d.isUpdating() { + t.Fatalf("update ownership must be cleared after a failed upgrade so pollers resume claiming") } } // TestTryEnterClaim_RespectsBarrier asserts the poller-side helper returns -// false while pauseClaims is held and that pairs of enter/exit balance the +// false while a lifecycle owner is held and that pairs of enter/exit balance the // counter so a later barrier set sees idle. func TestTryEnterClaim_RespectsBarrier(t *testing.T) { d := &Daemon{} @@ -144,15 +147,27 @@ func TestTryEnterClaim_RespectsBarrier(t *testing.T) { t.Fatalf("claimsInFlight not balanced: %d", d.claimsInFlight) } - if !d.trySetClaimBarrier() { - t.Fatal("trySetClaimBarrier should succeed when idle") + if !d.tryBeginUpdate(true) { + t.Fatal("tryBeginUpdate should acquire the barrier when idle") } if d.tryEnterClaim() { t.Fatal("tryEnterClaim must refuse while barrier is held") } - d.releaseClaimBarrier() + d.releaseUpdate() + if !d.tryBeginDrain() { + t.Fatal("tryBeginDrain should acquire an idle claim barrier") + } + if d.tryBeginUpdate(true) { + t.Fatal("auto-update must not steal a manual drain barrier") + } + d.releaseUpdate() + if d.tryEnterClaim() { + d.exitClaim() + t.Fatal("a mismatched update release must not clear the drain barrier") + } + d.releaseDrain() if !d.tryEnterClaim() { - t.Fatal("tryEnterClaim should succeed after barrier release") + t.Fatal("tryEnterClaim should succeed after barriers are released") } d.exitClaim() } @@ -197,8 +212,8 @@ func TestTryAutoUpdate_RunsUpgradeAndRestartsOnNewer(t *testing.T) { if restartCalls.Load() != 1 { t.Fatalf("triggerRestart fired %d times, want 1", restartCalls.Load()) } - if !d.updating.Load() { - t.Fatalf("updating flag should remain set across the restart kick; got cleared") + if !d.isUpdating() { + t.Fatalf("update owner should remain held across the restart kick; got cleared") } } @@ -215,8 +230,8 @@ func TestTryAutoUpdate_DoesNotRestartOnUpgradeFailure(t *testing.T) { if restartCalls.Load() != 0 { t.Fatalf("triggerRestart fired despite upgrade failure") } - if d.updating.Load() { - t.Fatalf("updating flag must be released after a failed upgrade so the next tick can retry") + if d.isUpdating() { + t.Fatalf("update owner must be released after a failed upgrade so the next tick can retry") } } diff --git a/server/internal/daemon/daemon.go b/server/internal/daemon/daemon.go index c1053241c38..b7e30f78a2d 100644 --- a/server/internal/daemon/daemon.go +++ b/server/internal/daemon/daemon.go @@ -215,6 +215,14 @@ type repoCacheBackend interface { CreateWorktree(params repocache.WorktreeParams) (*repocache.WorktreeResult, error) } +type claimBarrierOwner uint8 + +const ( + claimBarrierNone claimBarrierOwner = iota + claimBarrierDrain + claimBarrierUpdate +) + // Daemon is the local agent runtime that polls for and executes tasks. type Daemon struct { cfg Config @@ -299,25 +307,20 @@ type Daemon struct { cancelFunc context.CancelFunc // set by Run(); called by triggerRestart rootCtx context.Context // set by Run(); used by long-running recoveries that must survive per-runtime ctx cancellation restartBinary string // non-empty after a successful update; path to the new binary - updating atomic.Bool // prevents concurrent update attempts activeTasks atomic.Int64 // number of tasks currently in handleTask; exposed via /health ready atomic.Bool // false until preflight completes; gates /health status (starting -> running) - // claimMu guards pauseClaims and claimsInFlight. It is held only for the - // microseconds it takes to make a decision; ClaimTask itself runs without - // the lock so a slow per-runtime claim cannot stall auto-update or any - // other poller. + // claimMu guards claimBarrierOwner and claimsInFlight. It is held only for + // the microseconds it takes to make a decision; ClaimTask itself runs + // without the lock so a slow per-runtime claim cannot stall lifecycle work. // - // The pair is the auto-update path's barrier against the issue's - // requirement that "升级过程中如果有 task 进来,会延后升级而不是中断 task": - // runRuntimePoller refuses to call ClaimTask while pauseClaims is set, and - // tryAutoUpdate refuses to flip pauseClaims while any poller is mid-claim - // or any task is in handleTask. Together that closes the fetch-then-claim - // race where a new task slipping in during the release-metadata fetch - // would be cancelled by triggerRestart's root-ctx cancel. - claimMu sync.Mutex - pauseClaims bool // when true, the batch poller skips claiming - claimsInFlight int // pollers that have decided to claim but haven't yet handed the task off to handleTask + // One explicit owner coordinates manual drain, periodic auto-update, and + // heartbeat-triggered update. Any non-none owner pauses new claims, and + // owner-checked release prevents one operation from clearing another's + // barrier. + claimMu sync.Mutex + claimBarrierOwner claimBarrierOwner + claimsInFlight int // pollers that have decided to claim but haven't yet handed the task off to handleTask activeEnvRootsMu sync.Mutex activeEnvRootsCond *sync.Cond // signalled when an in-flight env-root GC mutation finishes @@ -2606,12 +2609,19 @@ func (d *Daemon) handleUpdate(ctx context.Context, runtimeID string, update *Pen return } - // Prevent concurrent update attempts. - if !d.updating.CompareAndSwap(false, true) { - d.logger.Warn("update already in progress, ignoring", "runtime_id", runtimeID, "update_id", update.ID) + // Claim the same lifecycle owner used by drain and periodic auto-update. + // A skipped heartbeat update remains pending server-side and will be + // offered again; barrier contention is not an update failure. + if !d.tryBeginUpdate(false) { + d.logger.Warn("update deferred: another lifecycle operation owns the claim barrier", "runtime_id", runtimeID, "update_id", update.ID) return } - defer d.updating.Store(false) + keepBarrier := false + defer func() { + if !keepBarrier { + d.releaseUpdate() + } + }() d.logger.Info("CLI update requested", "runtime_id", runtimeID, "update_id", update.ID, "target_version", update.TargetVersion) @@ -2638,13 +2648,14 @@ func (d *Daemon) handleUpdate(ctx context.Context, runtimeID string, update *Pen // Trigger daemon restart with the new binary. d.triggerRestart() + keepBarrier = d.RestartBinary() != "" } // runUpdate executes the brew-or-download upgrade against targetVersion and // returns the human-readable output (always populated, even on failure when // brew gives us a useful diagnostic). The caller is responsible for the -// `updating` CAS guard and for reporting status back to the server / triggering -// the restart — extracted so the server-triggered path (handleUpdate) and the +// lifecycle owner and for reporting status back to the server / triggering the +// restart — extracted so the server-triggered path (handleUpdate) and the // auto-update poller (autoUpdateLoop) share the exact same execution body. func (d *Daemon) runUpdate(targetVersion string) (string, error) { if cli.IsBrewInstall() { @@ -2720,14 +2731,14 @@ func (d *Daemon) reportUpdateResultWithRetry(ctx context.Context, runtimeID, upd } // tryEnterClaim records the intent to call ClaimTask. Returns true if the -// caller may proceed, false if the auto-update barrier is in effect. Every +// caller may proceed, false if a lifecycle barrier is in effect. Every // successful call MUST be paired with an exitClaim() on every exit path — // either right after a failed/empty claim, or via the handleTask goroutine's // defer once the task is handed off. func (d *Daemon) tryEnterClaim() bool { d.claimMu.Lock() defer d.claimMu.Unlock() - if d.pauseClaims { + if d.claimBarrierOwner != claimBarrierNone { return false } d.claimsInFlight++ @@ -2741,31 +2752,77 @@ func (d *Daemon) exitClaim() { d.claimsInFlight-- } -// trySetClaimBarrier atomically pauses new ClaimTask calls if the daemon is -// fully idle (no claims in flight, no tasks running). Returns true if the -// caller now holds the barrier and must release it with releaseClaimBarrier -// on every non-restart exit path; false if the daemon is busy and the caller -// should defer to the next tick. Used by tryAutoUpdate to close the race -// where a task slips in between the cheap pre-fetch idle check and the -// actual upgrade kick-off. -func (d *Daemon) trySetClaimBarrier() bool { +// tryAcquireClaimBarrier atomically assigns the claim barrier to owner. When +// requireIdle is true, it also rejects claims in flight and active tasks. +func (d *Daemon) tryAcquireClaimBarrier(owner claimBarrierOwner, requireIdle bool) bool { + if owner == claimBarrierNone { + return false + } d.claimMu.Lock() defer d.claimMu.Unlock() - if d.claimsInFlight > 0 || d.activeTasks.Load() > 0 { + if d.claimBarrierOwner != claimBarrierNone { + return false + } + if requireIdle && (d.claimsInFlight > 0 || d.activeTasks.Load() > 0) { return false } - d.pauseClaims = true + d.claimBarrierOwner = owner return true } -// releaseClaimBarrier clears the auto-update claim barrier so pollers may -// resume claiming. Called on failure paths only — a successful upgrade leaves -// the barrier set because triggerRestart is about to take the process down -// and clearing it would open a window for new claims during shutdown. -func (d *Daemon) releaseClaimBarrier() { +// releaseClaimBarrier clears the barrier only when the caller still owns it. +// A mismatched release is a no-op so one lifecycle path cannot resume claims +// underneath another. +func (d *Daemon) releaseClaimBarrier(owner claimBarrierOwner) { + d.claimMu.Lock() + defer d.claimMu.Unlock() + if d.claimBarrierOwner == owner { + d.claimBarrierOwner = claimBarrierNone + } +} + +func (d *Daemon) claimBarrierOwnedBy(owner claimBarrierOwner) bool { d.claimMu.Lock() defer d.claimMu.Unlock() - d.pauseClaims = false + return d.claimBarrierOwner == owner +} + +// tryBeginDrain atomically pauses new claims for a client-owned drain +// shutdown. Unlike the auto-update barrier, an active task or in-flight claim +// is expected here: the drain waits for both to finish naturally. A false +// return means another drain or auto-update already owns the claim pause. +func (d *Daemon) tryBeginDrain() bool { + return d.tryAcquireClaimBarrier(claimBarrierDrain, false) +} + +// drainIdle reports whether every claim that entered before the drain barrier +// has completed its handoff and every resulting task has finished. The batch +// poller increments activeTasks before exitClaim, so observing both values at +// zero under claimMu cannot miss a task between the two counters. +func (d *Daemon) drainIdle() bool { + d.claimMu.Lock() + defer d.claimMu.Unlock() + return d.claimBarrierOwner == claimBarrierDrain && d.claimsInFlight == 0 && d.activeTasks.Load() == 0 +} + +func (d *Daemon) releaseDrain() { + d.releaseClaimBarrier(claimBarrierDrain) +} + +func (d *Daemon) isDraining() bool { + return d.claimBarrierOwnedBy(claimBarrierDrain) +} + +func (d *Daemon) tryBeginUpdate(requireIdle bool) bool { + return d.tryAcquireClaimBarrier(claimBarrierUpdate, requireIdle) +} + +func (d *Daemon) releaseUpdate() { + d.releaseClaimBarrier(claimBarrierUpdate) +} + +func (d *Daemon) isUpdating() bool { + return d.claimBarrierOwnedBy(claimBarrierUpdate) } // triggerRestart initiates a graceful daemon restart after a successful CLI update. diff --git a/server/internal/daemon/health.go b/server/internal/daemon/health.go index e62da8164e2..4c9670ede6b 100644 --- a/server/internal/daemon/health.go +++ b/server/internal/daemon/health.go @@ -32,6 +32,7 @@ type HealthResponse struct { ServerURL string `json:"server_url"` CLIVersion string `json:"cli_version"` ActiveTaskCount int64 `json:"active_task_count"` + Draining bool `json:"draining"` Agents []string `json:"agents"` Workspaces []healthWorkspace `json:"workspaces"` } @@ -103,6 +104,7 @@ func (d *Daemon) healthHandler(startedAt time.Time) http.HandlerFunc { ServerURL: d.cfg.ServerBaseURL, CLIVersion: d.cfg.CLIVersion, ActiveTaskCount: d.activeTasks.Load(), + Draining: d.isDraining(), Agents: agents, Workspaces: wsList, } @@ -134,12 +136,71 @@ func (d *Daemon) shutdownHandler() http.HandlerFunc { } } +// drainShutdownHandler pauses new claims and waits for work already accepted +// by the daemon to finish before cancelling its root context. This is a +// dedicated endpoint rather than a /shutdown query parameter so a newer CLI +// fails safely with 404 against an older daemon instead of having the old +// handler ignore the parameter and perform an immediate shutdown. +func (d *Daemon) drainShutdownHandler() http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + if !d.tryBeginDrain() { + http.Error(w, "another operation is already pausing task claims", http.StatusConflict) + return + } + keepBarrier := false + defer func() { + if !keepBarrier { + d.releaseDrain() + } + }() + + ticker := time.NewTicker(100 * time.Millisecond) + defer ticker.Stop() + for { + select { + case <-r.Context().Done(): + return + default: + } + if d.drainIdle() { + break + } + select { + case <-r.Context().Done(): + return + case <-ticker.C: + } + } + // Give a client cancellation that raced the final idle observation one + // last chance to release the barrier rather than stop the daemon. + select { + case <-r.Context().Done(): + return + default: + } + + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(map[string]string{"status": "drained"}); err != nil { + return + } + keepBarrier = true + if d.cancelFunc != nil { + go d.cancelFunc() + } + } +} + // serveHealth runs the health HTTP server on the given listener. // Blocks until ctx is cancelled. func (d *Daemon) serveHealth(ctx context.Context, ln net.Listener, startedAt time.Time) { mux := http.NewServeMux() mux.HandleFunc("/health", d.healthHandler(startedAt)) mux.HandleFunc("/shutdown", d.shutdownHandler()) + mux.HandleFunc("/shutdown/drain", d.drainShutdownHandler()) mux.HandleFunc("/repo/checkout", d.repoCheckoutHandler()) srv := &http.Server{Handler: mux} diff --git a/server/internal/daemon/health_test.go b/server/internal/daemon/health_test.go index 5c5c31e6ed4..70a188f88cc 100644 --- a/server/internal/daemon/health_test.go +++ b/server/internal/daemon/health_test.go @@ -53,6 +53,9 @@ func TestHealthHandlerReportsCLIVersionAndActiveTaskCount(t *testing.T) { if got, want := raw["active_task_count"], float64(3); got != want { t.Errorf("active_task_count key: got %v, want %v", got, want) } + if got, want := raw["draining"], false; got != want { + t.Errorf("draining key: got %v, want %v", got, want) + } if got, want := raw["status"], "running"; got != want { t.Errorf("status key: got %v, want %q", got, want) } @@ -157,6 +160,157 @@ func TestShutdownHandlerPostCancelsDaemonContext(t *testing.T) { } } +func TestShutdownHandlerDrainWaitsForActiveTasks(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + d := &Daemon{cancelFunc: cancel} + d.activeTasks.Store(1) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/shutdown/drain", nil) + done := make(chan struct{}) + go func() { + d.drainShutdownHandler().ServeHTTP(rec, req) + close(done) + }() + + waitForDrainState(t, d, true) + if d.tryEnterClaim() { + d.exitClaim() + t.Fatal("new claims must be paused while a drain shutdown is pending") + } + select { + case <-ctx.Done(): + t.Fatal("drain shutdown cancelled the daemon while a task was active") + default: + } + + d.activeTasks.Add(-1) + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("drain shutdown did not finish after active tasks reached zero") + } + if rec.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", rec.Code) + } + select { + case <-ctx.Done(): + case <-time.After(time.Second): + t.Fatal("daemon context was not cancelled after the drain completed") + } +} + +func TestShutdownHandlerDrainCoversClaimToActiveHandoff(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + d := &Daemon{cancelFunc: cancel} + if !d.tryEnterClaim() { + t.Fatal("initial claim should enter before the drain barrier") + } + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/shutdown/drain", nil) + done := make(chan struct{}) + go func() { + d.drainShutdownHandler().ServeHTTP(rec, req) + close(done) + }() + waitForDrainState(t, d, true) + + // Mirror runBatchPoller's handoff invariant: activeTasks is incremented + // before the in-flight claim is released. The drain must observe the task, + // not a transient zero between the two counters. + d.activeTasks.Add(1) + d.exitClaim() + select { + case <-ctx.Done(): + t.Fatal("drain shutdown cancelled during the claim-to-active handoff") + case <-time.After(30 * time.Millisecond): + } + + d.activeTasks.Add(-1) + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("drain shutdown did not finish after the handed-off task completed") + } +} + +func TestShutdownHandlerDrainCancellationResumesClaims(t *testing.T) { + t.Parallel() + + d := &Daemon{} + d.activeTasks.Store(1) + reqCtx, cancelRequest := context.WithCancel(context.Background()) + req := httptest.NewRequest(http.MethodPost, "/shutdown/drain", nil).WithContext(reqCtx) + rec := httptest.NewRecorder() + done := make(chan struct{}) + go func() { + d.drainShutdownHandler().ServeHTTP(rec, req) + close(done) + }() + + waitForDrainState(t, d, true) + cancelRequest() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("drain request did not stop after client cancellation") + } + waitForDrainState(t, d, false) + if !d.tryEnterClaim() { + t.Fatal("claims did not resume after the drain requester disconnected") + } + d.exitClaim() +} + +func TestShutdownHandlerRejectsConcurrentDrain(t *testing.T) { + t.Parallel() + + d := &Daemon{} + d.activeTasks.Store(1) + firstCtx, cancelFirst := context.WithCancel(context.Background()) + firstDone := make(chan struct{}) + go func() { + d.drainShutdownHandler().ServeHTTP( + httptest.NewRecorder(), + httptest.NewRequest(http.MethodPost, "/shutdown/drain", nil).WithContext(firstCtx), + ) + close(firstDone) + }() + waitForDrainState(t, d, true) + + rec := httptest.NewRecorder() + d.drainShutdownHandler().ServeHTTP(rec, httptest.NewRequest(http.MethodPost, "/shutdown/drain", nil)) + if rec.Code != http.StatusConflict { + t.Fatalf("concurrent drain status = %d, want %d", rec.Code, http.StatusConflict) + } + + cancelFirst() + select { + case <-firstDone: + case <-time.After(time.Second): + t.Fatal("first drain did not stop after cancellation") + } +} + +func waitForDrainState(t *testing.T, d *Daemon, want bool) { + t.Helper() + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + if d.isDraining() == want { + return + } + time.Sleep(time.Millisecond) + } + t.Fatalf("draining did not become %v", want) +} + func TestShutdownHandlerRejectsNonPost(t *testing.T) { t.Parallel() diff --git a/server/internal/daemon/update_drain_test.go b/server/internal/daemon/update_drain_test.go new file mode 100644 index 00000000000..a9f30acd697 --- /dev/null +++ b/server/internal/daemon/update_drain_test.go @@ -0,0 +1,138 @@ +package daemon + +import ( + "context" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestHandleUpdateDefersWhileDrainOwnsBarrier(t *testing.T) { + d, reportCalls := updateReportDaemon(t, func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + }) + d.activeTasks.Store(1) + + var updateCalls atomic.Int32 + var restartCalls atomic.Int32 + d.runUpdateFn = func(string) (string, error) { + updateCalls.Add(1) + return "upgraded", nil + } + d.cancelFunc = func() { + restartCalls.Add(1) + } + + drainCtx, cancelDrain := context.WithCancel(context.Background()) + drainDone := make(chan struct{}) + go func() { + d.drainShutdownHandler().ServeHTTP( + httptest.NewRecorder(), + httptest.NewRequest(http.MethodPost, "/shutdown/drain", nil).WithContext(drainCtx), + ) + close(drainDone) + }() + t.Cleanup(func() { + cancelDrain() + select { + case <-drainDone: + case <-time.After(time.Second): + t.Error("drain handler did not stop during cleanup") + } + }) + + waitForDrainState(t, d, true) + d.handleUpdate(context.Background(), "runtime-1", &PendingUpdate{ + ID: "update-1", + TargetVersion: "v0.4.5", + }) + + if got := updateCalls.Load(); got != 0 { + t.Fatalf("heartbeat update ran %d time(s) while drain owned the barrier", got) + } + if got := restartCalls.Load(); got != 0 { + t.Fatalf("heartbeat update restarted %d time(s) while drain owned the barrier", got) + } + if got := atomic.LoadInt32(reportCalls); got != 0 { + t.Fatalf("deferred heartbeat update reported %d status update(s), want none", got) + } + if !d.isDraining() { + t.Fatal("heartbeat update stole the drain barrier") + } +} + +func TestDrainReturnsConflictWhileHandleUpdateOwnsBarrier(t *testing.T) { + d, _ := updateReportDaemon(t, func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + }) + d.activeTasks.Store(1) + + updateStarted := make(chan struct{}) + releaseUpdate := make(chan struct{}) + updateDone := make(chan struct{}) + var releaseOnce sync.Once + release := func() { + releaseOnce.Do(func() { close(releaseUpdate) }) + } + t.Cleanup(release) + + var restartCalls atomic.Int32 + d.cancelFunc = func() { + restartCalls.Add(1) + } + d.runUpdateFn = func(string) (string, error) { + close(updateStarted) + <-releaseUpdate + return "upgraded", nil + } + + go func() { + d.handleUpdate(context.Background(), "runtime-1", &PendingUpdate{ + ID: "update-1", + TargetVersion: "v0.4.5", + }) + close(updateDone) + }() + + select { + case <-updateStarted: + case <-time.After(time.Second): + t.Fatal("heartbeat update did not start") + } + + drainCtx, cancelDrain := context.WithCancel(context.Background()) + rec := httptest.NewRecorder() + drainDone := make(chan struct{}) + go func() { + d.drainShutdownHandler().ServeHTTP( + rec, + httptest.NewRequest(http.MethodPost, "/shutdown/drain", nil).WithContext(drainCtx), + ) + close(drainDone) + }() + + select { + case <-drainDone: + case <-time.After(200 * time.Millisecond): + cancelDrain() + <-drainDone + t.Fatal("drain blocked instead of rejecting an in-progress heartbeat update") + } + cancelDrain() + if rec.Code != http.StatusConflict { + t.Fatalf("drain status = %d, want %d while update owns the barrier", rec.Code, http.StatusConflict) + } + + release() + select { + case <-updateDone: + case <-time.After(time.Second): + t.Fatal("heartbeat update did not finish after release") + } + if got := restartCalls.Load(); got != 1 { + t.Fatalf("heartbeat update restarted %d time(s), want 1", got) + } +}