Fix: fence register/unregister and make shutdown sticky

[ChaoWao]

Summary

Two coupled admission defects: a window on the host side, and the consequence it produces on the wire.

register / unregister mutate the registry outside the operation lease. close() publishes CLOSED and then drains _active_ops before it touches the tree, so any registry mutation made outside a lease is invisible to that drain.

• register() published the handle before taking its lease on the L3+ post-init ChipCallable path and on the L2 pre-warm path.
• unregister() held no lease at all — it popped _live_handles, decremented ref_count, set the tombstone, and then broadcast bare.

A close() could therefore pass admission mid-transaction and start tearing the tree down underneath either one. The existing _rollback_handle_locked did not cover it: it fires only when the broadcast itself raises, and a close() winning the race does not make the broadcast raise.

SHUTDOWN is overwritable. The mailbox state word has three writers — the parent's CONTROL_REQUEST, the child's CONTROL_DONE, and run_control_command's return to IDLE. The CONTROL_DONE a child publishes for an in-flight control command lands after a concurrent SHUTDOWN store and erases it, so the child polls a mailbox whose request is gone until the parent's reap deadline expires and reports it as a survivor. The C++ side already worked around this by repeating the store (poisoned_progress_quiesced), which narrows the window rather than closing it.

Changes

• register(): post-READY L3+ and L2 paths take _operation_lease("register") before publication, so publication and the child-facing call are one transaction. The three lease sites (remote / python / chip) now enter at the same point. Pre-start (NEW) and INITIALIZING registrations are untouched — they keep _register_into_snapshot_or_wait, since the lease admits only READY and would break pre-start registration.
• unregister(): the live-target paths (L3+ child broadcast, L2 device-slot release) take a symmetric lease. Whether the target is reachable is now decided once by the caller under that lease instead of being re-read from _initialized / _hierarchical_started across the broadcast. A registry-only decrement keeps the no-lease path so it still works before init().
• L2 is not exempt (called out explicitly since it was an open question): it has no chip subtree to broadcast to, but its pre-warm does drive the live ChipWorker, and pre-fix it had the same publish-then-lease window with no rollback — a close() landing there left the entry published on a CLOSED worker. It is now fenced too. Its failure semantics are unchanged; no rollback was added.
• MAILBOX_OFF_SHUTDOWN: a sticky one-way word written only by a terminating parent, 0 -> 1, never cleared. Both terminating writers set it — Python _request_child_shutdown (now the sole writer of the request, used by _broadcast_child_shutdown and the startup-rollback graceful phase) and C++ LocalMailboxEndpoint::shutdown_child — sticky word first, then the state word. Both serve loops (_run_mailbox_loop, the chip two-frame loop) read it at the top of every iteration alongside the state word.

Wire delta

