Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: significantly speed up writing outgoing dns records #1260

Merged
merged 8 commits into from
Sep 7, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
20 changes: 16 additions & 4 deletions src/zeroconf/_protocol/outgoing.pxd
Original file line number Diff line number Diff line change
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
Original file line number Diff line number Diff line change
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