fix(operator): hand an apply back cleanly when a process shuts down

[aparajon]

Why this matters

An engine that runs its schema change inside the SchemaBot process holds resources on the target — for Spirit, an advisory lock on the table it is copying — for as long as that work lives. That work outlives the drive that started it: the drive runs behind a detached context, so cancelling the drive stops the heartbeat, not the copy.

Shutdown cancelled the drives and exited. The process stopped renewing its applies' leases while the targets stayed held. A minute later the leases went stale, peer drivers reclaimed the applies, every one of them was refused the lock, and each refusal charged a recovery attempt against the apply's budget.

Before

shutdown
  └── cancel drive ctx ──┬── heartbeat stops          ← lease starts aging
                         └── engine keeps copying     ← target stays locked
                                    │
                   60s later        ▼
        lease goes stale ──> peer driver claims ──> refused the lock
                                                     └── recovery attempt burned
                                                         (repeats until exhausted)

After

shutdown
  ├── 1. stop claiming            no driver picks up new work
  ├── 2. cancel drive ctx         drives return, apply left active — not stopped
  ├── 3. halt in-process engines  checkpoint, stop the copy, target released
  └── 4. hand the claims back     apply claimable on the very next poll

What it does

Tells a drive's shutdown apart from an operator stop. Both ended the sequential drive early through the same outcome, so a drive that lost its context was on course to be recorded as an operator stop — parking every apply a restart interrupted until a human ran start. Cancellation now has its own outcome: the drive exits leaving the apply active, and logs which of the two happened.

Brings in-process engines down before shutdown ends. Stopping the operator is phased: stop claiming, end the in-flight drives, then halt the engines and wait for them to release their targets. Halting is checkpointed and resumable, and records no operator intent — it is a handover, not a stop. An engine whose work runs elsewhere declares no halt capability, so its schema change is left alone and the lease handover stays the only thing that happens. A halt that does not complete inside its bound is logged with the endpoint and counted (schemabot.operator.shutdown_halt_failures), so a shutdown that leaves a target held is visible before the refusals it causes are.

Hands claimed applies back. A departing process left its applies leased with nothing renewing them, so the work sat idle for the whole staleness window before any peer could pick it up — a restart deferred every in-flight schema change by a minute. Shutdown now releases the claims once the targets are free. A driver registers the claim it drives under and deregisters it on the way out, guarded on the lease token so it never hands back a claim a peer has since rotated onto itself. An apply that settled while shutdown was in progress is left alone: it has nothing to hand over, and backdating its heartbeat would misreport when it finished.

The ordering is the invariant: the claims are released only after the engines are down, so a peer driver is never invited onto a target this process still holds.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

This PR improves SchemaBot’s operator shutdown behavior so in-process schema-change engines (notably Spirit) are halted before the process exits, and so active applies are handed back immediately instead of waiting for lease staleness—preventing lock refusals that burn recovery attempts and reducing downtime after restarts.

Changes:

• Distinguish drive-context cancellation (shutdown handover) from an operator stop so shutdown doesn’t park applies in stopped.
• Add shutdown halting for in-process engines (router + local client + Spirit implementation) and record halt failures via a new metric.
• Add storage support to explicitly release apply claims (backdate heartbeat + clear lease fields) and wire shutdown to release held claims after engines halt.

Reviewed changes

Copilot reviewed 19 out of 19 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
pkg/tern/target_router.go	Adds HaltForShutdown to halt cached routed clients that support shutdown halting.
pkg/tern/local_control_resume.go	Treats taskHandover as an early-return outcome during resume paths.
pkg/tern/local_client.go	Implements client-level HaltForShutdown by delegating to engine shutdown-halting capability.
pkg/tern/local_apply_sequential.go	Introduces taskHandover and updates sequential execution/polling to hand over on drive cancellation.
pkg/tern/local_apply_handover_test.go	Adds tests ensuring drive cancellation is not recorded as an operator stop and leaves applies active.
pkg/tern/local_apply_grouped.go	Logs and returns on drive cancellation while polling grouped/atomic applies (handover semantics).
pkg/tern/client.go	Defines optional tern.ShutdownHalter capability for clients with in-process engines.
pkg/storage/storage.go	Extends ApplyStore with ReleaseClaim for explicit lease handback.
pkg/storage/internal/sqlstore/applies.go	Implements ApplyStore.ReleaseClaim guarded on lease token with stale backdate.
pkg/storage/internal/sqlstore/applies_test.go	Adds tests for reclaimability, token mismatch no-op, and invalid-lease refusal.
pkg/metrics/metrics.go	Adds schemabot.operator.shutdown_halt_failures counter and recorder.
pkg/engine/spirit/spirit.go	Implements engine.ShutdownHalter to checkpoint, cancel, and wait for Spirit copy goroutine to exit.
pkg/engine/spirit/shutdown_halt_test.go	Adds unit tests for Spirit shutdown halting behavior and bounded failure on deadline.
pkg/engine/engine.go	Adds engine.ShutdownHalter interface and HaltEngineForShutdown helper.
pkg/engine/engine_test.go	Tests HaltEngineForShutdown behavior for engines with/without shutdown-halting capability.
pkg/api/service.go	Adds service tracking for held apply leases to support explicit shutdown handback.
pkg/api/operator.go	Implements phased shutdown: stop claiming, cancel drives, halt engines, then release held claims; tracks leases per drive.
pkg/api/operator_shutdown_halt_test.go	Adds tests for shutdown halt ordering, failure fanout, skipping remote clients, and claim handback semantics.
pkg/api/handlers_test.go	Extends capturingApplyStore test double to support Get/ReleaseClaim for shutdown tests.

Suppressed comments (2)

pkg/tern/local_apply_sequential.go:138

• checkTaskReady can mis-handle a drive cancellation race: if ctx is canceled after the initial ctx.Err() check but before Tasks().Get returns, the context-canceled error currently maps to taskSkip. In the single-task case this can fall through to finalizeSequentialApply and incorrectly mark the apply completed even though no task ran. Treat storage.Get errors caused by ctx cancellation as taskHandover instead.

	if err != nil {
		logger.Error("failed to fetch task state",
			"task_id", task.TaskIdentifier, "table", task.TableName, "state", task.State, "error", err)
		return taskSkip
	}

pkg/api/operator.go:208

• This debug log attribute includes the internal numeric apply row ID (apply_row_lookup_id), which the repo conventions discourage logging (AGENTS.md:219). Since apply is nil here, log the held lease token/owner instead.

		if apply == nil {
			s.logger.Debug("operator: a claimed apply no longer exists; nothing to hand back",
				"apply_row_lookup_id", lease.ApplyID, "lease_owner", lease.Owner)
			continue

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Kiran01bm]

🤖 Review findings - created by Kiran's code review agent - for pull/1019, 188aea5.

Verdict: 8 findings — 2 blocking (the operation lease is never handed back, and the ordering invariant is untested), 5 non-blocking, 1 suggestion. CI is fully green: 34/34 checks pass at head 188aea5 (state OPEN) — build, lint x5, unit, integration, all E2E matrices (MySQL, Vitess x3, K8s x3, gRPC x3, LocalScale x3), Semgrep, zizmor, DCO all green at verification time.

Blocking

1. Shutdown hands back only the apply lease, never the operation lease the mainline drive actually claims under. Normal drives enter via FindNextApplyOperation and heartbeat the operation row, but releaseHeldClaims calls only Applies().ReleaseClaim (operator.go:225), and a peer can only reclaim an active operation once updated_at is ~60s stale (apply_operations.go:933-934). So after a restart mid-apply the work idles the full staleness window exactly as pre-PR, while the log at :236 claims next-poll pickup; ApplyOperationStore.ReleaseClaim already exists (apply_operations.go:1371) and should be invoked here too.

2. The headline invariant — engines halted before claims are handed back — is pinned by no test. The one ordering assertion is vacuous: the service is built with an empty store and a 1h poll interval, so driversSeen is 0 under any StopOperator ordering (operator_shutdown_halt_test.go:45-54), and both releaseHeldClaims tests call the method directly, bypassing StopOperator (:96). Reordering the release ahead of the halt — or deleting the release call — passes the whole suite; add an end-to-end test that drives StopOperator with a registered held claim and asserts the halt happens first.

Non-blocking

3. trackHeldClaim registers whatever lease token the parent apply row happens to carry, with no ownership check. On operation-lease-only drives the parent is loaded with a plain Get (operator.go:760) yet resumeClaimedApplyWithOptions unconditionally tracks it (:1405), so shutdown can release a residual token this process never owned, backdating the parent 61s while peers still drive sibling operations. Gate registration on storage.LeaseOwnedByThisProcess (exists, unused here); today a live-peer double-drive is fenced only incidentally by the claim SQL's active-operation exclusions.

4. The claims are handed back even when the engine halt failed. haltEnginesForShutdown returns void — a wedged runner past the 20s bound is only logged and counted (operator.go:253) — yet StopOperator releases every claim unconditionally (:146), inviting a peer onto a target this process provably still holds within seconds instead of after 60s. That is the exact refusal-and-burned-attempt loop the PR set out to eliminate, made deterministic on the halt-failure path; skip (or defer) the release for claims whose halt failed.

5. A claim taken between the snapshot and the cancel is silently never handed back. heldClaimsSnapshot() runs after close(stop) but before cancel() (operator.go:120-131), so a drive already past its stop-channel check can register a claim after the snapshot, then deregister it on return (:182-184) — leaving it in neither the snapshot nor the map. The freshest claim falls back to the pre-PR 60s window with no log; either snapshot after quiescence or have the shutdown-path deregister feed a to-release set.

