fix: bound _expire_heap growth from re-adds and DRY test helpers

bdraco · bdraco · commit 7c6be63a9fc8 · 2026-05-17T20:48:14.000-07:00
Address remaining review feedback on #1718. **Re-add heap growth (Copilot inline at L100, real bug separate from the orphan).** ``_async_add``'s cap check is gated on ``is_new``, but a peer that just replays cached records with a different ``created``/ TTL hits the ``is_new == False`` path: no cap check, no ``_total_records`` change — yet the changed ``when`` makes ``_expirations.get(record) != when`` true, so a fresh ``(when, record)`` tuple is heappushed while the prior tuple stays behind as stale. ``_expire_heap`` grows unbounded between cleanups even though ``cache`` and ``_total_records`` stay flat. At ~10k pps for 10 s that's ~1 M stale heap entries × ~100 B ≈ 100 MB transient. Factor the rebuild check (``len(heap) > _MIN_SCHEDULED_RECORD_EXPIRATION and len(heap) > 2 * len(expirations)`` → filter-and-heapify) into a new ``_maybe_rebuild_heap`` cdef helper. Call it from three sites: - ``_async_add`` immediately after the conditional heappush, so a re-add flood is bounded at every step (amortized O(1) per push, O(N) rebuild fires at most once per N stale pushes) - ``async_expire`` after the pop loop, replacing the existing inline rebuild — same behavior, slightly more aggressive (uses post-pop heap length instead of pre-pop), no test regression - (removed from ``_async_evict_oldest`` — the eager rebuild in ``_async_add`` keeps the heap bounded before eviction is ever triggered, so the eviction-side call was dead code) Cython codegen: the rebuild call from ``_async_add`` only fires when the conditional heappush already fired, so the hot path on duplicate re-adds (same ``when``) is unchanged. The check inside ``_maybe_rebuild_heap`` is two ``len()`` calls + an integer compare — score-0 in ``cython -a``. **Test DRY-up.** Introduce a module-level ``_addr(name, idx, ...)`` helper that builds the boilerplate ``DNSAddress(name, _TYPE_A, _CLASS_IN, ttl, bytes((idx & 0xFF, (idx >> 8) & 0xFF, 0, 1)), created =...)`` so each new test reads top-down instead of being dominated by constructor noise. Test docstrings collapsed to one sentence per ``CLAUDE.md``. **New regression test:** ``test_cache_re_add_flood_does_not_grow_heap_unbounded`` adds 200 records below the cap, then replays them through 10 cycles of varying TTLs (2000 stale pushes total) and asserts the heap stays near the rebuild threshold. The previous code would have grown it to ~11x the record count; the helper keeps it within ``2 * expirations``. **Removed test:** ``test_cache_eviction_rebuilds_heap_when_mostly_stale`` — its premise (build up stale entries, then trigger eviction-side rebuild) is no longer reachable now that ``_async_add`` rebuilds proactively. The same code path is covered by the re-add flood test. **Existing tests updated:** ``test_cache_heap_multi_name_cleanup`` and ``test_cache_heap_pops_order`` previously asserted pre-cleanup heap size (``min_records_to_cleanup + 5``). With proactive rebuild firing inside ``_async_add``, the heap is cleaned up earlier and the post- cleanup invariant (``1 + 5`` after ``async_expire``) still holds, but the intermediate snapshot is no longer accurate. Replaced the stale intermediate assertion with a comment.
diff --git a/src/zeroconf/_cache.pxd b/src/zeroconf/_cache.pxd
@@ -67,9 +67,12 @@ cdef class DNSCache:
     )
     cdef bint _async_add(self, DNSRecord record)
 
-    @cython.locals(record=DNSRecord, when_record=tuple, expire_heap_len="unsigned int")
+    @cython.locals(record=DNSRecord, when_record=tuple)
     cdef void _async_evict_oldest(self)
 
+    @cython.locals(expire_heap_len="unsigned int")
+    cdef void _maybe_rebuild_heap(self)
+
     @cython.locals(service_record=DNSService)
     cdef void _async_remove(self, DNSRecord record)
 
diff --git a/src/zeroconf/_cache.py b/src/zeroconf/_cache.py
@@ -110,9 +110,14 @@ def _async_add(self, record: _DNSRecord) -> bool:
         store[record] = record
         when = record.created + (record.ttl * 1000)
         if self._expirations.get(record) != when:
-            # Avoid adding duplicates to the heap
             heappush(self._expire_heap, (when, record))
             self._expirations[record] = when
