mdns: don't stop a unicast scan before results stabilize

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,13 +204,18 @@ 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
         self.parser: ServiceParser = ServiceParser()
         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()

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