feat: improve performance of DNSCache backend (#1415)

bdraco · web-flow · commit 1df2e691ff11 · 2024-09-24T13:49:26.000-05:00
diff --git a/src/zeroconf/_cache.pxd b/src/zeroconf/_cache.pxd
@@ -13,9 +13,12 @@ from ._dns cimport (
 
 
 cdef object _UNIQUE_RECORD_TYPES
-cdef object _TYPE_PTR
+cdef unsigned int _TYPE_PTR
 cdef cython.uint _ONE_SECOND
 
+@cython.locals(
+    record_cache=dict,
+)
 cdef _remove_key(cython.dict cache, object key, DNSRecord record)
 
 
@@ -42,7 +45,7 @@ cdef class DNSCache:
         records=cython.dict,
         record=DNSRecord,
     )
-    cpdef list async_all_by_details(self, str name, object type_, object class_)
+    cpdef list async_all_by_details(self, str name, unsigned int type_, unsigned int class_)
 
     cpdef cython.dict async_entries_with_name(self, str name)
 
@@ -51,23 +54,35 @@ cdef class DNSCache:
     @cython.locals(
         cached_entry=DNSRecord,
     )
-    cpdef DNSRecord get_by_details(self, str name, object type_, object class_)
+    cpdef DNSRecord get_by_details(self, str name, unsigned int type_, unsigned int class_)
 
     @cython.locals(
         records=cython.dict,
         entry=DNSRecord,
     )
-    cpdef cython.list get_all_by_details(self, str name, object type_, object class_)
+    cpdef cython.list get_all_by_details(self, str name, unsigned int type_, unsigned int class_)
 
     @cython.locals(
         store=cython.dict,
+        service_record=DNSService
     )
     cdef bint _async_add(self, DNSRecord record)
 
+    @cython.locals(
+        service_record=DNSService
+    )
     cdef void _async_remove(self, DNSRecord record)
 
     @cython.locals(
         record=DNSRecord,
         created_double=double,
     )
     cpdef void async_mark_unique_records_older_than_1s_to_expire(self, cython.set unique_types, object answers, double now)
+
+    cpdef entries_with_name(self, str name)
+
+    @cython.locals(
+        record=DNSRecord,
+        now=double
+    )
+    cpdef current_entry_with_name_and_alias(self, str name, str alias)
diff --git a/src/zeroconf/_cache.py b/src/zeroconf/_cache.py
@@ -49,8 +49,9 @@ def _remove_key(cache: _DNSRecordCacheType, key: _str, record: _DNSRecord) -> No
 
     This function must be run in from event loop.
     """
-    del cache[key][record]
-    if not cache[key]:
+    record_cache = cache[key]
+    del record_cache[record]
+    if not record_cache:
         del cache[key]
 
 
@@ -81,7 +82,8 @@ def _async_add(self, record: _DNSRecord) -> bool:
         new = record not in store and not isinstance(record, DNSNsec)
         store[record] = record
         if isinstance(record, DNSService):
-            self.service_cache.setdefault(record.server_key, {})[record] = record
+            service_record = record
+            self.service_cache.setdefault(record.server_key, {})[service_record] = service_record
         return new
 
     def async_add_records(self, entries: Iterable[DNSRecord]) -> bool:
@@ -103,7 +105,8 @@ def _async_remove(self, record: _DNSRecord) -> None:
         This function must be run in from event loop.
         """
         if isinstance(record, DNSService):
-            _remove_key(self.service_cache, record.server_key, record)
+            service_record = record
+            _remove_key(self.service_cache, service_record.server_key, service_record)
         _remove_key(self.cache, record.key, record)
 
     def async_remove_records(self, entries: Iterable[DNSRecord]) -> None:
