Skip to content

Commit 55cf4cc

Browse files
bdracobdraco
authored
feat: speed up outgoing packet writer (#1313)
1 parent 9caeabb commit 55cf4cc

3 files changed

Lines changed: 60 additions & 36 deletions

File tree

‎bench/outgoing.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -158,9 +158,10 @@ def generate_packets() -> DNSOutgoing:
158158

159159

160160
def make_outgoing_message() -> None:
161+
out.packets()
161162
out.state = State.init.value
162163
out.finished = False
163-
out.packets()
164+
out._reset_for_next_packet()
164165

165166

166167
count = 100000

‎src/zeroconf/_protocol/outgoing.pxd‎

Lines changed: 30 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,11 @@ cdef cython.uint _FLAGS_TC
1515
cdef cython.uint _MAX_MSG_ABSOLUTE
1616
cdef cython.uint _MAX_MSG_TYPICAL
1717

18+
1819
cdef bint TYPE_CHECKING
1920

21+
cdef unsigned int SHORT_CACHE_MAX
22+
2023
cdef object PACK_BYTE
2124
cdef object PACK_SHORT
2225
cdef object PACK_LONG
@@ -28,6 +31,7 @@ cdef object LOGGING_IS_ENABLED_FOR
2831
cdef object LOGGING_DEBUG
2932

3033
cdef cython.tuple BYTE_TABLE
34+
cdef cython.tuple SHORT_LOOKUP
3135

3236
cdef class DNSOutgoing:
3337

@@ -46,13 +50,15 @@ cdef class DNSOutgoing:
4650
cdef public cython.list authorities
4751
cdef public cython.list additionals
4852

49-
cdef _reset_for_next_packet(self)
53+
cpdef _reset_for_next_packet(self)
5054

51-
cdef _write_byte(self, object value)
55+
cdef _write_byte(self, cython.uint value)
5256

53-
cdef _insert_short_at_start(self, object value)
57+
cdef void _insert_short_at_start(self, unsigned int value)
5458

55-
cdef _replace_short(self, object index, object value)
59+
cdef _replace_short(self, cython.uint index, cython.uint value)
60+
61+
cdef _get_short(self, cython.uint value)
5662

5763
cdef _write_int(self, object value)
5864

@@ -61,24 +67,29 @@ cdef class DNSOutgoing:
6167
@cython.locals(
6268
d=cython.bytes,
6369
data_view=cython.list,
70+
index=cython.uint,
6471
length=cython.uint
6572
)
6673
cdef cython.bint _write_record(self, DNSRecord record, object now)
6774

75+
@cython.locals(class_=cython.uint)
6876
cdef _write_record_class(self, DNSEntry record)
6977

7078
@cython.locals(
7179
start_size_int=object
7280
)
7381
cdef cython.bint _check_data_limit_or_rollback(self, cython.uint start_data_length, cython.uint start_size)
7482

75-
cdef _write_questions_from_offset(self, object questions_offset)
83+
@cython.locals(questions_written=cython.uint)
84+
cdef cython.uint _write_questions_from_offset(self, unsigned int questions_offset)
7685

77-
cdef _write_answers_from_offset(self, object answer_offset)
86+
@cython.locals(answers_written=cython.uint)
87+
cdef cython.uint _write_answers_from_offset(self, unsigned int answer_offset)
7888

79-
cdef _write_records_from_offset(self, cython.list records, object offset)
89+
@cython.locals(records_written=cython.uint)
90+
cdef cython.uint _write_records_from_offset(self, cython.list records, unsigned int offset)
8091

81-
cdef _has_more_to_add(self, object questions_offset, object answer_offset, object authority_offset, object additional_offset)
92+
cdef bint _has_more_to_add(self, unsigned int questions_offset, unsigned int answer_offset, unsigned int authority_offset, unsigned int additional_offset)
8293

8394
cdef _write_ttl(self, DNSRecord record, object now)
8495

@@ -93,23 +104,25 @@ cdef class DNSOutgoing:
93104

94105
cdef _write_link_to_name(self, unsigned int index)
95106

96-
cpdef write_short(self, object value)
107+
cpdef write_short(self, cython.uint value)
97108

98109
cpdef write_string(self, cython.bytes value)
99110

111+
@cython.locals(utfstr=bytes)
100112
cpdef _write_utf(self, cython.str value)
101113

102114
@cython.locals(
103115
debug_enable=bint,
104116
made_progress=bint,
105-
questions_offset=object,
106-
answer_offset=object,
107-
authority_offset=object,
108-
additional_offset=object,
109-
questions_written=object,
110-
answers_written=object,
111-
authorities_written=object,
112-
additionals_written=object,
117+
has_more_to_add=bint,
118+
questions_offset="unsigned int",
119+
answer_offset="unsigned int",
120+
authority_offset="unsigned int",
121+
additional_offset="unsigned int",
122+
questions_written="unsigned int",
123+
answers_written="unsigned int",
124+
authorities_written="unsigned int",
125+
additionals_written="unsigned int",
113126
)
114127
cpdef packets(self)
115128

‎src/zeroconf/_protocol/outgoing.py‎

Lines changed: 28 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,10 @@
5353
PACK_SHORT = Struct('>H').pack
5454
PACK_LONG = Struct('>L').pack
5555

56+
SHORT_CACHE_MAX = 128
57+
5658
BYTE_TABLE = tuple(PACK_BYTE(i) for i in range(256))
59+
SHORT_LOOKUP = tuple(PACK_SHORT(i) for i in range(SHORT_CACHE_MAX))
5760

5861

5962
class State(enum.Enum):
@@ -220,17 +223,21 @@ def _write_byte(self, value: int_) -> None:
220223
self.data.append(BYTE_TABLE[value])
221224
self.size += 1
222225

226+
def _get_short(self, value: int_) -> bytes:
227+
"""Convert an unsigned short to 2 bytes."""
228+
return SHORT_LOOKUP[value] if value < SHORT_CACHE_MAX else PACK_SHORT(value)
229+
223230
def _insert_short_at_start(self, value: int_) -> None:
224231
"""Inserts an unsigned short at the start of the packet"""
225-
self.data.insert(0, PACK_SHORT(value))
232+
self.data.insert(0, self._get_short(value))
226233

227234
def _replace_short(self, index: int_, value: int_) -> None:
228235
"""Replaces an unsigned short in a certain position in the packet"""
229-
self.data[index] = PACK_SHORT(value)
236+
self.data[index] = self._get_short(value)
230237

231238
def write_short(self, value: int_) -> None:
232239
"""Writes an unsigned short to the packet"""
233-
self.data.append(PACK_SHORT(value))
240+
self.data.append(self._get_short(value))
234241
self.size += 2
235242

236243
def _write_int(self, value: Union[float, int]) -> None:
@@ -323,10 +330,11 @@ def _write_question(self, question: DNSQuestion_) -> bool:
323330

324331
def _write_record_class(self, record: Union[DNSQuestion_, DNSRecord_]) -> None:
325332
"""Write out the record class including the unique/unicast (QU) bit."""
326-
if record.unique and self.multicast:
327-
self.write_short(record.class_ | _CLASS_UNIQUE)
333+
class_ = record.class_
334+
if record.unique is True and self.multicast is True:
335+
self.write_short(class_ | _CLASS_UNIQUE)
328336
else:
329-
self.write_short(record.class_)
337+
self.write_short(class_)
330338

331339
def _write_ttl(self, record: DNSRecord_, now: float_) -> None:
332340
"""Write out the record ttl."""
@@ -417,21 +425,20 @@ def packets(self) -> List[bytes]:
417425
will be written out to a single oversized packet no more than
418426
_MAX_MSG_ABSOLUTE in length (and hence will be subject to IP
419427
fragmentation potentially)."""
428+
packets_data = self.packets_data
429+
420430
if self.state == STATE_FINISHED:
421-
return self.packets_data
431+
return packets_data
422432

423433
questions_offset = 0
424434
answer_offset = 0
425435
authority_offset = 0
426436
additional_offset = 0
427437
# we have to at least write out the question
428-
first_time = True
429-
debug_enable = LOGGING_IS_ENABLED_FOR(LOGGING_DEBUG)
438+
debug_enable = LOGGING_IS_ENABLED_FOR(LOGGING_DEBUG) is True
439+
has_more_to_add = True
430440

431-
while first_time or self._has_more_to_add(
432-
questions_offset, answer_offset, authority_offset, additional_offset
433-
):
434-
first_time = False
441+
while has_more_to_add:
435442
if debug_enable:
436443
log.debug(
437444
"offsets = questions=%d, answers=%d, authorities=%d, additionals=%d",
@@ -473,9 +480,11 @@ def packets(self) -> List[bytes]:
473480
additional_offset,
474481
)
475482

476-
if self.is_query() and self._has_more_to_add(
483+
has_more_to_add = self._has_more_to_add(
477484
questions_offset, answer_offset, authority_offset, additional_offset
478-
):
485+
)
486+
487+
if has_more_to_add and self.is_query():
479488
# https://datatracker.ietf.org/doc/html/rfc6762#section-7.2
480489
if debug_enable: # pragma: no branch
481490
log.debug("Setting TC flag")
@@ -488,7 +497,7 @@ def packets(self) -> List[bytes]:
488497
else:
489498
self._insert_short_at_start(self.id)
490499

491-
self.packets_data.append(b''.join(self.data))
500+
packets_data.append(b''.join(self.data))
492501

493502
if not made_progress:
494503
# Generating an empty packet is not a desirable outcome, but currently
@@ -498,7 +507,8 @@ def packets(self) -> List[bytes]:
498507
log.warning("packets() made no progress adding records; returning")
499508
break
500509

501-
self._reset_for_next_packet()
510+
if has_more_to_add:
511+
self._reset_for_next_packet()
502512

503513
self.state = STATE_FINISHED
504-
return self.packets_data
514+
return packets_data

0 commit comments

Comments
 (0)