Skip to content

Commit 2bfbcbe

Browse files
authored
Make add_listener and remove_listener threadsafe (#794)
1 parent 6aac0eb commit 2bfbcbe

6 files changed

Lines changed: 59 additions & 31 deletions

File tree

tests/services/test_browser.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -315,6 +315,7 @@ def mock_incoming_msg(service_state_change: r.ServiceStateChange) -> r.DNSIncomi
315315
finally:
316316
assert len(zeroconf.listeners) == 1
317317
service_browser.cancel()
318+
time.sleep(0.2)
318319
assert len(zeroconf.listeners) == 0
319320
zeroconf.remove_all_service_listeners()
320321
zeroconf.close()
@@ -422,6 +423,7 @@ def _mock_get_expiration_time(self, percent):
422423
finally:
423424
assert len(zeroconf.listeners) == 1
424425
service_browser.cancel()
426+
time.sleep(0.2)
425427
assert len(zeroconf.listeners) == 0
426428
zeroconf.remove_all_service_listeners()
427429
zeroconf.close()

zeroconf/_core.py

Lines changed: 30 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -548,12 +548,38 @@ def add_listener(
548548
) -> None:
549549
"""Adds a listener for a given question. The listener will have
550550
its update_record method called when information is available to
551-
answer the question(s)."""
552-
self.record_manager.add_listener(listener, question)
551+
answer the question(s).
552+
553+
This function is threadsafe
554+
"""
555+
assert self.loop is not None
556+
self.loop.call_soon_threadsafe(self.record_manager.async_add_listener, listener, question)
553557

554558
def remove_listener(self, listener: RecordUpdateListener) -> None:
555-
"""Removes a listener."""
556-
self.record_manager.remove_listener(listener)
559+
"""Removes a listener.
560+
561+
This function is threadsafe
562+
"""
563+
assert self.loop is not None
564+
self.loop.call_soon_threadsafe(self.record_manager.async_remove_listener, listener)
565+
566+
def async_add_listener(
567+
self, listener: RecordUpdateListener, question: Optional[Union[DNSQuestion, List[DNSQuestion]]]
568+
) -> None:
569+
"""Adds a listener for a given question. The listener will have
570+
its update_record method called when information is available to
571+
answer the question(s).
572+
573+
This function is not threadsafe and must be called in the eventloop.
574+
"""
575+
self.record_manager.async_add_listener(listener, question)
576+
577+
def async_remove_listener(self, listener: RecordUpdateListener) -> None:
578+
"""Removes a listener.
579+
580+
This function is not threadsafe and must be called in the eventloop.
581+
"""
582+
self.record_manager.async_remove_listener(listener)
557583

558584
def handle_response(self, msg: DNSIncoming) -> None:
559585
"""Deal with incoming response packets. All answers

zeroconf/_handlers.py

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -387,20 +387,23 @@ def _async_mark_unique_cached_records_older_than_1s_to_expire(
387387
# Expire in 1s
388388
entry.set_created_ttl(now, 1)
389389

390-
def add_listener(
390+
def async_add_listener(
391391
self, listener: RecordUpdateListener, question: Optional[Union[DNSQuestion, List[DNSQuestion]]]
392392
) -> None:
393393
"""Adds a listener for a given question. The listener will have
394394
its update_record method called when information is available to
395-
answer the question(s)."""
395+
answer the question(s).
396+
397+
This function is not threadsafe and must be called in the eventloop.
398+
"""
396399
self.listeners.append(listener)
397400

398401
if question is None:
399402
return
400403

401404
questions = [question] if isinstance(question, DNSQuestion) else question
402405
assert self.zc.loop is not None
403-
self.zc.loop.call_soon_threadsafe(self._async_update_matching_records, listener, questions)
406+
self._async_update_matching_records(listener, questions)
404407

405408
def _async_update_matching_records(
406409
self, listener: RecordUpdateListener, questions: List[DNSQuestion]
@@ -422,10 +425,13 @@ def _async_update_matching_records(
422425
listener.async_update_records_complete()
423426
self.zc.async_notify_all()
424427

425-
def remove_listener(self, listener: RecordUpdateListener) -> None:
426-
"""Removes a listener."""
428+
def async_remove_listener(self, listener: RecordUpdateListener) -> None:
429+
"""Removes a listener.
430+
431+
This function is not threadsafe and must be called in the eventloop.
432+
"""
427433
try:
428434
self.listeners.remove(listener)
429-
self.zc.notify_all()
435+
self.zc.async_notify_all()
430436
except ValueError as e:
431437
log.exception('Failed to remove listener: %r', e)

zeroconf/_services/browser.py

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -242,14 +242,16 @@ def __init__(
242242
for h in handlers:
243243
self.service_state_changed.register_handler(h)
244244

245-
def _setup(self) -> None:
245+
def _async_start(self) -> None:
246246
"""Generate the next time and setup listeners.
247247
248248
Must be called by uses of this base class after they
249249
have finished setting their properties.
250250
"""
251251
self._generate_first_next_time()
252-
self.zc.add_listener(self, [DNSQuestion(type_, _TYPE_PTR, _CLASS_IN) for type_ in self.types])
252+
self.zc.async_add_listener(self, [DNSQuestion(type_, _TYPE_PTR, _CLASS_IN) for type_ in self.types])
253+
# Only start queries after the listener is installed
254+
self._browser_task = cast(asyncio.Task, asyncio.ensure_future(self.async_browser_task()))
253255

254256
def _generate_first_next_time(self) -> None:
255257
"""Generate the initial next query times.
@@ -374,10 +376,10 @@ def _fire_service_state_changed_event(self, event: Tuple[Tuple[str, str], Servic
374376
state_change=state_change,
375377
)
376378

377-
def cancel(self) -> None:
379+
def _async_cancel(self) -> None:
378380
"""Cancel the browser."""
379381
self.done = True
380-
self.zc.remove_listener(self)
382+
self.zc.async_remove_listener(self)
381383

382384
def generate_ready_queries(self) -> List[DNSOutgoing]:
383385
"""Generate the service browser query for any type that is due."""
@@ -454,20 +456,15 @@ def __init__(
454456
self.queue = get_best_available_queue()
455457
self.daemon = True
456458
self.start()
457-
self._setup()
458-
# Start queries after the listener is installed in _setup
459-
zc.loop.call_soon_threadsafe(self._async_start_browser)
459+
zc.loop.call_soon_threadsafe(self._async_start)
460460
self.name = "zeroconf-ServiceBrowser-%s-%s" % (
461461
'-'.join([type_[:-7] for type_ in self.types]),
462462
getattr(self, 'native_id', self.ident),
463463
)
464464

465-
def _async_start_browser(self) -> None:
466-
"""Start the browser from the event loop."""
467-
self._browser_task = cast(asyncio.Task, asyncio.ensure_future(self.async_browser_task()))
468-
469-
def _async_cancel_browser_soon(self) -> None:
465+
def _async_cancel_soon(self) -> None:
470466
"""Cancel the browser from the event loop."""
467+
self._async_cancel()
471468
if self._browser_task:
472469
asyncio.ensure_future(self._async_cancel_browser())
473470

@@ -476,8 +473,7 @@ def cancel(self) -> None:
476473
assert self.zc.loop is not None
477474
assert self.queue is not None
478475
self.queue.put(None)
479-
self.zc.loop.call_soon_threadsafe(self._async_cancel_browser_soon)
480-
super().cancel()
476+
self.zc.loop.call_soon_threadsafe(self._async_cancel_soon)
481477
self.join()
482478

483479
def run(self) -> None:

zeroconf/_services/info.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -421,7 +421,7 @@ async def async_request(
421421
last = now + timeout
422422
await zc.async_wait_for_start()
423423
try:
424-
zc.add_listener(self, None)
424+
zc.async_add_listener(self, None)
425425
while not self._is_complete:
426426
if last <= now:
427427
return False
@@ -436,7 +436,7 @@ async def async_request(
436436
await zc.async_wait(min(next_, last) - now)
437437
now = current_time_millis()
438438
finally:
439-
zc.remove_listener(self)
439+
zc.async_remove_listener(self)
440440

441441
return True
442442

zeroconf/aio.py

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
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, cast
25+
from typing import Awaitable, Callable, Dict, List, Optional, Tuple, Type, Union
2626

2727
from ._core import Zeroconf
2828
from ._dns import DNSQuestionType
@@ -86,15 +86,13 @@ def __init__(
8686
question_type: Optional[DNSQuestionType] = None,
8787
) -> None:
8888
super().__init__(zeroconf, type_, handlers, listener, addr, port, delay, question_type)
89-
self._setup()
90-
# Start queries after the listener is installed in _setup
91-
self._browser_task = cast(asyncio.Task, asyncio.ensure_future(self.async_browser_task()))
89+
self._async_start()
9290

9391
async def async_cancel(self) -> None:
9492
"""Cancel the browser."""
93+
self._async_cancel()
9594
with contextlib.suppress(asyncio.CancelledError):
9695
await self._async_cancel_browser()
97-
super().cancel()
9896

9997

10098
class AsyncZeroconfServiceTypes(ZeroconfServiceTypes):

0 commit comments

Comments
 (0)