Skip to content

Commit 4d0a8f3

Browse files
bdracobdraco
authored
Run ServiceBrowser queries in the event loop (#752)
1 parent e7adce2 commit 4d0a8f3

8 files changed

Lines changed: 95 additions & 75 deletions

File tree

‎tests/services/test_browser.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -348,7 +348,7 @@ def test_backoff():
348348
zeroconf_browser = Zeroconf(interfaces=['127.0.0.1'])
349349

350350
# we are going to patch the zeroconf send to check query transmission
351-
old_send = zeroconf_browser.send
351+
old_send = zeroconf_browser.async_send
352352

353353
time_offset = 0.0
354354
start_time = time.time() * 1000
@@ -366,7 +366,7 @@ def send(out, addr=const._MDNS_ADDR, port=const._MDNS_PORT):
366366
# patch the zeroconf send
367367
# patch the zeroconf current_time_millis
368368
# patch the backoff limit to prevent test running forever
369-
with unittest.mock.patch.object(zeroconf_browser, "send", send), unittest.mock.patch.object(
369+
with unittest.mock.patch.object(zeroconf_browser, "async_send", send), unittest.mock.patch.object(
370370
_services_browser, "current_time_millis", current_time_millis
371371
), unittest.mock.patch.object(_services_browser, "_BROWSER_BACKOFF_LIMIT", 10):
372372
# dummy service callback
@@ -432,7 +432,7 @@ def on_service_state_change(zeroconf, service_type, state_change, name):
432432
zeroconf_browser = Zeroconf(interfaces=['127.0.0.1'])
433433

434434
# we are going to patch the zeroconf send to check packet sizes
435-
old_send = zeroconf_browser.send
435+
old_send = zeroconf_browser.async_send
436436

437437
time_offset = 0.0
438438

@@ -459,7 +459,7 @@ def send(out, addr=const._MDNS_ADDR, port=const._MDNS_PORT):
459459
# patch the zeroconf send
460460
# patch the zeroconf current_time_millis
461461
# patch the backoff limit to ensure we always get one query every 1/4 of the DNS TTL
462-
with unittest.mock.patch.object(zeroconf_browser, "send", send), unittest.mock.patch.object(
462+
with unittest.mock.patch.object(zeroconf_browser, "async_send", send), unittest.mock.patch.object(
463463
_services_browser, "current_time_millis", current_time_millis
464464
), unittest.mock.patch.object(_services_browser, "_BROWSER_BACKOFF_LIMIT", int(expected_ttl / 4)):
465465
service_added = Event()

‎tests/test_aio.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,7 @@ def update_service(self, zeroconf: Zeroconf, type: str, name: str) -> None:
109109
calls.append(("update", type, name))
110110

111111
listener = MyListener()
112+
112113
aiozc.zeroconf.add_service_listener(type_, listener)
113114

114115
desc = {'path': '/~paulsm/'}

‎tests/test_init.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ def test_lots_of_names(self):
8282
self.verify_name_change(zc, type_, name, server_count)
8383

8484
# we are going to patch the zeroconf send to check packet sizes
85-
old_send = zc.send
85+
old_send = zc.async_send
8686

8787
longest_packet_len = 0
8888
longest_packet = None # type: Optional[r.DNSOutgoing]
@@ -97,7 +97,7 @@ def send(out, addr=const._MDNS_ADDR, port=const._MDNS_PORT):
9797
old_send(out, addr=addr, port=port)
9898

9999
# patch the zeroconf send
100-
with unittest.mock.patch.object(zc, "send", send):
100+
with unittest.mock.patch.object(zc, "async_send", send):
101101

102102
# dummy service callback
103103
def on_service_state_change(zeroconf, service_type, state_change, name):

‎zeroconf/_core.py‎

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -291,7 +291,6 @@ def __init__(
291291
self.query_handler = QueryHandler(self.registry, self.cache)
292292
self.record_manager = RecordManager(self)
293293

294-
self.condition = threading.Condition()
295294
self.async_condition: Optional[asyncio.Condition] = None
296295
self.loop: Optional[asyncio.AbstractEventLoop] = None
297296
self._loop_thread: Optional[threading.Thread] = None
@@ -338,10 +337,9 @@ def listeners(self) -> List[RecordUpdateListener]:
338337
return self.record_manager.listeners
339338

340339
def wait(self, timeout: float) -> None:
341-
"""Calling thread waits for a given number of milliseconds or
342-
until notified."""
343-
with self.condition:
344-
self.condition.wait(millis_to_seconds(timeout))
340+
"""Calling task waits for a given number of milliseconds or until notified."""
341+
assert self.loop is not None
342+
asyncio.run_coroutine_threadsafe(self.async_wait(timeout), self.loop).result()
345343

346344
async def async_wait(self, timeout: float) -> None:
347345
"""Calling task waits for a given number of milliseconds or until notified."""
@@ -361,10 +359,8 @@ def async_notify_all(self) -> None:
361359
async def _async_notify_all(self) -> None:
362360
"""Notify all async listeners."""
363361
assert self.async_condition is not None
364-
with self.condition:
365-
self.condition.notify_all()
366-
async with self.async_condition:
367-
self.async_condition.notify_all()
362+
async with self.async_condition:
363+
self.async_condition.notify_all()
368364

369365
def get_service_info(self, type_: str, name: str, timeout: int = 3000) -> Optional[ServiceInfo]:
370366
"""Returns network's service information for a particular

‎zeroconf/_handlers.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -410,7 +410,7 @@ def _async_update_matching_records(
410410
return
411411
listener.async_update_records(self.zc, now, records)
412412
listener.async_update_records_complete()
413-
self.zc.notify_all()
413+
self.zc.async_notify_all()
414414

415415
def remove_listener(self, listener: RecordUpdateListener) -> None:
416416
"""Removes a listener."""

‎zeroconf/_services/browser.py‎

Lines changed: 70 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,9 @@
2020
USA
2121
"""
2222

23+
import asyncio
24+
import contextlib
25+
import queue
2326
import threading
2427
import warnings
2528
from collections import OrderedDict
@@ -35,6 +38,7 @@
3538
Signal,
3639
SignalRegistrationInterface,
3740
)
41+
from .._utils.aio import get_best_available_queue, get_running_loop, wait_condition_or_timeout
3842
from .._utils.name import service_type_name
3943
from .._utils.time import current_time_millis, millis_to_seconds
4044
from ..const import (
@@ -180,6 +184,7 @@ def __init__(
180184
for check_type_ in self.types:
181185
# Will generate BadTypeInNameException on a bad name
182186
service_type_name(check_type_, strict=False)
187+
self._browser_task: Optional[asyncio.Task] = None
183188
self.zc = zc
184189
self.addr = addr
185190
self.port = port
@@ -190,7 +195,7 @@ def __init__(
190195
self._pending_handlers: OrderedDict[Tuple[str, str], ServiceStateChange] = OrderedDict()
191196
self._handlers_to_call: OrderedDict[Tuple[str, str], ServiceStateChange] = OrderedDict()
192197
self._service_state_changed = Signal()
193-
198+
self.queue: Optional[queue.Queue] = None
194199
self.done = False
195200

196201
if hasattr(handlers, 'add_service'):
@@ -341,6 +346,47 @@ def _seconds_to_wait(self) -> Optional[float]:
341346

342347
return millis_to_seconds(next_time - now)
343348

349+
async def async_browser_task(self) -> None:
350+
"""Run the browser task."""
351+
await self.zc.async_wait_for_start()
352+
assert self.zc.async_condition is not None
353+
while True:
354+
timeout = self._seconds_to_wait()
355+
if timeout:
356+
async with self.zc.async_condition:
357+
# We must check again while holding the condition
358+
# in case the other thread has added to _handlers_to_call
359+
# between when we checked above when we were not
360+
# holding the condition
361+
if not self._handlers_to_call:
362+
await wait_condition_or_timeout(self.zc.async_condition, timeout)
363+
364+
outs = self.generate_ready_queries()
365+
for out in outs:
366+
self.zc.async_send(out, addr=self.addr, port=self.port)
367+
368+
if not self._handlers_to_call:
369+
continue
370+
371+
(name_type, state_change) = self._handlers_to_call.popitem(False)
372+
if self.queue:
373+
self.queue.put((name_type, state_change))
374+
continue
375+
376+
self._service_state_changed.fire(
377+
zeroconf=self.zc,
378+
service_type=name_type[1],
379+
name=name_type[0],
380+
state_change=state_change,
381+
)
382+
383+
async def _async_cancel_browser(self) -> None:
384+
"""Cancel the browser."""
385+
assert self._browser_task is not None
386+
self._browser_task.cancel()
387+
with contextlib.suppress(asyncio.CancelledError):
388+
await self._browser_task
389+
344390

345391
class ServiceBrowser(_ServiceBrowserBase, threading.Thread):
346392
"""Used to browse for a service of a specific type.
@@ -361,42 +407,45 @@ def __init__(
361407
) -> None:
362408
threading.Thread.__init__(self)
363409
super().__init__(zc, type_, handlers=handlers, listener=listener, addr=addr, port=port, delay=delay)
410+
self.queue = get_best_available_queue()
364411
self.daemon = True
365412
self.start()
366413
self.name = "zeroconf-ServiceBrowser-%s-%s" % (
367414
'-'.join([type_[:-7] for type_ in self.types]),
368415
getattr(self, 'native_id', self.ident),
369416
)
417+
assert self.zc.loop is not None
418+
if get_running_loop() == self.zc.loop:
419+
self._browser_task = cast(asyncio.Task, asyncio.ensure_future(self.async_browser_task()))
420+
return
421+
self._browser_task = cast(
422+
asyncio.Task,
423+
asyncio.run_coroutine_threadsafe(self._async_browser_task(), self.zc.loop).result(),
424+
)
425+
426+
async def _async_browser_task(self) -> asyncio.Task:
427+
return cast(asyncio.Task, asyncio.ensure_future(self.async_browser_task()))
370428

371429
def cancel(self) -> None:
372430
"""Cancel the browser."""
431+
assert self.zc.loop is not None
432+
assert self.queue is not None
433+
self.queue.put(None)
434+
if get_running_loop() == self.zc.loop:
435+
asyncio.ensure_future(self._async_cancel_browser())
436+
else:
437+
asyncio.run_coroutine_threadsafe(self._async_cancel_browser(), self.zc.loop).result()
373438
super().cancel()
374439
self.join()
375440

376441
def run(self) -> None:
377442
"""Run the browser thread."""
443+
assert self.queue is not None
378444
while True:
379-
timeout = self._seconds_to_wait()
380-
if timeout:
381-
with self.zc.condition:
382-
# We must check again while holding the condition
383-
# in case the other thread has added to _handlers_to_call
384-
# between when we checked above when we were not
385-
# holding the condition
386-
if not self._handlers_to_call:
387-
self.zc.condition.wait(timeout)
388-
389-
if self.zc.done or self.done:
445+
event = self.queue.get()
446+
if event is None:
390447
return
391-
392-
outs = self.generate_ready_queries()
393-
for out in outs:
394-
self.zc.send(out, addr=self.addr, port=self.port)
395-
396-
if not self._handlers_to_call:
397-
continue
398-
399-
(name_type, state_change) = self._handlers_to_call.popitem(False)
448+
name_type, state_change = event
400449
self._service_state_changed.fire(
401450
zeroconf=self.zc,
402451
service_type=name_type[1],

‎zeroconf/_utils/aio.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,17 @@
2222

2323
import asyncio
2424
import contextlib
25+
import queue
2526
from typing import Optional, Set, cast
2627

2728

29+
def get_best_available_queue() -> queue.Queue:
30+
"""Create the best available queue type."""
31+
if hasattr(queue, "SimpleQueue"):
32+
return queue.SimpleQueue() # type: ignore # pylint: disable=all
33+
return queue.Queue()
34+
35+
2836
# Switch to asyncio.wait_for once https://bugs.python.org/issue39032 is fixed
2937
async def wait_condition_or_timeout(condition: asyncio.Condition, timeout: float) -> None:
3038
"""Wait for a condition or timeout."""

‎zeroconf/aio.py‎

Lines changed: 4 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -22,14 +22,13 @@
2222
import asyncio
2323
import contextlib
2424
from types import TracebackType # noqa # used in type hints
25-
from typing import Awaitable, Callable, Dict, List, Optional, Tuple, Type, Union
25+
from typing import Awaitable, Callable, Dict, List, Optional, Tuple, Type, Union, cast
2626

2727
from ._core import Zeroconf
2828
from ._exceptions import NonUniqueNameException
2929
from ._services.browser import _ServiceBrowserBase
3030
from ._services.info import ServiceInfo, instance_name_from_service_info
3131
from ._services.types import ZeroconfServiceTypes
32-
from ._utils.aio import wait_condition_or_timeout
3332
from ._utils.net import IPVersion, InterfaceChoice, InterfacesType
3433
from ._utils.time import millis_to_seconds
3534
from .const import (
@@ -83,46 +82,13 @@ def __init__(
8382
port: int = _MDNS_PORT,
8483
delay: int = _BROWSER_TIME,
8584
) -> None:
86-
self.aiozc = aiozc
8785
super().__init__(aiozc.zeroconf, type_, handlers, listener, addr, port, delay) # type: ignore
88-
self._browser_task = asyncio.ensure_future(self.async_run())
86+
self._browser_task = cast(asyncio.Task, asyncio.ensure_future(self.async_browser_task()))
8987

9088
async def async_cancel(self) -> None:
9189
"""Cancel the browser."""
92-
self.cancel()
93-
self._browser_task.cancel()
94-
with contextlib.suppress(asyncio.CancelledError):
95-
await self._browser_task
96-
97-
async def async_run(self) -> None:
98-
"""Run the browser task."""
99-
await self.aiozc.zeroconf.async_wait_for_start()
100-
assert self.aiozc.zeroconf.async_condition is not None
101-
while True:
102-
timeout = self._seconds_to_wait()
103-
if timeout:
104-
async with self.aiozc.zeroconf.async_condition:
105-
# We must check again while holding the condition
106-
# in case the other thread has added to _handlers_to_call
107-
# between when we checked above when we were not
108-
# holding the condition
109-
if not self._handlers_to_call:
110-
await wait_condition_or_timeout(self.aiozc.zeroconf.async_condition, timeout)
111-
112-
outs = self.generate_ready_queries()
113-
for out in outs:
114-
self.aiozc.zeroconf.async_send(out, addr=self.addr, port=self.port)
115-
116-
if not self._handlers_to_call:
117-
continue
118-
119-
(name_type, state_change) = self._handlers_to_call.popitem(False)
120-
self._service_state_changed.fire(
121-
zeroconf=self.aiozc,
122-
service_type=name_type[1],
123-
name=name_type[0],
124-
state_change=state_change,
125-
)
90+
await self._async_cancel_browser()
91+
super().cancel()
12692

12793

12894
class AsyncZeroconfServiceTypes(ZeroconfServiceTypes):

0 commit comments

Comments
 (0)