Prevent stale device lifecycle commits after removal

[dmulcahey]

Summary

This PR fixes a device lifecycle race where initialization or reinterview work could outlive device removal and later recreate state or emit DeviceFullInitEvent.

It establishes this invariant:

Once DeviceRemovedEvent is published for a device lifecycle, no older initialization or reinterview operation for that IEEE may recreate its wrapper state or emit DeviceFullInitEvent.

Shutdown applies the same ownership rule: once shutdown begins, device lifecycle work must quiesce before wrapper or controller teardown proceeds.

Why cancellation alone is insufficient

asyncio.Task.cancel() requests cancellation; it does not synchronously stop the task.

CancelledError is delivered when the task reaches a cancellation point. Before the task finishes, it may:

• still be running before its next await;
• catch CancelledError to perform asynchronous cleanup;
• remain blocked in code that delays or suppresses cancellation.

The previous removal path cancelled initialization and immediately removed the wrapper. That left a window in which the old task could resume after removal, commit state, or emit a full-initialization event.

The synchronization boundary is therefore not the call to cancel(). It is completion of the exact task that owned the old lifecycle.

Implementation

The gateway now coordinates initialization, reinterview, removal, and shutdown using explicit per-device lifecycle ownership:

• Each IEEE has a monotonically increasing lifecycle generation.
• Initialization and reinterview operations capture their generation when scheduled.
• Removal advances the generation before requesting cancellation, immediately revoking the old operation’s authority to commit.
• A removal-in-progress task acts as a tombstone and synchronization barrier for later work using the same IEEE.
• Removal cancels and awaits the exact superseded initialization or reinterview task before removing its captured wrapper and publishing DeviceRemovedEvent.
• A per-IEEE lock serializes lifecycle mutations and prevents a new join from overlapping an older removal and teardown.
• Initialization and factory-changing reinterview paths revalidate their generation and exact wrapper identity after awaited work and before state changes or lifecycle events.
• Task completion callbacks use task identity, so an older task cannot remove a newer task’s bookkeeping entry.
• Delayed removal callbacks are checked against the current zigpy-device identity, preventing an old callback from removing or cancelling a replacement device with the same IEEE.
• Shutdown invalidates every active lifecycle, cancels and awaits initialization tasks, waits for removals, and crosses each lifecycle lock before tearing down wrappers and the controller.
• Lifecycle callbacks received after shutdown begins are ignored.
• Lifecycle tracking is reset when a fully shut down Gateway is initialized again, binding identity checks to the new controller’s device objects.

Initialization tasks are intentionally not eager-started: the gateway must store the task handle before lifecycle code can re-enter through removal or another join.

Real-world failure scenarios

Device removed during pairing

A sleepy device may remain inside async_configure() while waiting for a response. If the user removes it, cancellation can be observed but delayed while the coroutine performs cleanup.

Previously, removal could delete the wrapper and publish DeviceRemovedEvent while that initialization task was still alive. When the task resumed, it could continue initialization or emit DeviceFullInitEvent for a device already reported as removed.

Now removal invalidates the generation, cancels the exact task, and waits for it to finish. Generation checks prevent that task from committing even if it handles cancellation before exiting.

Removal during a factory-changing reinterview

A reinterview can discover that a different quirk factory now applies. That path tears down the old wrapper and later creates and installs a replacement.

If removal occurs while teardown is awaiting, the stale reinterview must not resume and call Device.new(), reinstall the device, or publish full initialization.

The reinterview now checks both lifecycle generation and exact wrapper identity after teardown and before replacement, configuration, initialization, group-subscription refresh, fallback notification, and event emission.

Immediate re-pair with the same IEEE

A device may be removed and then rejoin quickly using the same IEEE while the old lifecycle is still cleaning up.

Without serialization, the old removal could delete the new wrapper, or an old task’s completion callback could clear the new initialization task’s bookkeeping.

The new lifecycle waits for the captured removal task before entering the per-IEEE lock. Removal only deletes the exact wrapper it captured, and completion callbacks only clear entries that still point to their own task.

A delayed removal callback for the old zigpy object is also rejected without changing the replacement’s generation or cancelling its in-flight initialization.

Reload or shutdown during initialization

Reload can begin while a device is still configuring. Merely cancelling that task and proceeding with wrapper and controller teardown leaves the task able to resume against partially dismantled state.

Shutdown now invalidates lifecycle generations, cancels and awaits initialization, waits for removal work, and crosses the per-device locks before wrapper or controller teardown. No stale full-initialization event can be emitted after shutdown begins.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 97.28%. Comparing base (82a14cd) to head (2a16b7f).

