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

@@ -1463,6 +1463,19 @@ def _run_async_call_action(
     hass.async_run_hass_job(job, time_tracker_utcnow())
 
 
+@callback
+def _async_schedule_and_track_handle(
+    hass: HomeAssistant,
+    when: float,
+    job: HassJob[[datetime], Coroutine[Any, Any, None] | None],
+) -> CALLBACK_TYPE:
+    """Schedule the callback and keep track of the TimeHandle."""
+    handle = hass.loop.call_at(when, _run_async_call_action, hass, job)
+    if job.cancel_on_shutdown:
+        hass.async_track_timer_handle(handle)
+    return handle.cancel
+
+
 @callback
 @bind_hass
 def async_call_at(
@@ -1480,7 +1493,7 @@ 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
+    return _async_schedule_and_track_handle(hass, loop_time, job)
 
 
 @callback
@@ -1502,8 +1515,7 @@ def async_call_later(
         if isinstance(action, HassJob)
         else HassJob(action, f"call_later {delay}")
     )
-    loop = hass.loop
-    return loop.call_at(loop.time() + delay, _run_async_call_action, hass, job).cancel
+    return _async_schedule_and_track_handle(hass, hass.loop.time() + delay, job)
 
 
 call_later = threaded_listener_factory(async_call_later)

tests/test_core.py

@@ -2673,6 +2673,8 @@ def run_job(job: HassJob) -> None:
     )
     timer2 = hass.loop.call_later(60, run_job, HassJob(ha.callback(job)))
 
+    hass.async_track_timer_handle(timer1)
+
     await hass.async_stop()
 
     assert timer1.cancelled()