Keep track of TimerHandles independent from event loop

homeassistant/core.py

@@ -42,6 +42,7 @@
 )
 from urllib.parse import urlparse
 
+from _weakrefset import WeakSet
 import voluptuous as vol
 import yarl
 
@@ -376,6 +377,7 @@ def __init__(self, config_dir: str) -> None:
         self.loop = asyncio.get_running_loop()
         self._tasks: set[asyncio.Future[Any]] = set()
         self._background_tasks: set[asyncio.Future[Any]] = set()
+        self._timers: WeakSet[asyncio.TimerHandle] = WeakSet()
         self.bus = EventBus(self)
         self.services = ServiceRegistry(self)
         self.states = StateMachine(self.bus, self.loop)
@@ -985,22 +987,19 @@ async def async_stop(self, exit_code: int = 0, *, force: bool = False) -> None:
 
     def _cancel_cancellable_timers(self) -> None:
         """Cancel timer handles marked as cancellable."""
-        # pylint: disable-next=protected-access
-        handles: Iterable[asyncio.TimerHandle] = self.loop._scheduled  # type: ignore[attr-defined]
-        for handle in handles:
-            if (
-                not handle.cancelled()
-                and (args := handle._args)  # pylint: disable=protected-access
-                and type(job := args[0]) is HassJob  # noqa: E721
-                and job.cancel_on_shutdown
-            ):
+        for handle in self._timers:
+            if not handle.cancelled():
                 handle.cancel()
 
     def _async_log_running_tasks(self, stage: str) -> None:
         """Log all running tasks."""
         for task in self._tasks:
             _LOGGER.warning("Shutdown stage '%s': still running: %s", stage, task)
 
+    def async_track_timer_handle(self, timer_handle: asyncio.TimerHandle) -> None:
+        """Add a timer handle to be canceled on shutdown."""
+        self._timers.add(timer_handle)
+
 
 class Context:
     """The context that triggered something."""

homeassistant/helpers/event.py

@@ -1480,7 +1480,10 @@ def async_call_at(
         if isinstance(action, HassJob)
         else HassJob(action, f"call_at {loop_time}")
     )
-    return hass.loop.call_at(loop_time, _run_async_call_action, hass, job).cancel
+    handle = hass.loop.call_at(loop_time, _run_async_call_action, hass, job)
+    if job.cancel_on_shutdown:
+        hass.async_track_timer_handle(handle)
+    return handle.cancel
 
 
 @callback
@@ -1503,7 +1506,10 @@ def async_call_later(
         else HassJob(action, f"call_later {delay}")
     )
     loop = hass.loop
-    return loop.call_at(loop.time() + delay, _run_async_call_action, hass, job).cancel
+    handle = loop.call_at(loop.time() + delay, _run_async_call_action, hass, job)
+    if job.cancel_on_shutdown:
+        hass.async_track_timer_handle(handle)
+    return handle.cancel
 
 
 call_later = threaded_listener_factory(async_call_later)