Cancel an in-flight entity initialization on teardown and rebuild

[zigpy-review-bot]

Note

Stacked on #898 (base branch zigpy-bot/serialize-device-entity-init). Do not merge before #898; GitHub retargets this PR to dev once #898 merges and its branch is deleted. #878 is independent but touches the same file, so a rebase may be needed if it lands first.

Closes the remaining gap noted in #897 and #898.

Problem

async_teardown and async_rebuild_from_zigpy_device do not take the per-device entity lifecycle lock that #898 adds, so a device removal or a re-interview rebuild can land while an initialization round (startup mains polling, an availability refresh) is still registering entities.

The window is the registration step of that round: _add_pending_entities awaits entity.on_remove() for every entity it drops, which suspends while that entity's owned tasks are cancelled (lights, cluster pollers, the device tracker own one). A teardown landing there clears the pending list and runs on_remove() on everything in it, but cannot take back the entities the round has already collected, and those are registered afterwards. On a removed device they are announced to consumers after the device is gone (some with live listeners, some already half torn down but still registered). After a rebuild they are bound to the old zigpy device, and the re-interview's own initialization drops its fresh entities as duplicates, so the stale ones stay registered. Only unregistered entities get collected, so the exposure is mainly a startup poll or availability refresh of a device that has just joined (recompute_entities collecting newly supported entities on a live device is the same race); a teardown landing during the round's device reads is harmless.

Change

• A teardown marks the device torn down before it waits for the lock, and cancels the task currently holding it. It cancels rather than waits because the round may be mid-way through its device reads, and waiting would stall device removal and gateway shutdown (which tears devices down before it cancels background tasks) behind zigpy request timeouts. This mirrors device_reinterviewed cancelling a previous initialization task.
• async_initialize and recompute_entities record the holding task and skip themselves when they get the lock after a teardown was requested. The lock is FIFO, so without the early mark a round already queued ahead of the teardown would run a full round of reads first.
• _add_pending_entities re-checks the mark right before registering, and drops what it collected if a teardown was requested meanwhile. This only matters if the cancellation got swallowed, e.g. by an on_remove() override suppressing CancelledError.
• async_configure and _discover_new_entities skip a torn-down device, so a join racing a removal cannot activate and queue entities that nothing would drain, including a configuration that was already past its first await when the teardown landed.
• A teardown whose wait for the lock is cancelled (a rejoin cancels a queued re-interview task) clears the mark again, since nothing was torn down; otherwise the device would keep its entities but skip every later round until another re-interview. The locked teardown sets the mark again itself, so a rebuild that held the lock while the teardown was requested cannot erase it. The reset is unconditional, so it would also clear the mark of another teardown still in progress on the same object; no production path queues a teardown behind an in-progress one (a re-interview cancels the previous task first, removal pops the device first, shutdown awaits every removal before cancelling), so this is left simple.
• async_initialize returns before setting _initialized when a teardown was requested during its round; that flag only gates whether later rounds emit entity-added events.
• async_rebuild_from_zigpy_device does the same mark + cancel + locked teardown and clears the mark while still holding the lock, so a round queued behind it runs against the rebuilt device and the re-interview's configure + initialize proceed as before.
• BaseEntity.on_remove no longer suppresses CancelledError around the gather of its owned tasks. With return_exceptions=True the children never raise there, so the suppression only ever swallowed the caller's own cancellation, which is exactly the teardown's cancellation of a round that is dropping an entity. Without this change the cancel is eaten at the one suspension point that matters and the teardown degrades to waiting for the round.

Behaviour change to be aware of

Re-initializing a device after on_remove() is no longer supported: the device is gone for good, and only a rebuild re-creates its entities. No production path did that; test_reinitialize_after_on_remove_emits_events did, and is now test_reinitialize_after_rebuild_emits_events using async_rebuild_from_zigpy_device, which keeps its intent (all entities re-added while _initialized stays set still emit events).

Tests

The new tests use a not-yet-initialized IKEA GU10 bulb. The three round-holding tests (cancels-in-flight, skips-queued, rebuild) start an initialization and hold it in the registration step by giving every entity an owned task, so the round suspends in the real on_remove() gather rather than a stand-in; the two swallowed-cancellation tests use an on_remove() override that sleeps and suppresses the cancel; the configure test parks the quirk's custom configuration instead. Then:

• test_teardown_cancels_in_flight_initialize: removes the device. The round is cancelled, nothing is registered, no pending entities remain, and a configure, initialize and recompute started afterwards all skip. On the base it fails with the full entity set registered on the removed device.
• test_teardown_skips_queued_initialize: a second round is queued on the lock before the removal. It skips without a round of reads (one initialize_cluster_configs call in total, not two). On the base it fails with two.
• test_teardown_with_swallowed_cancellation: the round's on_remove() suppresses CancelledError, so it survives the cancel and reaches the registration step; it drops what it collected. On the base it fails with the full entity set registered.
• test_teardown_during_started_configure: a configuration suspended in the quirk's custom configuration when the device is removed resumes without discovering, activating or queueing anything. On the base it fails with the full set discovered and queued.
• test_cancelled_teardown_wait_lets_rounds_run_again: a rebuild queued behind a round whose cancellation got swallowed is itself cancelled; the mark is cleared, the round registers as usual and later rounds work. On the base the mark does not exist.
• test_rebuild_cancels_in_flight_initialize: rebuilds from a new zigpy device object, then runs the re-interview's configure + initialize. Every resulting entity's cluster belongs to the new zigpy device. On the base it fails with the stale set registered right after the rebuild.

Full suite passes, pre-commit clean, venv mypy zha/ clean, the new tests stable over repeated runs.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 97.21%. Comparing base (137812a) to head (09b97d1).

Additional details and impacted files

@@                            Coverage Diff                             @@
##           zigpy-bot/serialize-device-entity-init     #899      +/-   ##
==========================================================================
+ Coverage                                   97.19%   97.21%   +0.01%     
==========================================================================
  Files                                          57       57              
  Lines                                       10569    10618      +49     
==========================================================================
+ Hits                                        10273    10322      +49     
  Misses                                        296      296              

☔ 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.