fix(core): release sockets when close runs before engine setup completes (#1706)

bdraco · web-flow · commit 0deb56b78fe6 · 2026-05-17T15:18:26.000-07:00
diff --git a/src/zeroconf/_engine.py b/src/zeroconf/_engine.py
@@ -114,10 +114,17 @@ async def _async_create_endpoints(self) -> None:
                 lambda: AsyncListener(self.zc),  # type: ignore[arg-type, return-value]
                 sock=s,
             )
+            # Register the wrapped transport before releasing the engine's
+            # handle so a concurrent shutdown always sees ``s`` in exactly
+            # one place; do not add an ``await`` between these two steps.
             self.protocols.append(cast(AsyncListener, protocol))
             self.readers.append(make_wrapped_transport(cast(asyncio.DatagramTransport, transport)))
             if s in sender_sockets:
                 self.senders.append(make_wrapped_transport(cast(asyncio.DatagramTransport, transport)))
+            if s is self._listen_socket:
+                self._listen_socket = None
+            if s in self._respond_sockets:
+                self._respond_sockets.remove(s)
 
     def _async_cache_cleanup(self) -> None:
         """Periodic cache cleanup."""
@@ -139,19 +146,37 @@ def _async_schedule_next_cache_cleanup(self) -> None:
     async def _async_close(self) -> None:
         """Cancel and wait for the cleanup task to finish."""
         assert self._setup_task is not None
-        await self._setup_task
+        # Swallow CancelledError only if the setup task itself was
+        # cancelled (close-before-start); outer-task cancellation must
+        # propagate.
+        try:
+            await self._setup_task
+        except asyncio.CancelledError:
+            if not self._setup_task.cancelled():
+                raise
         self._async_shutdown()
         await asyncio.sleep(0)  # flush out any call soons
-        assert self._cleanup_timer is not None
-        self._cleanup_timer.cancel()
+        if self._cleanup_timer is not None:
+            self._cleanup_timer.cancel()
 
     def _async_shutdown(self) -> None:
-        """Shutdown transports and sockets."""
+        """Shutdown transports and sockets; safe to call repeatedly."""
         assert self.running_future is not None
         assert self.loop is not None
         self.running_future = self.loop.create_future()
+        # Cancel pending setup so it can't wrap fresh transports after
+        # shutdown has started.
+        if self._setup_task is not None and not self._setup_task.done():
+            self._setup_task.cancel()
         for wrapped_transport in itertools.chain(self.senders, self.readers):
             wrapped_transport.transport.close()