+            # Re-adds of an existing record with a new TTL push a fresh
+            # entry but leave the prior tuple behind as stale, so a peer
+            # that just replays cached records can grow ``_expire_heap``
+            # without ever tripping the cap. Rebuild when stale entries
+            # dominate.
+            self._maybe_rebuild_heap()
 
         if isinstance(record, DNSService):
             service_record = record
@@ -122,16 +127,28 @@ def _async_add(self, record: _DNSRecord) -> bool:
         return new
 
     def _async_evict_oldest(self) -> None:
-        """Drop the closest-to-expiration record to make room for a new one.
-
-        Skips stale heap entries (records re-added with a different TTL),
-        which mirrors the staleness check in ``async_expire``. If the
-        heap is mostly stale (long stale prefix from sustained TTL
-        re-adds), rebuild it once up front so the pop loop below
-        doesn't do O(stale_prefix * log n) work on a single add. Same
-        heuristic / threshold as ``async_expire``.
+        """Drop the closest-to-expiration record to make room for a new one."""
+        while self._expire_heap:
+            when_record = heappop(self._expire_heap)
+            record = when_record[1]
+            if self._expirations.get(record) != when_record[0]:
+                continue
+            self._async_remove(record)
+            return
 
-        This function must be run in from event loop.
+    def _maybe_rebuild_heap(self) -> None:
+        """Rebuild ``_expire_heap`` if stale entries dominate.
+
+        Re-adds of an existing record with a new TTL append a fresh
+        ``(when, record)`` and leave the prior tuple behind as stale;
+        eviction's stale-skip loop and ``async_expire`` already absorb
+        these, but unchecked accumulation lets a peer that just replays
+        cached records grow the heap arbitrarily between cleanups.
+        Same threshold as the long-standing rebuild in ``async_expire``:
+        only fire when stale entries outweigh live ones (heap > 2 x
+        expirations), and only above a minimum floor so a small cache
+        isn't rebuilt for nothing. Amortized cost is O(1) per push;
+        the O(N) rebuild fires at most once per N stale pushes.
         """
         expire_heap_len = len(self._expire_heap)
         if (
@@ -142,13 +159,6 @@ def _async_evict_oldest(self) -> None:
                 entry for entry in self._expire_heap if self._expirations.get(entry[1]) == entry[0]
             ]
             heapify(self._expire_heap)
-        while self._expire_heap:
-            when_record = heappop(self._expire_heap)
-            record = when_record[1]
-            if self._expirations.get(record) != when_record[0]:
-                continue
-            self._async_remove(record)
-            return
 
     def async_add_records(self, entries: Iterable[DNSRecord]) -> bool:
         """Add multiple records.
@@ -190,43 +200,23 @@ def async_expire(self, now: _float) -> list[DNSRecord]:
 
         :param now: The current time in milliseconds.
         """
-        if not (expire_heap_len := len(self._expire_heap)):
+        if not self._expire_heap:
             return []
 
         expired: list[DNSRecord] = []
-        # Find any expired records and add them to the to-delete list
         while self._expire_heap:
             when_record = self._expire_heap[0]
             when = when_record[0]
             if when > now:
                 break
             heappop(self._expire_heap)
-            # Check if the record hasn't been re-added to the heap
-            # with a different expiration time as it will be removed
-            # later when it reaches the top of the heap and its
-            # expiration time is met.
+            # Skip entries left behind by a TTL re-add; the live tuple is
+            # later in the heap and will be removed when it reaches the top.
             record = when_record[1]
             if self._expirations.get(record) == when:
                 expired.append(record)
 
-        # If the expiration heap grows larger than the number expirations
-        # times two, we clean it up to avoid keeping expired entries in
-        # the heap and consuming memory. We guard this with a minimum
-        # threshold to avoid cleaning up the heap too often when there are
-        # only a few scheduled expirations.
-        if (
-            expire_heap_len > _MIN_SCHEDULED_RECORD_EXPIRATION
-            and expire_heap_len > len(self._expirations) * 2
-        ):
-            # Remove any expired entries from the expiration heap
-            # that do not match the expiration time in the expirations
-            # as it means the record has been re-added to the heap
-            # with a different expiration time.
-            self._expire_heap = [
-                entry for entry in self._expire_heap if self._expirations.get(entry[1]) == entry[0]
-            ]
-            heapify(self._expire_heap)
-
+        self._maybe_rebuild_heap()
         self.async_remove_records(expired)
         return expired
 
diff --git a/tests/test_cache.py b/tests/test_cache.py
