Switch ServiceBrowser query scheduling to use call_later instead of a loop (#849)

bdraco · web-flow · commit a8c16231881d · 2021-06-27T07:08:54.000-10:00
- Simplifies scheduling as there is no more need to sleep in a loop as
  we now schedule future callbacks with call_later

- Simplifies cancelation as there is no more coroutine to cancel, only a timer handle
  We no longer have to handle the canceled error and cleaning up the awaitable

- Solves the infrequent test failures in test_backoff and test_integration
diff --git a/tests/services/test_browser.py b/tests/services/test_browser.py
@@ -496,7 +496,7 @@ def on_service_state_change(zeroconf, service_type, state_change, name):
                 else:
                     assert not got_query.is_set()
                 time_offset += initial_query_interval
-                zeroconf_browser.loop.call_soon_threadsafe(browser.query_scheduler.set_schedule_changed)
+                zeroconf_browser.loop.call_soon_threadsafe(browser.schedule_changed)
 
         finally:
             browser.cancel()
@@ -984,8 +984,8 @@ async def test_query_scheduler():
 
     # Test query interval is increasing
     assert query_scheduler.millis_to_wait(now - 1) == 1
-    assert query_scheduler.millis_to_wait(now) is None
-    assert query_scheduler.millis_to_wait(now + 1) is None
+    assert query_scheduler.millis_to_wait(now) is 0
+    assert query_scheduler.millis_to_wait(now + 1) is 0
 
     assert set(query_scheduler.process_ready_types(now)) == types_
     assert set(query_scheduler.process_ready_types(now)) == set()
@@ -1018,41 +1018,3 @@ async def test_query_scheduler():
     assert set(query_scheduler.process_ready_types(now + delay * 20)) == set()
 
     assert set(query_scheduler.process_ready_types(now + delay * 31)) == set(["_http._tcp.local."])
-
-
-@pytest.mark.asyncio
-async def test_query_scheduler_triggers_async_wait_ready_on_reschedule():
-    """Test that a reschedule wakes up the async_wait_ready."""
-    delay = const._BROWSER_TIME
-    types_ = set(["_hap._tcp.local.", "_http._tcp.local."])
-    query_scheduler = _services_browser.QueryScheduler(types_, delay, (0, 0))
-
-    now = current_time_millis()
-    query_scheduler.start(now)
-    assert set(query_scheduler.process_ready_types(now)) == types_
-    assert query_scheduler.millis_to_wait(now) == delay
-
-    task = asyncio.ensure_future(query_scheduler.async_wait_ready(now))
-    await asyncio.sleep(0)  # Start the task
-    await asyncio.sleep(0)  # Make sure its waiting
-    assert not task.done()
-    assert query_scheduler.millis_to_wait(now + 1) == delay - 1
-    query_scheduler.reschedule_type("_hap._tcp.local.", now + 1)
-    assert query_scheduler.millis_to_wait(now + 1) is None
-    await asyncio.wait_for(task, timeout=0.1)
-    assert task.done()
-
-    task2 = asyncio.ensure_future(query_scheduler.async_wait_ready(now + 10000))
-    assert set(query_scheduler.process_ready_types(now + 1)) == set(["_hap._tcp.local."])
-    assert not task2.done()
-    assert query_scheduler.millis_to_wait(now + 2) == delay - 2
-    query_scheduler.reschedule_type("_hap._tcp.local.", now + 2)
-    assert query_scheduler.millis_to_wait(now + 2) is None
-    await asyncio.wait_for(task2, timeout=0.1)
-    assert task2.done()
-    assert set(query_scheduler.process_ready_types(now + 10000)) == types_
-    assert query_scheduler.millis_to_wait(now + 10000) == delay * 2
-
-    task3 = asyncio.ensure_future(query_scheduler.async_wait_ready(now + 10000))
-    with pytest.raises(asyncio.TimeoutError):
-        await asyncio.wait_for(task3, timeout=0.1)
diff --git a/tests/test_aio.py b/tests/test_aio.py
@@ -688,7 +688,7 @@ def on_service_state_change(zeroconf, service_type, state_change, name):
 
     time_offset = 0.0
 
-    def current_time_millis():
+    def _new_current_time_millis():
         """Current system time in milliseconds"""
         return (time.time() * 1000) + (time_offset * 1000)
 
@@ -705,7 +705,6 @@ def send(out, addr=const._MDNS_ADDR, port=const._MDNS_PORT, v6_flow_scope=()):
                 unexpected_ttl.set()
 
         got_query.set()
-        got_query.clear()
 
         old_send(out, addr=addr, port=port, v6_flow_scope=v6_flow_scope)
 
@@ -734,7 +733,7 @@ def send(out, addr=const._MDNS_ADDR, port=const._MDNS_PORT, v6_flow_scope=()):
     ), patch.object(
         zeroconf_browser, "async_send", send
     ), patch(
-        "zeroconf._services.browser.current_time_millis", current_time_millis
+        "zeroconf._services.browser.current_time_millis", _new_current_time_millis
     ), patch.object(
         _services_browser, "_BROWSER_BACKOFF_LIMIT", int(expected_ttl / 4)
     ):