6. Four of the six releaseHeldClaims branches are untested — the Get-error and missing-apply continues, the ReleaseClaim error, and released==false (operator.go:211-234). A regression like a continue becoming a return would silently abort the handback of all remaining claims after one failed read and still pass the suite.

7. Spirit's checkpoint-before-cancel behavior is untested. Both halt tests construct runningSchemaChange with no runners, so the DumpCheckpoint call and its warn-and-continue error path (spirit.go:344-348) never execute; reordering the cancel ahead of the checkpoint, or early-returning on checkpoint error (leaving the copy goroutine and advisory lock alive), passes everything.

General suggestions

8. Docs still describe the pre-PR behavior. architecture.md:733 says a running apply becomes claimable only after >1 minute of heartbeat staleness, and apply-lifecycle.md:158 still lists a pod restart mid-drive as a transient failure that burns a recovery attempt — both now describe only unclean kills.

The one thing that could have broken, verified

The riskiest mechanism is shutdown releasing an apply's claim while something still holds the work — an invitation for a peer to double-drive. Tracing StopOperator end-to-end (operator.go:118-146) shows the success path is sound: close(stop) blocks new claims, cancel() plus recoveryWg.Wait() joins every drive and heartbeat goroutine, the 20s-bounded engine halt releases the advisory lock well inside the 60s lease window, and only then does ReleaseClaim run — with the token-scoped UPDATE (applies.go:1958) and the terminal-state reload preventing handback of rotated or settled work. But the safety is one-directional: the handback still runs when the halt failed (finding 4), registers tokens this process never owned on operation-lease-only drives (finding 3), and is largely ineffective on the mainline path because the operation lease that actually gates re-claim is never released (finding 1). All three leaks are reported above, and the ordering itself is guarded by no non-vacuous test (finding 2).

Verified correct

• StopOperator success-path ordering is sound and no post-release heartbeat can re-freshen a released row — all drive/heartbeat goroutines are joined before the release (operator.go:118-146).
• The handback fires only on shutdown, never on drive failure — releaseHeldClaims is referenced solely from StopOperator (operator.go:146).
• applyStore.ReleaseClaim fails closed, is token-scoped, and its backdate beats the claim query's strict < on both dialects (applies.go:1945-1967).
• trackHeldClaim deregistration is token-guarded and invalid leases are never registered (operator.go:182-184).
• releaseHeldClaims leaves settled and missing applies alone, so a terminal apply's settle timestamp is never backdated (operator.go:216-224).
• haltEnginesForShutdown attempts every client despite failures and skips remote (non-ShutdownHalter) clients (operator.go:266-280).
• The halt is a handover, not a stop: Spirit's HaltForShutdown checkpoints before cancelling and never writes rm.state (spirit.go:326-372), and tern returns taskHandover on ctx cancellation while storage-observed stops still park the apply (local_apply_sequential.go:128-146).
• CI: 34/34 checks pass at head 188aea5 — build, lint x5, unit, integration, all E2E matrices (MySQL, Vitess x3, K8s x3, gRPC x3, LocalScale x3), Semgrep, zizmor, DCO green, including the integration test proving immediate re-claim after ReleaseClaim on both dialects.

This review was generated by Claude Code (claude-fable-5).

[aparajon]

🤖 Thanks — findings 1 and 2 were both right, and 1 was the load-bearing one: releasing only the apply lease left the mainline drive's work unclaimable for the full staleness window, since the operation claim is what the next poll looks at first. Both leases are now handed back, registered once in recoverApplyOperation right after the operation lease validates, so every drive path below it is covered.

• 1 — operation claims tracked and released (51c4f3a6), ordered after the apply claims so a peer never hits reconcileUnclaimableParent.
• 2 — TestStopOperatorHandsClaimsBackOnlyAfterEnginesAreDown drives the real StopOperator with a registered claim and asserts the recorded sequence. I mutation-tested it: moving the release ahead of the halt fails with the swapped order.
• 3 — gated on storage.LeaseOwnedByThisProcess (c4b0c553).
• 4 — the halt now returns an outcome per deployment; a claim whose engine did not come down is left to go stale, with a warn saying why. An unattributable failure (the unkeyed default client) holds every claim.
• 5 — a drain flag replaces the snapshot: a drive returning during shutdown leaves its claim registered instead of deregistering into the gap.
• 6 — all four branches covered.
• 8 — both docs updated.

7 I have not done in this PR, and will follow up with a dedicated one. runningSchemaChange.runners is []*spiritmigration.Runner, a concrete Spirit type with no injection seam, and nothing in pkg/engine/spirit drives a real runner mid-copy today — the integration tests there are plan-level. Covering the checkpoint-before-cancel ordering needs a new real-MySQL mid-flight harness, which I will build as its own change rather than bolt onto this one.

Reviewed and addressed by Claude Code (claude-opus-5).