feat: significantly speed up writing outgoing dns records (#1260)

bdraco · web-flow · commit bf2f3660a1f3 · 2023-09-07T14:55:52.000-05:00
diff --git a/src/zeroconf/_protocol/outgoing.pxd b/src/zeroconf/_protocol/outgoing.pxd
@@ -17,14 +17,18 @@ cdef cython.uint _MAX_MSG_TYPICAL
 
 cdef object TYPE_CHECKING
 
+cdef object PACK_BYTE
+cdef object PACK_SHORT
+cdef object PACK_LONG
+
 cdef class DNSOutgoing:
 
     cdef public unsigned int flags
     cdef public object finished
     cdef public object id
     cdef public bint multicast
     cdef public cython.list packets_data
-    cdef public object names
+    cdef public cython.dict names
     cdef public cython.list data
     cdef public unsigned int size
     cdef public object allow_long
@@ -53,24 +57,32 @@ cdef class DNSOutgoing:
     )
     cdef _write_record(self, DNSRecord record, object now)
 
-    cdef _write_record_class(self, object record)
+    cdef _write_record_class(self, DNSEntry record)
 
     cdef _check_data_limit_or_rollback(self, object start_data_length, object start_size)
 
     cdef _write_questions_from_offset(self, object questions_offset)
 
     cdef _write_answers_from_offset(self, object answer_offset)
 
-    cdef _write_records_from_offset(self, object records, object offset)
+    cdef _write_records_from_offset(self, cython.list records, object offset)
 
     cdef _has_more_to_add(self, object questions_offset, object answer_offset, object authority_offset, object additional_offset)
 
     cdef _write_ttl(self, DNSRecord record, object now)
 
-    cpdef write_name(self, object name)
+    @cython.locals(
+        labels=cython.list,
+        label=cython.str,
+    )
+    cpdef write_name(self, cython.str name)
 
     cpdef write_short(self, object value)
 
+    cpdef write_string(self, cython.bytes value)
+
+    cpdef _write_utf(self, cython.str value)
+
     @cython.locals(
         questions_offset=object,
         answer_offset=object,
diff --git a/src/zeroconf/_protocol/outgoing.py b/src/zeroconf/_protocol/outgoing.py
@@ -22,6 +22,7 @@
 
 import enum
 import logging
+from struct import Struct
 from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Union
 
 from .._cache import DNSCache
@@ -43,10 +44,16 @@
 str_ = str
 float_ = float
 int_ = int
+bytes_ = bytes
 DNSQuestion_ = DNSQuestion
 DNSRecord_ = DNSRecord
 
 
+PACK_BYTE = Struct('>B').pack
+PACK_SHORT = Struct('>H').pack
+PACK_LONG = Struct('>L').pack
+
+
 class State(enum.Enum):
     init = 0
     finished = 1
@@ -200,35 +207,35 @@ def add_question_or_all_cache(
 
     def _write_byte(self, value: int_) -> None:
         """Writes a single byte to the packet"""
-        self.data.append(value.to_bytes(1, 'big'))
+        self.data.append(PACK_BYTE(value))
         self.size += 1
 
     def _insert_short_at_start(self, value: int_) -> None:
         """Inserts an unsigned short at the start of the packet"""
-        self.data.insert(0, value.to_bytes(2, 'big'))
+        self.data.insert(0, PACK_SHORT(value))
 
     def _replace_short(self, index: int_, value: int_) -> None:
         """Replaces an unsigned short in a certain position in the packet"""
-        self.data[index] = value.to_bytes(2, 'big')
+        self.data[index] = PACK_SHORT(value)
 
     def write_short(self, value: int_) -> None:
         """Writes an unsigned short to the packet"""
-        self.data.append(value.to_bytes(2, 'big'))
+        self.data.append(PACK_SHORT(value))
         self.size += 2
 
     def _write_int(self, value: Union[float, int]) -> None:
         """Writes an unsigned integer to the packet"""
-        self.data.append(int(value).to_bytes(4, 'big'))
+        self.data.append(PACK_LONG(int(value)))
         self.size += 4
 
-    def write_string(self, value: bytes) -> None:
+    def write_string(self, value: bytes_) -> None:
         """Writes a string to the packet"""
         if TYPE_CHECKING:
             assert isinstance(value, bytes)
         self.data.append(value)
         self.size += len(value)
 
-    def _write_utf(self, s: str) -> None:
+    def _write_utf(self, s: str_) -> None:
         """Writes a UTF-8 string of a given length to the packet"""
         utfstr = s.encode('utf-8')
         length = len(utfstr)
@@ -446,7 +453,8 @@ def _packets(self) -> List[bytes]:
                 questions_offset, answer_offset, authority_offset, additional_offset
             ):
                 # https://datatracker.ietf.org/doc/html/rfc6762#section-7.2
-                log.debug("Setting TC flag")
+                if debug_enable:  # pragma: no branch
+                    log.debug("Setting TC flag")
                 self._insert_short_at_start(self.flags | _FLAGS_TC)
             else:
                 self._insert_short_at_start(self.flags)
@@ -459,9 +467,13 @@ def _packets(self) -> List[bytes]:
             self.packets_data.append(b''.join(self.data))
             self._reset_for_next_packet()
 
-            if (questions_written + answers_written + authorities_written + additionals_written) == 0 and (
-                len(self.questions) + len(self.answers) + len(self.authorities) + len(self.additionals)
-            ) > 0:
+            if (
+                not questions_written
+                and not answers_written
+                and not authorities_written
+                and not additionals_written
+                and (self.questions or self.answers or self.authorities or self.additionals)
+            ):
                 log.warning("packets() made no progress adding records; returning")
                 break
         self.state = State.finished
