Trigger Home Assistant shutdown automations right before the stop event instead of during it (home-assistant#91165)

Co-authored-by: Erik <[email protected]>

Author: tetele

Dockerfile

@@ -6,7 +6,7 @@ FROM ${BUILD_FROM}
 
 # Synchronize with homeassistant/core.py:async_stop
 ENV \
-    S6_SERVICES_GRACETIME=220000
+    S6_SERVICES_GRACETIME=240000
 
 ARG QEMU_CPU
 

homeassistant/components/homeassistant/triggers/homeassistant.py

@@ -1,8 +1,8 @@
 """Offer Home Assistant core automation rules."""
 import voluptuous as vol
 
-from homeassistant.const import CONF_EVENT, CONF_PLATFORM, EVENT_HOMEASSISTANT_STOP
-from homeassistant.core import CALLBACK_TYPE, HassJob, HomeAssistant, callback
+from homeassistant.const import CONF_EVENT, CONF_PLATFORM
+from homeassistant.core import CALLBACK_TYPE, HassJob, HomeAssistant
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.trigger import TriggerActionType, TriggerInfo
 from homeassistant.helpers.typing import ConfigType
@@ -30,24 +30,17 @@ async def async_attach_trigger(
     job = HassJob(action, f"homeassistant trigger {trigger_info}")
 
     if event == EVENT_SHUTDOWN:
-
-        @callback
-        def hass_shutdown(event):
-            """Execute when Home Assistant is shutting down."""
-            hass.async_run_hass_job(
-                job,
-                {
-                    "trigger": {
-                        **trigger_data,
-                        "platform": "homeassistant",
-                        "event": event,
-                        "description": "Home Assistant stopping",
-                    }
-                },
-                event.context,
-            )
-
-        return hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, hass_shutdown)
+        return hass.async_add_shutdown_job(
+            job,
+            {
+                "trigger": {
+                    **trigger_data,
+                    "platform": "homeassistant",
+                    "event": event,
+                    "description": "Home Assistant stopping",
+                }
+            },
+        )
 
     # Automation are enabled while hass is starting up, fire right away
     # Check state because a config reload shouldn't trigger it.

homeassistant/core.py

@@ -18,6 +18,7 @@
 )
 import concurrent.futures
 from contextlib import suppress
+from dataclasses import dataclass
 import datetime
 import enum
 import functools
@@ -107,9 +108,10 @@
     from .helpers.entity import StateInfo
 
 
-STAGE_1_SHUTDOWN_TIMEOUT = 100
-STAGE_2_SHUTDOWN_TIMEOUT = 60
-STAGE_3_SHUTDOWN_TIMEOUT = 30
+STOPPING_STAGE_SHUTDOWN_TIMEOUT = 20
+STOP_STAGE_SHUTDOWN_TIMEOUT = 100
+FINAL_WRITE_STAGE_SHUTDOWN_TIMEOUT = 60
+CLOSE_STAGE_SHUTDOWN_TIMEOUT = 30
 
 block_async_io.enable()
 
@@ -299,6 +301,14 @@ def __repr__(self) -> str:
         return f"<Job {self.name} {self.job_type} {self.target}>"
 
 
+@dataclass(frozen=True)
+class HassJobWithArgs:
+    """Container for a HassJob and arguments."""
+
+    job: HassJob[..., Coroutine[Any, Any, Any] | Any]
+    args: Iterable[Any]
+
+
 def _get_hassjob_callable_job_type(target: Callable[..., Any]) -> HassJobType:
     """Determine the job type from the callable."""
     # Check for partials to properly determine if coroutine function
@@ -370,6 +380,7 @@ def __init__(self, config_dir: str) -> None:
         # Timeout handler for Core/Helper namespace
         self.timeout: TimeoutManager = TimeoutManager()
         self._stop_future: concurrent.futures.Future[None] | None = None
+        self._shutdown_jobs: list[HassJobWithArgs] = []
 
     @property
     def is_running(self) -> bool:
