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