One new int32 at MAILBOX_OFF_SHUTDOWN = MAILBOX_OFF_FRAME_PROTOCOL - 8, reserved on every frame by shrinking MAILBOX_ARGS_CAPACITY by 8 bytes (64024 -> 64016 on the post-#1648 64 KiB frame; the max-blob static_assert has ~30 KiB of slack). No existing offset moves, and run_control_command's IDLE / CONTROL_REQUEST writes need no change — the new word is not in their write set, which is why this is preferable to converting the state word to CAS.

No path reuses a mailbox after a shutdown request: mailbox shms are created fresh per child, and every shutdown_child() caller (stop_workers, fail_progress_driver, poisoned_progress_quiesced) is terminal.

Testing

New tests/ut/py/test_worker/test_admission_fence.py — 5 device-free tests (L3 worker, one SUB child, no chips), each with its own hard timeout. 4 of the 5 fail against the previous code.

1. test_publication_holds_the_lease — register publishes under a lease (red before)
2. test_close_drains_a_register_in_flight — close() cannot tear down mid-broadcast
3. test_broadcast_holds_the_lease — unregister broadcasts under a lease (red before)
4. test_close_cannot_slip_into_the_unregister_broadcast — close() returned in 3.7 ms before; now fenced (red before)
5. test_shutdown_survives_an_in_flight_control_command — forks a real _run_mailbox_loop, holds its control handler open, requests shutdown underneath it, releases; the child must exit. Asserts the state word really did end up CONTROL_DONE, so it cannot pass by missing the race (red before — the child hung to the test budget)

Test 2 is the one that passes both ways, and deliberately so: a park at publication is held under _registry_lock, which close()'s registry detach also takes, so it would delay close() even with no lease at all and cannot discriminate. It is retargeted at the broadcast and kept as the symmetric partner of test 4; test 1 is the actual register-side barrier.

test_host_buffer_registration.py scripts the serve loop's poll sequence through a monkeypatched _mailbox_load_i32; the loop now reads two words per iteration, so its fake answers the shutdown address separately.

• tests/ut/py — 1057 passed, 13 skipped
• C++ unit tests — 76/76 (ctest -LE requires_hardware)
• Hardware tests — onboard a2a3 sweep (-m "not sdma", via task-submit) exit 0, 62 groups PASS / 0 fail, L2 host_build_graph and L2 tensormap_and_ringbuffer both PASS; quarantined -m sdma set exit 0
• pre-commit clean on the changed files

All of the above re-run after rebasing onto upstream/main (#1648 resized the mailbox frame 32 KiB -> 64 KiB, which these offsets derive from).

Notes

This lands the first two exit conditions of the P0.2 lifecycle-hardening follow-up (symmetric register / unregister fencing, and a SHUTDOWN that cannot be overwritten). Not in scope: the eligibility revalidation, the _CloseAttempt / cancellation cluster, and the retryable commit-barrier journal.

The lost-shutdown mechanism is a candidate, unproven cause of #1562; this PR is not claiming to fix it, and its acceptance is the regression tests above.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 979026bf-9455-42f1-8ca6-c6ad490738ef

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

The mailbox protocol now uses a sticky shutdown marker separate from mailbox state. Worker loops observe the marker during control processing. Registration and unregistration operations retain leases across publication and cleanup transactions. Tests cover admission fencing and shutdown races.

Changes

Worker lifecycle and mailbox coordination

Layer / File(s)	Summary
Mailbox shutdown protocol src/common/hierarchical/worker_manager.h, src/common/hierarchical/worker_manager.cpp, python/simpler/worker.py	Mailbox frames reserve a sticky shutdown word. Argument capacity ends before that word. Shutdown publication writes the marker before SHUTDOWN.
Shutdown detection and propagation python/simpler/worker.py, tests/ut/py/test_worker/test_host_buffer_registration.py, tests/ut/py/test_worker/test_admission_fence.py	Child loops and chip loops observe sticky shutdown requests. Startup rollback and group shutdown use the shutdown helper. Tests cover shutdown during an in-flight control command.
Registration and unregistration admission fences python/simpler/worker.py, tests/ut/py/test_worker/test_admission_fence.py	Registration and unregistration retain leases across publication, broadcast, device-slot, and cleanup transactions. Tests verify that close() waits for in-flight operations.
Mailbox protocol documentation docs/worker-manager.md	The mailbox layout and sticky shutdown behavior are documented.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant ParentWorker
  participant LocalMailboxEndpoint
  participant ChildMailboxLoop
  participant ControlCommand
  ParentWorker->>LocalMailboxEndpoint: request shutdown
  LocalMailboxEndpoint->>ChildMailboxLoop: publish sticky shutdown marker
  LocalMailboxEndpoint->>ChildMailboxLoop: publish SHUTDOWN state
  ChildMailboxLoop->>ControlCommand: finish in-flight control command
  ChildMailboxLoop->>ChildMailboxLoop: recheck sticky shutdown marker
  ChildMailboxLoop-->>ParentWorker: terminate successfully

Loading

Possibly related PRs

• hw-native-sys/simpler#1398: Modifies worker lifecycle, registration, unregistration, and close() concurrency handling.
• hw-native-sys/simpler#1495: Modifies forked-child mailbox loops and shutdown handling.
• hw-native-sys/simpler#1583: Modifies the hierarchical worker mailbox protocol and frame state handling.

Poem

A rabbit guards the shutdown sign,
A sticky mark stays safe in line.
Leases hold while workers tread,
No closing race can leap ahead.
Frames now rest when tasks are done.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes both primary changes: admission fencing for register/unregister and sticky shutdown handling.
Description check	✅ Passed	The description directly explains the lifecycle defects, implementation changes, wire layout, and validation results.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@python/simpler/worker.py`:
- Around line 5225-5241: Update _post_start_register_l2 so failures from
_chip_worker._register_callable_at_slot roll back the newly installed
registration, mirroring _post_start_register_chip. Only remove the registration
when is_new is true and prewarming raises, then propagate the original exception
while preserving successful registration behavior.

In `@tests/ut/py/test_worker/test_admission_fence.py`:
- Around line 63-66: Update the exception suppression in the worker cleanup
try/except around w.close() to also silence Ruff’s S110 warning, while retaining
the existing BLE001 suppression and its explanatory reason. Keep the intentional
swallowing of close() errors unchanged because the worker may already be closed.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 8c265d6d-5916-4d8f-9d34-ac8888cc9a90

📥 Commits

Reviewing files that changed from the base of the PR and between b5261a7 and 06c9ba6.

📒 Files selected for processing (6)

• docs/worker-manager.md
• python/simpler/worker.py
• src/common/hierarchical/worker_manager.cpp
• src/common/hierarchical/worker_manager.h
• tests/ut/py/test_worker/test_admission_fence.py
• tests/ut/py/test_worker/test_host_buffer_registration.py