fix: cancellation handling

bdraco · bdraco · commit 05577430bcb3 · 2025-01-31T13:27:43.000-06:00
diff --git a/src/zeroconf/_core.py b/src/zeroconf/_core.py
@@ -207,6 +207,7 @@ def started(self) -> bool:
             not self.done
             and self.engine.running_future
             and self.engine.running_future.done()
+            and not self.engine.running_future.cancelled()
             and not self.engine.running_future.exception()
         )
 
@@ -242,7 +243,12 @@ async def async_wait_for_start(self) -> None:
             raise NotRunningException
         assert self.engine.running_future is not None
         await wait_future_or_timeout(self.engine.running_future, timeout=_STARTUP_TIMEOUT)
-        if not self.engine.running_future.done() or self.engine.running_future.exception() or self.done:
+        if (
+            not self.engine.running_future.done()
+            or self.engine.running_future.cancelled()
+            or self.engine.running_future.exception()
+            or self.done
+        ):
             raise NotRunningException
 
     @property
diff --git a/src/zeroconf/_utils/asyncio.py b/src/zeroconf/_utils/asyncio.py
@@ -25,6 +25,7 @@
 import asyncio
 import concurrent.futures
 import contextlib
+import sys
 from typing import Any, Awaitable, Coroutine
 
 from .._exceptions import EventLoopBlocked
@@ -70,6 +71,9 @@ async def wait_future_or_timeout(future: asyncio.Future[bool | None], timeout: f
     handle = loop.call_later(timeout, _set_future_none_if_not_done, future)
     try:
         await future
+    except asyncio.CancelledError:
+        if sys.version_info >= (3, 11) and (task := asyncio.current_task()) and task.cancelling():
+            raise
     finally:
         handle.cancel()
 
diff --git a/src/zeroconf/asyncio.py b/src/zeroconf/asyncio.py
@@ -29,6 +29,7 @@
 
 from ._core import Zeroconf
 from ._dns import DNSQuestionType
+from ._exceptions import NotRunningException
 from ._services import ServiceListener
 from ._services.browser import _ServiceBrowserBase
 from ._services.info import AsyncServiceInfo, ServiceInfo
@@ -227,7 +228,7 @@ async def async_close(self) -> None:
         """Ends the background threads, and prevent this instance from
         servicing further queries."""
         if not self.zeroconf.done:
-            with contextlib.suppress(asyncio.TimeoutError):
+            with contextlib.suppress(asyncio.TimeoutError, NotRunningException):
                 await asyncio.wait_for(self.zeroconf.async_wait_for_start(), timeout=1)
         await self.async_remove_all_service_listeners()
         await self.async_unregister_all_services()
