From a98334ede80e6594581c9cdbaf6244a077ca644b Mon Sep 17 00:00:00 2001
From: Mike Degatano <michael.degatano@gmail.com>
Date: Fri, 4 Aug 2023 16:28:44 -0400
Subject: [PATCH] Cancel startup wait task on addon uninstallation (#4475)

* Cancel startup wait task on addon uninstallation

* Await startup task instead

* Suppress cancelled error
---
 supervisor/addons/addon.py   | 14 ++++++++++++-
 tests/addons/test_addon.py   |  1 +
 tests/addons/test_manager.py | 39 ++++++++++++++++++++++++++++++++++++
 tests/api/test_addons.py     |  1 +
 tests/api/test_store.py      |  1 +
 5 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/supervisor/addons/addon.py b/supervisor/addons/addon.py
index 57cbcdaa9..a3c79797b 100644
--- a/supervisor/addons/addon.py
+++ b/supervisor/addons/addon.py
@@ -129,6 +129,7 @@ class Addon(AddonModel):
         )
         self._listeners: list[EventListener] = []
         self._startup_event = asyncio.Event()
+        self._startup_task: asyncio.Task | None = None
 
         @Job(
             name=f"addon_{slug}_restart_after_problem",
@@ -608,6 +609,12 @@ class Addon(AddonModel):
 
     async def unload(self) -> None:
         """Unload add-on and remove data."""
+        if self._startup_task:
+            # If we were waiting on startup, cancel that and let the task finish before proceeding
+            self._startup_task.cancel(f"Removing add-on {self.name} from system")
+            with suppress(asyncio.CancelledError):
+                await self._startup_task
+
         for listener in self._listeners:
             self.sys_bus.remove_listener(listener)
 
@@ -699,13 +706,18 @@ class Addon(AddonModel):
     async def _wait_for_startup(self) -> None:
         """Wait for startup event to be set with timeout."""
         try:
-            await asyncio.wait_for(self._startup_event.wait(), STARTUP_TIMEOUT)
+            self._startup_task = self.sys_create_task(self._startup_event.wait())
+            await asyncio.wait_for(self._startup_task, STARTUP_TIMEOUT)
         except asyncio.TimeoutError:
             _LOGGER.warning(
                 "Timeout while waiting for addon %s to start, took more then %s seconds",
                 self.name,
                 STARTUP_TIMEOUT,
             )
+        except asyncio.CancelledError as err:
+            _LOGGER.info("Wait for addon startup task cancelled due to: %s", err)
+        finally:
+            self._startup_task = None
 
     async def start(self) -> Awaitable[None]:
         """Set options and start add-on.
diff --git a/tests/addons/test_addon.py b/tests/addons/test_addon.py
index 606aadb42..83e27cab7 100644
--- a/tests/addons/test_addon.py
+++ b/tests/addons/test_addon.py
@@ -523,6 +523,7 @@ async def test_restore(
     path_extern,
 ) -> None:
     """Test restoring an addon."""
+    coresys.hardware.disk.get_disk_free_space = lambda x: 5000
     install_addon_ssh.path_data.mkdir()
     await install_addon_ssh.load()
 
diff --git a/tests/addons/test_manager.py b/tests/addons/test_manager.py
index 628d3ea20..98c348d7c 100644
--- a/tests/addons/test_manager.py
+++ b/tests/addons/test_manager.py
@@ -277,3 +277,42 @@ async def test_rebuild(
         start_task = await coresys.addons.rebuild(TEST_ADDON_SLUG)
 
     assert bool(start_task) is (status == "running")
+
+
+async def test_start_wait_cancel_on_uninstall(
+    coresys: CoreSys,
+    install_addon_ssh: Addon,
+    container: MagicMock,
+    caplog: pytest.LogCaptureFixture,
+    tmp_supervisor_data,
+    path_extern,
+) -> None:
+    """Test the addon wait task is cancelled when addon is uninstalled."""
+    install_addon_ssh.path_data.mkdir()
+    container.attrs["Config"] = {"Healthcheck": "exists"}
+    await install_addon_ssh.load()
+    await asyncio.sleep(0)
+    assert install_addon_ssh.state == AddonState.STOPPED
+
+    start_task = asyncio.create_task(await install_addon_ssh.start())
+    assert start_task
+
+    coresys.bus.fire_event(
+        BusEvent.DOCKER_CONTAINER_STATE_CHANGE,
+        DockerContainerStateEvent(
+            name=f"addon_{TEST_ADDON_SLUG}",
+            state=ContainerState.RUNNING,
+            id="abc123",
+            time=1,
+        ),
+    )
+    await asyncio.sleep(0.01)
+
+    assert not start_task.done()
+    assert install_addon_ssh.state == AddonState.STARTUP
+
+    caplog.clear()
+    await coresys.addons.uninstall(TEST_ADDON_SLUG)
+    await asyncio.sleep(0.01)
+    assert start_task.done()
+    assert "Wait for addon startup task cancelled" in caplog.text
diff --git a/tests/api/test_addons.py b/tests/api/test_addons.py
index f45a0a73b..2875ab38b 100644
--- a/tests/api/test_addons.py
+++ b/tests/api/test_addons.py
@@ -171,6 +171,7 @@ async def test_api_addon_rebuild_healthcheck(
     path_extern,
 ):
     """Test rebuilding an addon waits for healthy."""
+    coresys.hardware.disk.get_disk_free_space = lambda x: 5000
     container.status = "running"
     install_addon_ssh.path_data.mkdir()
     container.attrs["Config"] = {"Healthcheck": "exists"}
diff --git a/tests/api/test_store.py b/tests/api/test_store.py
index b561464f3..29f844496 100644
--- a/tests/api/test_store.py
+++ b/tests/api/test_store.py
@@ -125,6 +125,7 @@ async def test_api_store_update_healthcheck(
     path_extern,
 ):
     """Test updating an addon with healthcheck waits for health status."""
+    coresys.hardware.disk.get_disk_free_space = lambda x: 5000
     container.status = "running"
     container.attrs["Config"] = {"Healthcheck": "exists"}
     install_addon_ssh.path_data.mkdir()