From 185e87cee82af82d59e5fd7eeffd77efb0807d1e Mon Sep 17 00:00:00 2001 From: Cory Forseth Date: Sat, 18 Jul 2026 16:02:41 -0500 Subject: [PATCH] core/scan: don't close unicast mDNS socket before results stabilize UnicastDnsSdClientProtocol previously finished as soon as it had received some reply for every outbound query packet, regardless of whether the reply content was complete. A device can take a moment after boot/wake to re-register a dynamic-port service (e.g. Companion) with its own mDNS responder, so this could close the socket before that service ever had a chance to answer. Completion now requires either every explicitly requested service to have answered (unambiguous, no wait needed), or the discovered service set to have gone unchanged for UNICAST_MIN_STABLE_ROUNDS consecutive rounds (reduces, doesn't eliminate, the risk of a still-registering service being missed), or the final resend round to be reached. --- pyatv/core/mdns.py | 62 +++++++++++++++- tests/core/test_mdns_functional.py | 115 +++++++++++++++++++++++++++++ 2 files changed, 175 insertions(+), 2 deletions(-) diff --git a/pyatv/core/mdns.py b/pyatv/core/mdns.py index f48f6ff85..89b468941 100644 --- a/pyatv/core/mdns.py +++ b/pyatv/core/mdns.py @@ -29,6 +29,22 @@ SLEEP_PROXY_SERVICE = "_sleep-proxy._udp.local" +# Number of consecutive resend rounds with no new (non-empty) services +# discovered before a unicast query is considered "stable" and finished early +# (unless every requested service has already answered, or the query has run +# out of resend rounds - see UnicastDnsSdClientProtocol.datagram_received). +# +# A device can take a moment after boot/wake to re-register a dynamic-port +# service (e.g. Companion) with its own mDNS responder, so a round with no new +# answers isn't reliable proof nothing else is coming - only repeated quiet +# rounds are. This value trades latency against that risk and is a heuristic, +# not a guarantee: raise it for more confidence that a slow-to-register +# service won't be missed (at the cost of waiting longer before concluding a +# device genuinely doesn't have the rest of what was queried); lower it +# (never below 1, which reduces to "one round of confirmation") for faster +# results when that risk matters less than speed. +UNICAST_MIN_STABLE_ROUNDS = 2 + # This module produces a lot of debug output, use a dedicated log level. # Maybe move this to top-level support later? TRAFFIC_LEVEL = logging.DEBUG - 5 @@ -188,6 +204,7 @@ class UnicastDnsSdClientProtocol(asyncio.Protocol): def __init__(self, services: typing.List[str], host: str, timeout: int): """Initialize a new UnicastDnsSdClientProtocol.""" self.queries = create_service_queries(services, QueryType.PTR) + self._requested_services: typing.FrozenSet[str] = frozenset(services) self.host = host self.timeout = timeout self.transport = None @@ -195,6 +212,10 @@ def __init__(self, services: typing.List[str], host: str, timeout: int): self.semaphore: asyncio.Semaphore = asyncio.Semaphore(value=0) self.received_responses: int = 0 self._task: typing.Optional[asyncio.Future] = None + self._total_rounds: int = math.ceil(timeout) + self._round_no: int = 0 + self._last_known_services: typing.FrozenSet[str] = frozenset() + self._last_growth_round: typing.Optional[int] = None async def get_response(self) -> Response: """Get respoonse with a maximum timeout.""" @@ -224,7 +245,9 @@ def connection_made(self, transport) -> None: self._task = asyncio.ensure_future(self._resend_loop()) async def _resend_loop(self): - for _ in range(math.ceil(self.timeout)): + for round_no in range(self._total_rounds): + self._round_no = round_no + for query in self.queries: log_binary( _LOGGER, @@ -250,7 +273,42 @@ def datagram_received(self, data: bytes, _) -> None: self.parser.add_message(DnsMessage().unpack(data)) self.received_responses += 1 - if self.received_responses == len(self.queries): + if self.received_responses < len(self.queries): + return + + current_services = frozenset(self.parser.ptrs) + + # If every service explicitly requested has been answered, there's + # nothing left that could still be pending (mDNS has no "not found" + # response, but it can't hide something we never asked about either) + # - safe to stop right away. + have_everything_requested = self._requested_services <= current_services + + # Otherwise, a device can take a moment after boot/wake to + # re-register a dynamic-port service (e.g. Companion) with its own + # mDNS responder, so a still-missing requested service doesn't + # necessarily mean it doesn't exist. Track the round new services + # were last seen in (rather than only comparing to the immediately + # preceding round, which a round with multiple query packets could + # evaluate more than once, or which could be fooled by one round of + # coincidental quiet before something else answers) and only finish + # early once at least UNICAST_MIN_STABLE_ROUNDS have passed with + # nothing new showing up, or once this is the last round available + # and there's nothing left to gain by waiting further. An empty + # result is never "stable" - it's just as likely evidence the device + # isn't ready yet as it is evidence nothing will ever show up. + if current_services != self._last_known_services: + self._last_growth_round = self._round_no + self._last_known_services = current_services + + is_final_round = self._round_no == self._total_rounds - 1 + is_stable = ( + current_services + and self._last_growth_round is not None + and self._round_no - self._last_growth_round >= UNICAST_MIN_STABLE_ROUNDS + ) + + if have_everything_requested or is_final_round or is_stable: self._finished() if self.transport: self.transport.close() diff --git a/tests/core/test_mdns_functional.py b/tests/core/test_mdns_functional.py index 730bc2daa..bb4f81d41 100644 --- a/tests/core/test_mdns_functional.py +++ b/tests/core/test_mdns_functional.py @@ -152,6 +152,121 @@ async def test_unicast_resend_if_no_response(udns_server): assert resp.services[0].port == service.port +async def test_unicast_catches_service_registered_after_first_round(udns_server): + """A service that only starts answering in a later resend round (e.g. a + device still re-registering a dynamic-port service with its own mDNS + responder after boot/wake) must still be found, not missed because the + first round already got replies for everything else that was asked. + + Registration is gated on request count so the test deterministically + lands the service in the second resend round. The test suite's global + fake asyncio.sleep is disabled here: it collapses the ~1s spacing between + resend rounds, which races the client's round bookkeeping ahead of the + (real, unpatched) UDP round-trip to the fake server and makes the + scenario this test targets unreproducible. + """ + udns_server.services = {} + original_datagram_received = udns_server.datagram_received + + def _datagram_received(data, addr): + if udns_server.request_count >= 1: + udns_server.services[MEDIAREMOTE_SERVICE] = TEST_SERVICES[ + MEDIAREMOTE_SERVICE + ] + original_datagram_received(data, addr) + + udns_server.datagram_received = _datagram_received + + utils.unstub_sleep() + try: + resp = await mdns.unicast( + asyncio.get_running_loop(), + "127.0.0.1", + [MEDIAREMOTE_SERVICE], + port=udns_server.port, + timeout=3, + ) + finally: + utils.stub_sleep() + + assert len(resp.services) == 1 + assert resp.services[0].type == MEDIAREMOTE_SERVICE + assert resp.services[0].name == TEST_SERVICES[MEDIAREMOTE_SERVICE].name + + +async def test_unicast_two_empty_rounds_do_not_count_as_stable(udns_server): + """Two consecutive rounds that both found nothing must not be treated as + "results have stabilized" and stop early - an empty round says nothing + about whether the device is still starting up versus genuinely lacking + the service, so only a stable *non-empty* result is meaningful evidence. + """ + udns_server.services = {} + original_datagram_received = udns_server.datagram_received + + def _datagram_received(data, addr): + if udns_server.request_count >= 2: + udns_server.services[MEDIAREMOTE_SERVICE] = TEST_SERVICES[ + MEDIAREMOTE_SERVICE + ] + original_datagram_received(data, addr) + + udns_server.datagram_received = _datagram_received + + utils.unstub_sleep() + try: + resp = await mdns.unicast( + asyncio.get_running_loop(), + "127.0.0.1", + [MEDIAREMOTE_SERVICE], + port=udns_server.port, + timeout=4, + ) + finally: + utils.stub_sleep() + + assert len(resp.services) == 1 + assert resp.services[0].type == MEDIAREMOTE_SERVICE + assert resp.services[0].name == TEST_SERVICES[MEDIAREMOTE_SERVICE].name + + +async def test_unicast_catches_service_appearing_after_another_stabilizes(udns_server): + """One requested service (MRP) answers immediately and stays unchanged for + a round; a second requested service (AirPlay) only starts answering + later. A single round of "nothing changed" must not be treated as final + stability when a still-requested service is missing - min_stable_rounds + requires that many *consecutive* unchanged rounds, so one quiet round + followed by a new arrival must reset the count and keep waiting. + """ + udns_server.services = {MEDIAREMOTE_SERVICE: TEST_SERVICES[MEDIAREMOTE_SERVICE]} + airplay_type, airplay_service = fake_udns.airplay_service( + "Kitchen", "AA:BB:CC:DD:EE:FF" + ) + + original_datagram_received = udns_server.datagram_received + + def _datagram_received(data, addr): + if udns_server.request_count >= 2: + udns_server.services[airplay_type] = airplay_service + original_datagram_received(data, addr) + + udns_server.datagram_received = _datagram_received + + utils.unstub_sleep() + try: + resp = await mdns.unicast( + asyncio.get_running_loop(), + "127.0.0.1", + [MEDIAREMOTE_SERVICE, airplay_type], + port=udns_server.port, + timeout=6, + ) + finally: + utils.stub_sleep() + + found_types = {service.type for service in resp.services} + assert found_types == {MEDIAREMOTE_SERVICE, airplay_type} + + async def test_unicast_specific_service(udns_server): resp, _ = await unicast(udns_server, [SERVICE_NAME + "." + MEDIAREMOTE_SERVICE]) assert len(resp.services) == 1