Additional details and impacted files

@@           Coverage Diff            @@
##              dev     #690    +/-   ##
========================================
  Coverage   97.27%   97.28%            
========================================
  Files          55       55            
  Lines       10941    11053   +112     
========================================
+ Hits        10643    10753   +110     
- Misses        298      300     +2     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull request overview

Improves ZHA gateway lifecycle robustness by adding per-device lifecycle generations/locks and by quiescing device lifecycle work during shutdown/removal, reducing race conditions between initialization, reinterview, removal, and shutdown.

Changes:

• Introduces per-device lifecycle generation tracking + locks to serialize initialization/reinterview/removal and prevent stale callbacks from committing.
• Refactors device init/reinterview/removal flows to cancel superseded work, await cancellation-resistant tasks appropriately, and harden shutdown teardown ordering.
• Expands gateway tests to cover lifecycle races (init vs remove, remove vs immediate rejoin, shutdown vs in-flight init).

Reviewed changes

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

File	Description
zha/application/gateway.py	Adds device lifecycle generation/locking + quiesce logic; refactors init/reinterview/removal/shutdown to avoid lifecycle races.
tests/test_gateway.py	Adds/updates tests validating the new lifecycle serialization and shutdown/removal race handling.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[zigpy-review-bot]

Reviewed the lifecycle machinery end-to-end and it holds up well — the generation counter (advanced before cancellation) + per-IEEE lock + removal-task barrier + post-await revalidation is a sound and internally consistent design, and the wait-before-lock ordering keeps the task wait-graph acyclic, so no lock-cycle deadlock is reachable. Test coverage maps cleanly onto each scenario in the description; locally all test_gateway.py tests pass with no wall-clock outliers under looptime, in-venv mypy is clean, and CI is green across 3.12–3.14.

I traced the paths this PR is most exposed on and didn't find a remaining race: an old removal can't clobber a newer lifecycle (the wrapper-identity check in _async_remove_device gates the del self._devices[ieee]), the done-callbacks only delete entries that still point at their own task, and shutdown's quiesce awaits (rather than races) removals before crossing the locks. An independent second-opinion pass reached the same conclusion.

No blockers — just a few optional observations:

• _device_lifecycle_zigpy_devices keeps a strong reference to each removed device's zigpy Device until the gateway is reinitialized (gateway.py:834 is read-only; the dict is only cleared in _async_initialize). It's bounded by distinct IEEEs so it's not a runaway leak, and I understand keeping the generations/locks for the run is deliberate for stale-callback rejection — but this dict retains the heaviest objects (endpoints/clusters/listeners). A weakref, or pruning just this dict on a successful removal, would trim the retention without weakening the ownership checks. Optional.
• Optional readability: _async_remove_device awaits the previous task with a plain asyncio.gather(..., return_exceptions=True) (gateway.py:889) while _async_run_device_lifecycle_operation shields the equivalent wait (:505). The asymmetry is safe (removal tasks are only cancelled at the final super().shutdown(), after quiesce has already awaited them), but a one-line comment on why removal doesn't need the shield would save the next reader the trace.
• Nit: in _async_run_device_lifecycle_operation, await asyncio.wait({previous_removal_task}) would be a simpler equivalent to the shield(gather(...)) wait — same don't-propagate / don't-cancel-the-inner semantics.

Nice work — the invariant is clearly stated and the tests back it up.

[dmulcahey]

Follow-up on the optional observations, addressed in 2a16b7f:

• Strong device retention: valid and corrected. The lifecycle identity map now uses weak values. A pending lifecycle task or active ZHA wrapper still owns the current zigpy device strongly, so same-IEEE stale-callback rejection remains intact, while the identity map no longer keeps removed endpoint, cluster, and listener graphs alive by itself. A focused regression test verifies that an otherwise unowned tracked zigpy device is collectable and that its weak map entry disappears.
• Removal wait asymmetry: the behavior is intentionally unchanged and now documented at the wait site. Initialization and reinterview tasks can be superseded and cancelled, so their wait on an older removal must be shielded. Removal tasks are chained rather than superseded by lifecycle callbacks, so their previous-removal wait does not need that shield.
• asyncio.wait simplification: not applied because it is not equivalent here. asyncio.wait does not retrieve the completed task outcome. The current shield plus gather with return_exceptions consumes a failed removal outcome without letting it poison the replacement lifecycle and without producing an unhandled task exception.