@@ -762,9 +761,11 @@ def send(out, addr=const._MDNS_ADDR, port=const._MDNS_PORT, v6_flow_scope=()):
             while nbr_answers < test_iterations:
                 # Increase simulated time shift by 1/4 of the TTL in seconds
                 time_offset += expected_ttl / 4
-                browser.query_scheduler.set_schedule_changed()
+                now = _new_current_time_millis()
+                browser.reschedule_type(type_, now)
                 sleep_count += 1
                 await asyncio.wait_for(got_query.wait(), 1)
+                got_query.clear()
                 # Prevent the test running indefinitely in an error condition
                 assert sleep_count < test_iterations * 4
             assert not unexpected_ttl.is_set()
diff --git a/zeroconf/_services/browser.py b/zeroconf/_services/browser.py
@@ -21,7 +21,6 @@
 """
 
 import asyncio
-import contextlib
 import queue
 import random
 import threading
@@ -39,7 +38,7 @@
     SignalRegistrationInterface,
 )
 from .._updates import RecordUpdate, RecordUpdateListener
-from .._utils.aio import get_best_available_queue, wait_event_or_timeout
+from .._utils.aio import get_best_available_queue
 from .._utils.name import service_type_name
 from .._utils.time import current_time_millis, millis_to_seconds
 from ..const import (
@@ -221,25 +220,17 @@ def _generate_first_next_time(self, now: float) -> None:
         next_time = now + delay
         self._next_time = {check_type_: next_time for check_type_ in self._types}
 
-    def millis_to_wait(self, now: float) -> Optional[float]:
+    def millis_to_wait(self, now: float) -> float:
         """Returns the number of milliseconds to wait for the next event."""
         # Wait for the type has the smallest next time
         next_time = min(self._next_time.values())
-        return None if next_time <= now else next_time - now
+        return 0 if next_time <= now else next_time - now
 
     def reschedule_type(self, type_: str, next_time: float) -> None:
         """Reschedule the query for a type to happen sooner."""
         if next_time >= self._next_time[type_]:
             return
-
         self._next_time[type_] = next_time
-        self.set_schedule_changed()
-
-    def set_schedule_changed(self) -> None:
-        """Set the event to unblock async_wait_ready to make sure the adjusted next time is seen."""
-        assert self._schedule_changed_event is not None
-        self._schedule_changed_event.set()
-        self._schedule_changed_event.clear()
 
     def process_ready_types(self, now: float) -> List[str]:
         """Generate a list of ready types that is due and schedule the next time."""
@@ -258,13 +249,6 @@ def process_ready_types(self, now: float) -> List[str]:
 
         return ready_types
 
-    async def async_wait_ready(self, now: float) -> None:
-        """Wait for at least one query to be ready."""
-        timeout = self.millis_to_wait(now)
-        if timeout:
-            assert self._schedule_changed_event is not None
-            await wait_event_or_timeout(self._schedule_changed_event, timeout=millis_to_seconds(timeout))
-
 
 class _ServiceBrowserBase(RecordUpdateListener):
     """Base class for ServiceBrowser."""
@@ -302,7 +286,6 @@ def __init__(
         for check_type_ in self.types:
             # Will generate BadTypeInNameException on a bad name
             service_type_name(check_type_, strict=False)
-        self._browser_task: Optional[asyncio.Task] = None
         self.zc = zc
         self.addr = addr
         self.port = port
@@ -313,6 +296,8 @@ def __init__(
         self.query_scheduler = QueryScheduler(self.types, delay, _FIRST_QUERY_DELAY_RANDOM_INTERVAL)
         self.queue: Optional[queue.Queue] = None
         self.done = False
+        self._first_request: bool = True
+        self._next_send_timer: Optional[asyncio.TimerHandle] = None
 
         if hasattr(handlers, 'add_service'):
             listener = cast('ServiceListener', handlers)
@@ -335,7 +320,7 @@ def _async_start(self) -> None:
         self.query_scheduler.start(current_time_millis())
         self.zc.async_add_listener(self, [DNSQuestion(type_, _TYPE_PTR, _CLASS_IN) for type_ in self.types])
         # Only start queries after the listener is installed
-        self._browser_task = cast(asyncio.Task, asyncio.ensure_future(self.async_browser_task()))
+        asyncio.ensure_future(self._async_start_query_sender())
 
     @property
     def service_state_changed(self) -> SignalRegistrationInterface:
@@ -378,9 +363,7 @@ def _async_process_record_update(
             elif expired:
                 self._enqueue_callback(ServiceStateChange.Removed, record.name, record.alias)
             else:
-                self.query_scheduler.reschedule_type(
-                    record.name, record.get_expiration_time(_EXPIRE_REFRESH_TIME_PERCENT)
-                )
+                self.reschedule_type(record.name, record.get_expiration_time(_EXPIRE_REFRESH_TIME_PERCENT))
             return
 
         # If its expired or already exists in the cache it cannot be updated.
@@ -448,6 +431,7 @@ def _fire_service_state_changed_event(self, event: Tuple[Tuple[str, str], Servic
     def _async_cancel(self) -> None:
         """Cancel the browser."""
         self.done = True
+        self._cancel_send_timer()
         self.zc.async_remove_listener(self)
 
     def _generate_ready_queries(self, first_request: bool) -> List[DNSOutgoing]:
@@ -464,28 +448,40 @@ def _generate_ready_queries(self, first_request: bool) -> List[DNSOutgoing]:
         question_type = DNSQuestionType.QU if not self.question_type and first_request else self.question_type
         return generate_service_query(self.zc, now, ready_types, self.multicast, question_type)
 
-    async def async_browser_task(self) -> None:
-        """Run the browser task."""
+    async def _async_start_query_sender(self) -> None:
+        """Start scheduling queries."""
         await self.zc.async_wait_for_start()
-        first_request = True
-        while True:
-            await self.query_scheduler.async_wait_ready(current_time_millis())
-            outs = self._generate_ready_queries(first_request)
-            if not outs:
-                continue
+        self._async_send_ready_queries_schedule_next()
+
+    def _cancel_send_timer(self) -> None:
+        """Cancel the next send."""
+        if self._next_send_timer:
+            self._next_send_timer.cancel()
 
-            first_request = False
+    def reschedule_type(self, type_: str, next_time: float) -> None:
+        """Reschedule a type to be refreshed in the future."""
+        self.query_scheduler.reschedule_type(type_, next_time)
+        self.schedule_changed()
+
+    def schedule_changed(self) -> None:
+        """Called when the schedule has changed."""
+        self._cancel_send_timer()
+        self._async_send_ready_queries_schedule_next()
+
+    def _async_send_ready_queries_schedule_next(self) -> None:
+        """Send any ready queries and scheule the next time."""
+        if self.done or self.zc.done:
+            return
+
+        outs = self._generate_ready_queries(self._first_request)
+        if outs:
+            self._first_request = False
             for out in outs:
                 self.zc.async_send(out, addr=self.addr, port=self.port)
 
-    async def _async_cancel_browser(self) -> None:
-        """Cancel the browser."""
-        assert self._browser_task is not None
-        self._browser_task.cancel()
-        browser_task = self._browser_task
-        self._browser_task = None
-        with contextlib.suppress(asyncio.CancelledError):
-            await browser_task
+        assert self.zc.loop is not None
+        delay = millis_to_seconds(self.query_scheduler.millis_to_wait(current_time_millis()))
+        self._next_send_timer = self.zc.loop.call_later(delay, self._async_send_ready_queries_schedule_next)
 
 
 class ServiceBrowser(_ServiceBrowserBase, threading.Thread):
@@ -523,18 +519,12 @@ def __init__(
             getattr(self, 'native_id', self.ident),
         )
 
-    def _async_cancel_soon(self) -> None:
-        """Cancel the browser from the event loop."""
-        self._async_cancel()
-        if self._browser_task:
-            asyncio.ensure_future(self._async_cancel_browser())
-
     def cancel(self) -> None:
         """Cancel the browser."""
         assert self.zc.loop is not None
         assert self.queue is not None
         self.queue.put(None)
-        self.zc.loop.call_soon_threadsafe(self._async_cancel_soon)
+        self.zc.loop.call_soon_threadsafe(self._async_cancel)
         self.join()
 
     def run(self) -> None:
diff --git a/zeroconf/aio.py b/zeroconf/aio.py
@@ -91,8 +91,6 @@ def __init__(
     async def async_cancel(self) -> None:
         """Cancel the browser."""
         self._async_cancel()
-        with contextlib.suppress(asyncio.CancelledError):
-            await self._async_cancel_browser()
 
 
 class AsyncZeroconfServiceTypes(ZeroconfServiceTypes):