@@ -766,6 +777,42 @@ async def _await_and_log_pending(
             for task in pending:
                 _LOGGER.debug("Waited %s seconds for task: %s", wait_time, task)
 
+    @overload
+    @callback
+    def async_add_shutdown_job(
+        self, hassjob: HassJob[..., Coroutine[Any, Any, Any]], *args: Any
+    ) -> CALLBACK_TYPE:
+        ...
+
+    @overload
+    @callback
+    def async_add_shutdown_job(
+        self, hassjob: HassJob[..., Coroutine[Any, Any, Any] | Any], *args: Any
+    ) -> CALLBACK_TYPE:
+        ...
+
+    @callback
+    def async_add_shutdown_job(
+        self, hassjob: HassJob[..., Coroutine[Any, Any, Any] | Any], *args: Any
+    ) -> CALLBACK_TYPE:
+        """Add a HassJob which will be executed on shutdown.
+
+        This method must be run in the event loop.
+
+        hassjob: HassJob
+        args: parameters for method to call.
+
+        Returns function to remove the job.
+        """
+        job_with_args = HassJobWithArgs(hassjob, args)
+        self._shutdown_jobs.append(job_with_args)
+
+        @callback
+        def remove_job() -> None:
+            self._shutdown_jobs.remove(job_with_args)
+
+        return remove_job
+
     def stop(self) -> None:
         """Stop Home Assistant and shuts down all threads."""
         if self.state == CoreState.not_running:  # just ignore
@@ -799,6 +846,26 @@ async def async_stop(self, exit_code: int = 0, *, force: bool = False) -> None:
                     "Stopping Home Assistant before startup has completed may fail"
                 )
 
+        # Stage 1 - Run shutdown jobs
+        try:
+            async with self.timeout.async_timeout(STOPPING_STAGE_SHUTDOWN_TIMEOUT):
+                tasks: list[asyncio.Future[Any]] = []
+                for job in self._shutdown_jobs:
+                    task_or_none = self.async_run_hass_job(job.job, *job.args)
+                    if not task_or_none:
+                        continue
+                    tasks.append(task_or_none)
+                if tasks:
+                    asyncio.gather(*tasks, return_exceptions=True)
+        except asyncio.TimeoutError:
+            _LOGGER.warning(
+                "Timed out waiting for shutdown jobs to complete, the shutdown will"
+                " continue"
+            )
+            self._async_log_running_tasks("run shutdown jobs")
+
+        # Stage 2 - Stop integrations
+
         # Keep holding the reference to the tasks but do not allow them
         # to block shutdown. Only tasks created after this point will
         # be waited for.
@@ -816,33 +883,32 @@ async def async_stop(self, exit_code: int = 0, *, force: bool = False) -> None:
 
         self.exit_code = exit_code
 
-        # stage 1
         self.state = CoreState.stopping
         self.bus.async_fire(EVENT_HOMEASSISTANT_STOP)
         try:
-            async with self.timeout.async_timeout(STAGE_1_SHUTDOWN_TIMEOUT):
+            async with self.timeout.async_timeout(STOP_STAGE_SHUTDOWN_TIMEOUT):
                 await self.async_block_till_done()
         except asyncio.TimeoutError:
             _LOGGER.warning(
-                "Timed out waiting for shutdown stage 1 to complete, the shutdown will"
+                "Timed out waiting for integrations to stop, the shutdown will"
                 " continue"
             )
-            self._async_log_running_tasks(1)
+            self._async_log_running_tasks("stop integrations")
 
-        # stage 2
+        # Stage 3 - Final write
         self.state = CoreState.final_write
         self.bus.async_fire(EVENT_HOMEASSISTANT_FINAL_WRITE)
         try:
-            async with self.timeout.async_timeout(STAGE_2_SHUTDOWN_TIMEOUT):
+            async with self.timeout.async_timeout(FINAL_WRITE_STAGE_SHUTDOWN_TIMEOUT):
                 await self.async_block_till_done()
         except asyncio.TimeoutError:
             _LOGGER.warning(
-                "Timed out waiting for shutdown stage 2 to complete, the shutdown will"
+                "Timed out waiting for final writes to complete, the shutdown will"
                 " continue"
             )
-            self._async_log_running_tasks(2)
+            self._async_log_running_tasks("final write")
 
-        # stage 3
+        # Stage 4 - Close
         self.state = CoreState.not_running
         self.bus.async_fire(EVENT_HOMEASSISTANT_CLOSE)
 
@@ -856,12 +922,12 @@ async def async_stop(self, exit_code: int = 0, *, force: bool = False) -> None:
                 # were awaiting another task
                 continue
             _LOGGER.warning(
-                "Task %s was still running after stage 2 shutdown; "
+                "Task %s was still running after final writes shutdown stage; "
                 "Integrations should cancel non-critical tasks when receiving "
                 "the stop event to prevent delaying shutdown",
                 task,
             )
-            task.cancel("Home Assistant stage 2 shutdown")
+            task.cancel("Home Assistant final writes shutdown stage")
             try:
                 async with asyncio.timeout(0.1):
                     await task
@@ -870,11 +936,11 @@ async def async_stop(self, exit_code: int = 0, *, force: bool = False) -> None:
             except asyncio.TimeoutError:
                 # Task may be shielded from cancellation.
                 _LOGGER.exception(
-                    "Task %s could not be canceled during stage 3 shutdown", task
+                    "Task %s could not be canceled during final shutdown stage", task
                 )
             except Exception as exc:  # pylint: disable=broad-except
                 _LOGGER.exception(
-                    "Task %s error during stage 3 shutdown: %s", task, exc
+                    "Task %s error during final shutdown stage: %s", task, exc
                 )
 
         # Prevent run_callback_threadsafe from scheduling any additional
@@ -885,14 +951,14 @@ async def async_stop(self, exit_code: int = 0, *, force: bool = False) -> None:
         shutdown_run_callback_threadsafe(self.loop)
 
         try:
-            async with self.timeout.async_timeout(STAGE_3_SHUTDOWN_TIMEOUT):
+            async with self.timeout.async_timeout(CLOSE_STAGE_SHUTDOWN_TIMEOUT):
                 await self.async_block_till_done()
         except asyncio.TimeoutError:
             _LOGGER.warning(
-                "Timed out waiting for shutdown stage 3 to complete, the shutdown will"
+                "Timed out waiting for close event to be processed, the shutdown will"
                 " continue"
             )
-            self._async_log_running_tasks(3)
+            self._async_log_running_tasks("close")
 
         self.state = CoreState.stopped
 
@@ -912,10 +978,10 @@ def _cancel_cancellable_timers(self) -> None:
             ):
                 handle.cancel()
 
-    def _async_log_running_tasks(self, stage: int) -> None:
+    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)
+            _LOGGER.warning("Shutdown stage '%s': still running: %s", stage, task)
 
 
 class Context:

script/hassfest/docker.py

@@ -59,9 +59,10 @@
 
 def _generate_dockerfile() -> str:
     timeout = (
-        core.STAGE_1_SHUTDOWN_TIMEOUT
-        + core.STAGE_2_SHUTDOWN_TIMEOUT
-        + core.STAGE_3_SHUTDOWN_TIMEOUT
+        core.STOPPING_STAGE_SHUTDOWN_TIMEOUT
+        + core.STOP_STAGE_SHUTDOWN_TIMEOUT
+        + core.FINAL_WRITE_STAGE_SHUTDOWN_TIMEOUT
+        + core.CLOSE_STAGE_SHUTDOWN_TIMEOUT
         + executor.EXECUTOR_SHUTDOWN_TIMEOUT
         + thread.THREADING_SHUTDOWN_TIMEOUT
         + 10

tests/test_core.py

@@ -36,6 +36,7 @@
 )
 import homeassistant.core as ha
 from homeassistant.core import (
+    CoreState,
     HassJob,
     HomeAssistant,
     ServiceCall,
@@ -399,6 +400,32 @@ async def test_stage_shutdown(hass: HomeAssistant) -> None:
     assert len(test_all) == 2
 
 
+async def test_stage_shutdown_timeouts(hass: HomeAssistant) -> None:
+    """Simulate a shutdown, test timeouts at each step."""
+
+    with patch.object(hass.timeout, "async_timeout", side_effect=asyncio.TimeoutError):
+        await hass.async_stop()
+
+    assert hass.state == CoreState.stopped
+
+
+async def test_stage_shutdown_generic_error(hass: HomeAssistant, caplog) -> None:
+    """Simulate a shutdown, test that a generic error at the final stage doesn't prevent it."""
+
+    task = asyncio.Future()
+    hass._tasks.add(task)
+
+    def fail_the_task(_):
+        task.set_exception(Exception("test_exception"))
+
+    with patch.object(task, "cancel", side_effect=fail_the_task) as patched_call:
+        await hass.async_stop()
+        assert patched_call.called
+
+    assert "test_exception" in caplog.text
+    assert hass.state == ha.CoreState.stopped
+
+
 async def test_stage_shutdown_with_exit_code(hass: HomeAssistant) -> None:
     """Simulate a shutdown, test calling stuff with exit code checks."""
     test_stop = async_capture_events(hass, EVENT_HOMEASSISTANT_STOP)
@@ -2566,3 +2593,30 @@ def not_callback_func():
         HassJob(not_callback_func, job_type=ha.HassJobType.Callback).job_type
         == ha.HassJobType.Callback
     )
+
+
+async def test_shutdown_job(hass: HomeAssistant) -> None:
+    """Test async_add_shutdown_job."""
+    evt = asyncio.Event()
+
+    async def shutdown_func() -> None:
+        evt.set()
+
+    job = HassJob(shutdown_func, "shutdown_job")
+    hass.async_add_shutdown_job(job)
+    await hass.async_stop()
+    assert evt.is_set()
+
+
+async def test_cancel_shutdown_job(hass: HomeAssistant) -> None:
+    """Test cancelling a job added to async_add_shutdown_job."""
+    evt = asyncio.Event()
+
+    async def shutdown_func() -> None:
+        evt.set()
+
+    job = HassJob(shutdown_func, "shutdown_job")
+    cancel = hass.async_add_shutdown_job(job)
+    cancel()
+    await hass.async_stop()
+    assert not evt.is_set()

tests/test_runner.py

@@ -13,17 +13,18 @@
 from homeassistant.util import executor, thread
 
 # https://github.com/home-assistant/supervisor/blob/main/supervisor/docker/homeassistant.py
-SUPERVISOR_HARD_TIMEOUT = 220
+SUPERVISOR_HARD_TIMEOUT = 240
 
 TIMEOUT_SAFETY_MARGIN = 10
 
 
 async def test_cumulative_shutdown_timeout_less_than_supervisor() -> None:
     """Verify the cumulative shutdown timeout is at least 10s less than the supervisor."""
     assert (
-        core.STAGE_1_SHUTDOWN_TIMEOUT
-        + core.STAGE_2_SHUTDOWN_TIMEOUT
-        + core.STAGE_3_SHUTDOWN_TIMEOUT
+        core.STOPPING_STAGE_SHUTDOWN_TIMEOUT
+        + core.STOP_STAGE_SHUTDOWN_TIMEOUT
+        + core.FINAL_WRITE_STAGE_SHUTDOWN_TIMEOUT
+        + core.CLOSE_STAGE_SHUTDOWN_TIMEOUT
         + executor.EXECUTOR_SHUTDOWN_TIMEOUT
         + thread.THREADING_SHUTDOWN_TIMEOUT
         + TIMEOUT_SAFETY_MARGIN