diff --git a/src/kvcr/local_dram.py b/src/kvcr/local_dram.py index e6ba87e..3c120dd 100644 --- a/src/kvcr/local_dram.py +++ b/src/kvcr/local_dram.py @@ -1032,8 +1032,9 @@ def short() -> set[str]: if len(self._free_slots[name]) + freed[name] < count } + deficient = short() with closing(self._evictable.candidates(protected)) as candidates: - while deficient := short(): + while deficient: key = next(candidates, None) if key is None: return None, [], False @@ -1054,7 +1055,8 @@ def short() -> set[str]: CacheTier.LOCAL_G2, deadline, ) - if not short(): + deficient = short() + if not deficient: break if eviction_pending: self._capacity_eviction_key = key @@ -1063,6 +1065,7 @@ def short() -> set[str]: continue victims.append((key, record, residency, size_bytes)) freed.update(name for name, _ in residency.slots) + deficient = short() for key, record, residency, size_bytes in victims: self._remove_evictable(key, residency) diff --git a/src/kvcr/policy_runtime.py b/src/kvcr/policy_runtime.py index 5cf70f6..8fe9c5c 100644 --- a/src/kvcr/policy_runtime.py +++ b/src/kvcr/policy_runtime.py @@ -151,9 +151,6 @@ class _Entry: class _EvictionQueue: - # TODO: Bound stale heap growth from DRAM/G3 claim/release cycles. - # Removal only invalidates _live; candidates() removes stale heap entries as - # it encounters them, so repeated cache use can grow _heap without eviction. def __init__(self) -> None: self._heap: list[tuple[float, int, BlockKey]] = [] self._live: dict[BlockKey, _Entry] = {} @@ -167,6 +164,18 @@ def insert(self, key: BlockKey, score: float) -> None: self._next_sequence += 1 self._live[key] = entry heapq.heappush(self._heap, (entry.score, entry.sequence, key)) + if len(self._heap) > 2 * len(self._live): + # Accumulated stale entries amortize this O(n) pass. Reusing tuples + # limits overhead, but this insert still pays the synchronous cost. + # If profiling shows significant pauses, use incremental compaction. + # Compact queued entries only; candidates() restores those it holds. + self._heap = [ + item + for item in self._heap + if (current := self._live.get(item[2])) is not None + and current.sequence == item[1] + ] + heapq.heapify(self._heap) def remove(self, key: BlockKey) -> bool: return self._live.pop(key, None) is not None diff --git a/tests/unit/test_kvcr_local_dram.py b/tests/unit/test_kvcr_local_dram.py index 423920b..d9b6abe 100644 --- a/tests/unit/test_kvcr_local_dram.py +++ b/tests/unit/test_kvcr_local_dram.py @@ -5,6 +5,7 @@ import ctypes import heapq import logging +from contextlib import closing from unittest.mock import Mock import pytest @@ -25,6 +26,7 @@ from kvcr.core import _BlockRecord from kvcr.local_dram import _LocalDramResidency, _LocalDramState from kvcr.policy import FIFOPolicy, LRUPolicy +from kvcr.policy_runtime import _EvictionQueue from kvcr.recovery_journal import ( RecoveryMirrorError, install_recovery_records, @@ -286,6 +288,27 @@ def test_group_allocation_evicts_enough_whole_keys(monkeypatch) -> None: ] +def test_eviction_heap_compaction_preserves_active_candidates() -> None: + queue = _EvictionQueue() + first, churned, excluded = ( + BlockKey(name) for name in (b"first", b"churned", b"excluded") + ) + queue.insert(first, 0) + queue.insert(churned, 1) + queue.insert(excluded, 1) + + with closing(queue.candidates({excluded})) as candidates: + assert next(candidates) == first + for _ in range(100): + queue.remove(churned) + queue.insert(churned, 1) + assert len(queue._heap) <= 2 * len(queue) + assert list(candidates) == [churned] + + assert len(queue._heap) == 3 + assert list(queue.candidates(set())) == [first, excluded, churned] + + @pytest.mark.parametrize( ("policy", "evicted_index"), [(FIFOPolicy(), 1), (LRUPolicy(), 0), (None, 0)],