feat: speed up service browser queries when browsing many types (#1311)

bdraco · web-flow · commit d192d33b1f05 · 2023-11-12T13:51:59.000-06:00
diff --git a/src/zeroconf/_services/browser.pxd b/src/zeroconf/_services/browser.pxd
@@ -2,6 +2,7 @@
 import cython
 
 from .._cache cimport DNSCache
+from .._history cimport QuestionHistory
 from .._protocol.outgoing cimport DNSOutgoing, DNSPointer, DNSQuestion, DNSRecord
 from .._record_update cimport RecordUpdate
 from .._updates cimport RecordUpdateListener
@@ -11,7 +12,7 @@ from . cimport Signal, SignalRegistrationInterface
 
 cdef bint TYPE_CHECKING
 cdef object cached_possible_types
-cdef cython.uint _EXPIRE_REFRESH_TIME_PERCENT
+cdef cython.uint _EXPIRE_REFRESH_TIME_PERCENT, _MAX_MSG_TYPICAL, _DNS_PACKET_HEADER_LEN
 cdef cython.uint _TYPE_PTR
 cdef object SERVICE_STATE_CHANGE_ADDED, SERVICE_STATE_CHANGE_REMOVED, SERVICE_STATE_CHANGE_UPDATED
 cdef cython.set _ADDRESS_RECORD_TYPES
@@ -24,8 +25,16 @@ cdef class _DNSPointerOutgoingBucket:
 
     cpdef add(self, cython.uint max_compressed_size, DNSQuestion question, cython.set answers)
 
+@cython.locals(cache=DNSCache, question_history=QuestionHistory, record=DNSRecord)
+cpdef generate_service_query(
+    object zc,
+    float now,
+    list type_,
+    bint multicast,
+    object question_type
+)
 
-@cython.locals(answer=DNSPointer)
+@cython.locals(answer=DNSPointer, query_buckets=list, question=DNSQuestion, max_compressed_size=cython.uint, max_bucket_size=cython.uint, query_bucket=_DNSPointerOutgoingBucket)
 cdef _group_ptr_queries_with_known_answers(object now, object multicast, cython.dict question_with_known_answers)
 
 cdef class QueryScheduler:
diff --git a/src/zeroconf/_services/browser.py b/src/zeroconf/_services/browser.py
@@ -164,33 +164,31 @@ def _group_ptr_queries_with_known_answers(
 
 
 def generate_service_query(
-    zc: 'Zeroconf',
-    now: float,
-    types_: List[str],
-    multicast: bool = True,
-    question_type: Optional[DNSQuestionType] = None,
+    zc: 'Zeroconf', now: float_, types_: List[str], multicast: bool, question_type: Optional[DNSQuestionType]
 ) -> List[DNSOutgoing]:
     """Generate a service query for sending with zeroconf.send."""
     questions_with_known_answers: _QuestionWithKnownAnswers = {}
     qu_question = not multicast if question_type is None else question_type == DNSQuestionType.QU
+    question_history = zc.question_history
+    cache = zc.cache
     for type_ in types_:
         question = DNSQuestion(type_, _TYPE_PTR, _CLASS_IN)
         question.unicast = qu_question
         known_answers = {
             record
-            for record in zc.cache.get_all_by_details(type_, _TYPE_PTR, _CLASS_IN)
-            if not record.is_stale(now)
+            for record in cache.get_all_by_details(type_, _TYPE_PTR, _CLASS_IN)
+            if record.is_stale(now) is False
         }
-        if not qu_question and zc.question_history.suppresses(question, now, known_answers):
+        if not qu_question and question_history.suppresses(question, now, known_answers):
             log.debug("Asking %s was suppressed by the question history", question)
             continue
         if TYPE_CHECKING:
             pointer_known_answers = cast(Set[DNSPointer], known_answers)
         else:
             pointer_known_answers = known_answers
         questions_with_known_answers[question] = pointer_known_answers
-        if not qu_question:
-            zc.question_history.add_question_at_time(question, now, known_answers)
+        if qu_question is False:
+            question_history.add_question_at_time(question, now, known_answers)
 
     return _group_ptr_queries_with_known_answers(now, multicast, questions_with_known_answers)
 
diff --git a/tests/services/test_browser.py b/tests/services/test_browser.py
@@ -1010,32 +1010,34 @@ async def test_generate_service_query_suppress_duplicate_questions():
     assert zc.question_history.suppresses(question, now, other_known_answers)
 
     # The known answer list is different, do not suppress
-    outs = _services_browser.generate_service_query(zc, now, [name], multicast=True)
+    outs = _services_browser.generate_service_query(zc, now, [name], multicast=True, question_type=None)
     assert outs
 
     zc.cache.async_add_records([answer])
     # The known answer list contains all the asked questions in the history
     # we should suppress
 
-    outs = _services_browser.generate_service_query(zc, now, [name], multicast=True)
+    outs = _services_browser.generate_service_query(zc, now, [name], multicast=True, question_type=None)
     assert not outs
 
     # We do not suppress once the question history expires
-    outs = _services_browser.generate_service_query(zc, now + 1000, [name], multicast=True)
+    outs = _services_browser.generate_service_query(
+        zc, now + 1000, [name], multicast=True, question_type=None
+    )
     assert outs
 
     # We do not suppress QU queries ever
-    outs = _services_browser.generate_service_query(zc, now, [name], multicast=False)
+    outs = _services_browser.generate_service_query(zc, now, [name], multicast=False, question_type=None)
     assert outs
 
     zc.question_history.async_expire(now + 2000)
     # No suppression after clearing the history
-    outs = _services_browser.generate_service_query(zc, now, [name], multicast=True)
+    outs = _services_browser.generate_service_query(zc, now, [name], multicast=True, question_type=None)
     assert outs
 
     # The previous query we just sent is still remembered and
     # the next one is suppressed
-    outs = _services_browser.generate_service_query(zc, now, [name], multicast=True)
+    outs = _services_browser.generate_service_query(zc, now, [name], multicast=True, question_type=None)
     assert not outs
 
     await aiozc.async_close()
