From bf2f3660a1f341e50ab0ae586dfbacbc5ddcc077 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Thu, 7 Sep 2023 14:55:52 -0500 Subject: [PATCH] feat: significantly speed up writing outgoing dns records (#1260) --- src/zeroconf/_protocol/outgoing.pxd | 20 +++++++++++++---- src/zeroconf/_protocol/outgoing.py | 34 +++++++++++++++++++---------- 2 files changed, 39 insertions(+), 15 deletions(-) diff --git a/src/zeroconf/_protocol/outgoing.pxd b/src/zeroconf/_protocol/outgoing.pxd index 0338cfd8..4caaf453 100644 --- a/src/zeroconf/_protocol/outgoing.pxd +++ b/src/zeroconf/_protocol/outgoing.pxd @@ -17,6 +17,10 @@ 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 @@ -24,7 +28,7 @@ cdef class DNSOutgoing: 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,7 +57,7 @@ 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) @@ -61,16 +65,24 @@ cdef class DNSOutgoing: 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 index 069b2936..4d17cce0 100644 --- 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