From 63c3d6e113ba1ee28b2a265387406ba2289a0f90 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 29 Feb 2024 04:47:36 -1000
Subject: [PATCH] Fix race in config entry setup again (#111800)

Because the setup again was scheduled as a task, it would
not unset self._async_cancel_retry_setup in time and we would
try to unsub self._async_cancel_retry_setup after it had already
fired. Change it to call a callback that runs right away so it
unsets self._async_cancel_retry_setup as soon as its called
so there is no race

fixes #111796
---
 homeassistant/config_entries.py | 37 ++++++++++++++++++---------------
 tests/test_config_entries.py    |  3 ++-
 2 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/homeassistant/config_entries.py b/homeassistant/config_entries.py
index 864ad90344a..1ca40886da2 100644
--- a/homeassistant/config_entries.py
+++ b/homeassistant/config_entries.py
@@ -33,6 +33,7 @@ from .core import (
     CoreState,
     Event,
     HassJob,
+    HassJobType,
     HomeAssistant,
     callback,
 )
@@ -363,7 +364,6 @@ class ConfigEntry:
 
         self._integration_for_domain: loader.Integration | None = None
         self._tries = 0
-        self._setup_again_job: HassJob | None = None
 
     def __repr__(self) -> str:
         """Representation of ConfigEntry."""
@@ -555,12 +555,18 @@ class ConfigEntry:
 
             if hass.state is CoreState.running:
                 self._async_cancel_retry_setup = async_call_later(
-                    hass, wait_time, self._async_get_setup_again_job(hass)
+                    hass,
+                    wait_time,
+                    HassJob(
+                        functools.partial(self._async_setup_again, hass),
+                        job_type=HassJobType.Callback,
+                    ),
                 )
             else:
-                self._async_cancel_retry_setup = hass.bus.async_listen_once(
+                self._async_cancel_retry_setup = hass.bus.async_listen(
                     EVENT_HOMEASSISTANT_STARTED,
                     functools.partial(self._async_setup_again, hass),
+                    run_immediately=True,
                 )
 
             await self._async_process_on_unload(hass)
@@ -585,28 +591,25 @@ class ConfigEntry:
         if not domain_is_integration:
             return
 
+        self.async_cancel_retry_setup()
+
         if result:
             self._async_set_state(hass, ConfigEntryState.LOADED, None)
         else:
             self._async_set_state(hass, ConfigEntryState.SETUP_ERROR, error_reason)
 
-    async def _async_setup_again(self, hass: HomeAssistant, *_: Any) -> None:
-        """Run setup again."""
+    @callback
+    def _async_setup_again(self, hass: HomeAssistant, *_: Any) -> None:
+        """Schedule setup again.
+
+        This method is a callback to ensure that _async_cancel_retry_setup
+        is unset as soon as its callback is called.
+        """
+        self._async_cancel_retry_setup = None
         # Check again when we fire in case shutdown
         # has started so we do not block shutdown
         if not hass.is_stopping:
-            self._async_cancel_retry_setup = None
-            await self.async_setup(hass)
-
-    @callback
-    def _async_get_setup_again_job(self, hass: HomeAssistant) -> HassJob:
-        """Get a job that will call setup again."""
-        if not self._setup_again_job:
-            self._setup_again_job = HassJob(
-                functools.partial(self._async_setup_again, hass),
-                cancel_on_shutdown=True,
-            )
-        return self._setup_again_job
+            hass.async_create_task(self.async_setup(hass), eager_start=True)
 
     @callback
     def async_shutdown(self) -> None:
diff --git a/tests/test_config_entries.py b/tests/test_config_entries.py
index ab1942b5c43..672dbb9ae64 100644
--- a/tests/test_config_entries.py
+++ b/tests/test_config_entries.py
@@ -1174,7 +1174,8 @@ async def test_setup_raise_not_ready(
     mock_setup_entry.side_effect = None
     mock_setup_entry.return_value = True
 
-    await hass.async_run_hass_job(p_setup, None)
+    hass.async_run_hass_job(p_setup, None)
+    await hass.async_block_till_done()
     assert entry.state is config_entries.ConfigEntryState.LOADED
     assert entry.reason is None