+        # Anything still here was never adopted by a transport.
+        if self._listen_socket is not None:
+            self._listen_socket.close()
+            self._listen_socket = None
+        for s in self._respond_sockets:
+            s.close()
+        self._respond_sockets = []
 
     def close(self) -> None:
         """Close from sync context.
diff --git a/tests/services/test_browser.py b/tests/services/test_browser.py
@@ -1033,7 +1033,7 @@ def update_service(self, zc, type_, name) -> None:  # type: ignore[no-untyped-de
 
 
 def test_service_browser_listeners_no_update_service():
-    """Test that the ServiceBrowser ServiceListener that does not implement update_service."""
+    """A listener that ignores update events records only add/remove callbacks."""
 
     # instantiate a zeroconf instance
     zc = Zeroconf(interfaces=["127.0.0.1"])
@@ -1051,6 +1051,9 @@ def remove_service(self, zc, type_, name) -> None:  # type: ignore[no-untyped-de
             if name == registration_name:
                 callbacks.append(("remove", type_, name))
 
+        def update_service(self, zc, type_, name) -> None:  # type: ignore[no-untyped-def]
+            pass
+
     listener = MyServiceListener()
 
     browser = r.ServiceBrowser(zc, type_, None, listener)
diff --git a/tests/services/test_info.py b/tests/services/test_info.py
@@ -1529,6 +1529,7 @@ async def test_bad_ip_addresses_ignored_in_cache():
     info = ServiceInfo(type_, registration_name)
     info.load_from_cache(aiozc.zeroconf)
     assert info.addresses_by_version(IPVersion.V4Only) == [b"\x7f\x00\x00\x01"]
+    await aiozc.async_close()
 
 
 @pytest.mark.asyncio
@@ -1804,6 +1805,7 @@ async def test_address_resolver():
     aiozc.zeroconf.async_send(outgoing)
     assert await resolve_task
     assert resolver.addresses == [b"\x7f\x00\x00\x01"]
+    await aiozc.async_close()
 
 
 @pytest.mark.asyncio
@@ -1828,6 +1830,7 @@ async def test_address_resolver_ipv4():
     aiozc.zeroconf.async_send(outgoing)
     assert await resolve_task
     assert resolver.addresses == [b"\x7f\x00\x00\x01"]
+    await aiozc.async_close()
 
 
 @pytest.mark.asyncio
@@ -1854,6 +1857,7 @@ async def test_address_resolver_ipv6():
     aiozc.zeroconf.async_send(outgoing)
     assert await resolve_task
     assert resolver.ip_addresses_by_version(IPVersion.All) == [ip_address("fe80::52e:c2f2:bc5f:e9c6")]
+    await aiozc.async_close()
 
 
 @pytest.mark.asyncio
diff --git a/tests/test_asyncio.py b/tests/test_asyncio.py
@@ -500,6 +500,7 @@ async def test_async_service_registration_name_strict_check(quick_timing: None)
 
     await aiozc.async_unregister_service(info)
     await aiozc.async_close()
+    zc.close()
 
 
 @pytest.mark.asyncio
diff --git a/tests/test_core.py b/tests/test_core.py
@@ -824,6 +824,9 @@ def _background_register():
 async def test_event_loop_blocked(mock_start):
     """Test we raise NotRunningException when waiting for startup that times out."""
     aiozc = AsyncZeroconf(interfaces=["127.0.0.1"])
-    with pytest.raises(NotRunningException):
-        await aiozc.zeroconf.async_wait_for_start(timeout=0)
-    assert aiozc.zeroconf.started is False
+    try:
+        with pytest.raises(NotRunningException):
+            await aiozc.zeroconf.async_wait_for_start(timeout=0)
+        assert aiozc.zeroconf.started is False
+    finally:
+        await aiozc.async_close()
diff --git a/tests/test_engine.py b/tests/test_engine.py
@@ -73,6 +73,42 @@ async def test_reaper():
         assert record_with_1s_ttl not in entries
 
 
+@pytest.mark.asyncio
+async def test_setup_releases_socket_ownership() -> None:
+    """Engine releases its pending-socket refs once each socket has a transport."""
+    aiozc = AsyncZeroconf(interfaces=["127.0.0.1"])
+    try:
+        await aiozc.zeroconf.async_wait_for_start()
+        engine = aiozc.zeroconf.engine
+        assert engine._listen_socket is None
+        assert engine._respond_sockets == []
+        assert engine.readers
+        assert engine.senders
+    finally:
+        await aiozc.async_close()
+
+
+@pytest.mark.asyncio
+async def test_async_close_propagates_outer_cancellation() -> None:
+    """Outer-task cancellation while awaiting setup propagates to the caller."""
+    aiozc = AsyncZeroconf(interfaces=["127.0.0.1"])
+    try:
+        await aiozc.zeroconf.async_wait_for_start()
+        engine = aiozc.zeroconf.engine
+        loop = asyncio.get_running_loop()
+        original_task = engine._setup_task
+        fake_task = loop.create_future()
+        fake_task.set_exception(asyncio.CancelledError())
+        engine._setup_task = fake_task  # type: ignore[assignment]
+        try:
+            with pytest.raises(asyncio.CancelledError):
+                await engine._async_close()
+        finally:
+            engine._setup_task = original_task
+    finally:
+        await aiozc.async_close()
+
+
 @pytest.mark.asyncio
 async def test_reaper_aborts_when_done():
     """Ensure cache cleanup stops when zeroconf is done."""
diff --git a/tests/test_handlers.py b/tests/test_handlers.py
@@ -1799,6 +1799,8 @@ async def test_response_aggregation_timings_multiple(run_isolated, disable_dupli
         zc.record_manager.async_updates_from_response(incoming)
         assert info2.dns_pointer() in incoming.answers()
 
+    await aiozc.async_close()
+
 
 @pytest.mark.asyncio
 async def test_response_aggregation_random_delay():
diff --git a/tests/utils/test_asyncio.py b/tests/utils/test_asyncio.py
@@ -105,6 +105,7 @@ def _run_coro() -> None:
 
     assert loop.is_running() is False
     runcoro_thread.join()
+    loop.close()
 
 
 def test_cumulative_timeouts_less_than_close_plus_buffer():
