feat: improve performance of ServiceBrowser outgoing query scheduler (#1170)

bdraco · web-flow · commit 963d022ef82b · 2023-05-03T19:40:04.000-05:00
diff --git a/src/zeroconf/_services/browser.py b/src/zeroconf/_services/browser.py
@@ -460,13 +460,18 @@ def _cancel_send_timer(self) -> None:
         """Cancel the next send."""
         if self._next_send_timer:
             self._next_send_timer.cancel()
+            self._next_send_timer = None
 
     def reschedule_type(self, type_: str, now: float, next_time: float) -> None:
         """Reschedule a type to be refreshed in the future."""
         if self.query_scheduler.reschedule_type(type_, next_time):
+            # We need to send the queries before rescheduling the next one
+            # otherwise we may be scheduling a query to go out in the next
+            # iteration of the event loop which should be sent now.
+            if now >= next_time:
+                self._async_send_ready_queries(now)
             self._cancel_send_timer()
             self._async_schedule_next(now)
-        self._async_send_ready_queries(now)
 
     def _async_send_ready_queries(self, now: float) -> None:
         """Send any ready queries."""
diff --git a/tests/test_asyncio.py b/tests/test_asyncio.py
@@ -996,6 +996,9 @@ def send(out, addr=const._MDNS_ADDR, port=const._MDNS_PORT, v6_flow_scope=()):
                 # Increase simulated time shift by 1/4 of the TTL in seconds
                 time_offset += expected_ttl / 4
                 now = _new_current_time_millis()
+                # Force the next query to be sent since we are testing
+                # to see if the query contains answers and not the scheduler
+                browser.query_scheduler._next_time[type_] = now + (1000 * expected_ttl)
                 browser.reschedule_type(type_, now, now)
                 sleep_count += 1
                 await asyncio.wait_for(got_query.wait(), 1)
@@ -1244,3 +1247,67 @@ def update_service(self, zc, type_, name) -> None:  # type: ignore[no-untyped-de
         ('add', '_http._tcp.local.', 'ShellyPro4PM-94B97EC07650._http._tcp.local.'),
         ('update', '_http._tcp.local.', 'ShellyPro4PM-94B97EC07650._http._tcp.local.'),
     ]
+
+
+@pytest.mark.asyncio
+async def test_service_browser_does_not_try_to_send_if_not_ready():
+    """Test that the service browser does not try to send if not ready when rescheduling a type."""
+    service_added = asyncio.Event()
+    type_ = "_http._tcp.local."
+    registration_name = "nosend.%s" % type_
+
+    def on_service_state_change(zeroconf, service_type, state_change, name):
+        if name == registration_name:
+            if state_change is ServiceStateChange.Added:
+                service_added.set()
+
+    aiozc = AsyncZeroconf(interfaces=['127.0.0.1'])
+    zeroconf_browser = aiozc.zeroconf
+    await zeroconf_browser.async_wait_for_start()
+
+    expected_ttl = const._DNS_HOST_TTL
+    time_offset = 0.0
+
+    def _new_current_time_millis():
+        """Current system time in milliseconds"""
+        return (time.monotonic() * 1000) + (time_offset * 1000)
+
+    assert len(zeroconf_browser.engine.protocols) == 2
+
+    aio_zeroconf_registrar = AsyncZeroconf(interfaces=['127.0.0.1'])
+    zeroconf_registrar = aio_zeroconf_registrar.zeroconf
+    await aio_zeroconf_registrar.zeroconf.async_wait_for_start()
+    assert len(zeroconf_registrar.engine.protocols) == 2
+    with patch("zeroconf._services.browser.current_time_millis", _new_current_time_millis):
+        service_added = asyncio.Event()
+        browser = AsyncServiceBrowser(zeroconf_browser, type_, [on_service_state_change])
+        desc = {'path': '/~paulsm/'}
+        info = ServiceInfo(
+            type_, registration_name, 80, 0, 0, desc, "ash-2.local.", addresses=[socket.inet_aton("10.0.1.2")]
+        )
+        task = await aio_zeroconf_registrar.async_register_service(info)
+        await task
+
+        try:
+            await asyncio.wait_for(service_added.wait(), 1)
+            time_offset = 1000 * expected_ttl  # set the time to the end of the ttl
+            now = _new_current_time_millis()
+            browser.query_scheduler._next_time[type_] = now + (1000 * expected_ttl)
+            # Make sure the query schedule is to a time in the future
+            # so we will reschedule
+            with patch.object(
+                browser, "_async_send_ready_queries"
+            ) as _async_send_ready_queries, patch.object(
+                browser, "_async_send_ready_queries_schedule_next"
+            ) as _async_send_ready_queries_schedule_next:
+                # Reschedule the type to be sent in 1ms in the future
+                # to make sure the query is not sent
+                browser.reschedule_type(type_, now, now + 1)
+                assert not _async_send_ready_queries.called
+                await asyncio.sleep(0.01)
+                # Make sure it does happen after the sleep
+                assert _async_send_ready_queries_schedule_next.called
+        finally:
+            await aio_zeroconf_registrar.async_close()
+            await browser.async_cancel()
+            await aiozc.async_close()
