feat: speed up writing name compression for outgoing packets (#1312)

bdraco · web-flow · commit 9caeabb6d465 · 2023-11-13T12:16:15.000-06:00
diff --git a/bench/outgoing.py b/bench/outgoing.py
@@ -158,7 +158,7 @@ def generate_packets() -> DNSOutgoing:
 
 
 def make_outgoing_message() -> None:
-    out.state = State.init
+    out.state = State.init.value
     out.finished = False
     out.packets()
 
diff --git a/src/zeroconf/_protocol/outgoing.pxd b/src/zeroconf/_protocol/outgoing.pxd
@@ -21,8 +21,8 @@ cdef object PACK_BYTE
 cdef object PACK_SHORT
 cdef object PACK_LONG
 
-cdef object STATE_INIT
-cdef object STATE_FINISHED
+cdef unsigned int STATE_INIT
+cdef unsigned int STATE_FINISHED
 
 cdef object LOGGING_IS_ENABLED_FOR
 cdef object LOGGING_DEBUG
@@ -40,7 +40,7 @@ cdef class DNSOutgoing:
     cdef public cython.list data
     cdef public unsigned int size
     cdef public bint allow_long
-    cdef public object state
+    cdef public unsigned int state
     cdef public cython.list questions
     cdef public cython.list answers
     cdef public cython.list authorities
@@ -91,13 +91,17 @@ cdef class DNSOutgoing:
     )
     cpdef write_name(self, cython.str name)
 
+    cdef _write_link_to_name(self, unsigned int index)
+
     cpdef write_short(self, object value)
 
     cpdef write_string(self, cython.bytes value)
 
     cpdef _write_utf(self, cython.str value)
 
     @cython.locals(
+        debug_enable=bint,
+        made_progress=bint,
         questions_offset=object,
         answer_offset=object,
         authority_offset=object,
@@ -107,7 +111,7 @@ cdef class DNSOutgoing:
         authorities_written=object,
         additionals_written=object,
     )
-    cdef _packets(self)
+    cpdef packets(self)
 
     cpdef add_question_or_all_cache(self, DNSCache cache, object now, str name, object type_, object class_)
 
@@ -124,6 +128,6 @@ cdef class DNSOutgoing:
 
     cpdef add_additional_answer(self, DNSRecord record)
 
-    cpdef is_query(self)
+    cpdef bint is_query(self)
 
-    cpdef is_response(self)
+    cpdef bint is_response(self)
diff --git a/src/zeroconf/_protocol/outgoing.py b/src/zeroconf/_protocol/outgoing.py
@@ -61,8 +61,8 @@ class State(enum.Enum):
     finished = 1
 
 
-STATE_INIT = State.init
-STATE_FINISHED = State.finished
+STATE_INIT = State.init.value
+STATE_FINISHED = State.finished.value
 
 LOGGING_IS_ENABLED_FOR = log.isEnabledFor
 LOGGING_DEBUG = logging.DEBUG
@@ -277,30 +277,41 @@ def write_name(self, name: str_) -> None:
         """
 
         # split name into each label
-        name_length = 0
         if name.endswith('.'):
-            name = name[: len(name) - 1]
-        labels = name.split('.')
-        # Write each new label or a pointer to the existing
-        # on in the packet
+            name = name[:-1]
+
+        index = self.names.get(name, 0)
+        if index:
+            self._write_link_to_name(index)
+            return
+
         start_size = self.size
-        for count in range(len(labels)):
-            label = name if count == 0 else '.'.join(labels[count:])
-            index = self.names.get(label, 0)
+        labels = name.split('.')
+        # Write each new label or a pointer to the existing one in the packet
+        self.names[name] = start_size
+        self._write_utf(labels[0])
+
+        name_length = 0
+        for count in range(1, len(labels)):
+            partial_name = '.'.join(labels[count:])
+            index = self.names.get(partial_name, 0)
             if index:
-                # If part of the name already exists in the packet,
-                # create a pointer to it
-                self._write_byte((index >> 8) | 0xC0)
-                self._write_byte(index & 0xFF)
+                self._write_link_to_name(index)
                 return
             if name_length == 0:
                 name_length = len(name.encode('utf-8'))
-            self.names[label] = start_size + name_length - len(label.encode('utf-8'))
+            self.names[partial_name] = start_size + name_length - len(partial_name.encode('utf-8'))
             self._write_utf(labels[count])
 
         # this is the end of a name
         self._write_byte(0)
 
+    def _write_link_to_name(self, index: int_) -> None:
+        # If part of the name already exists in the packet,
+        # create a pointer to it
+        self._write_byte((index >> 8) | 0xC0)
+        self._write_byte(index & 0xFF)
+
     def _write_question(self, question: DNSQuestion_) -> bool:
         """Writes a question to the packet"""
         start_data_length = len(self.data)
@@ -406,9 +417,6 @@ def packets(self) -> List[bytes]:
         will be written out to a single oversized packet no more than
         _MAX_MSG_ABSOLUTE in length (and hence will be subject to IP
         fragmentation potentially)."""
-        return self._packets()
-
-    def _packets(self) -> List[bytes]:
         if self.state == STATE_FINISHED:
             return self.packets_data
 
@@ -445,6 +453,8 @@ def _packets(self) -> List[bytes]:
             authorities_written = self._write_records_from_offset(self.authorities, authority_offset)
             additionals_written = self._write_records_from_offset(self.additionals, additional_offset)
 
+            made_progress = bool(self.data)
+
             self._insert_short_at_start(additionals_written)
             self._insert_short_at_start(authorities_written)
             self._insert_short_at_start(answers_written)
@@ -479,16 +489,16 @@ def _packets(self) -> List[bytes]:
                 self._insert_short_at_start(self.id)
 
             self.packets_data.append(b''.join(self.data))
-            self._reset_for_next_packet()
 
-            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)
-            ):
+            if not made_progress:
+                # Generating an empty packet is not a desirable outcome, but currently
+                # too many internals rely on this behavior.  So, we'll just return an
+                # empty packet and log a warning until this can be refactored at a later
+                # date.
                 log.warning("packets() made no progress adding records; returning")
                 break
+
+            self._reset_for_next_packet()
+
         self.state = STATE_FINISHED
         return self.packets_data
