Skip to content

Commit

Permalink
feat: significantly speed up writing outgoing dns records (#1260)
Browse files Browse the repository at this point in the history
  • Loading branch information
bdraco committed Sep 7, 2023
1 parent 72bd07d commit bf2f366
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 15 deletions.
20 changes: 16 additions & 4 deletions src/zeroconf/_protocol/outgoing.pxd
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
34 changes: 23 additions & 11 deletions src/zeroconf/_protocol/outgoing.py
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down

0 comments on commit bf2f366

Please sign in to comment.