mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 05:06:59 +00:00
e6a3c9ffb7
* fix(checkpoints): serialize shared memory history cache operations * fix(checkpoints): align memory cache contract with locking * test(checkpoints): make cache race regression deterministic * docs: remove overlapping cache note from memory-cache PR
126 lines
4.3 KiB
Python
126 lines
4.3 KiB
Python
"""Concurrent sync readers share the provider's process-local history cache."""
|
|
|
|
import asyncio
|
|
import threading
|
|
from collections import OrderedDict
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from queue import Queue
|
|
|
|
import pytest
|
|
|
|
from deerflow.runtime.checkpoint_cache.base import make_history_key
|
|
from deerflow.runtime.checkpoint_cache.memory import MemoryCheckpointHistoryCache
|
|
|
|
|
|
@pytest.mark.parametrize("mutation", ["evict", "purge", "close"])
|
|
def test_lookup_survives_concurrent_removal(mutation):
|
|
cache = MemoryCheckpointHistoryCache(max_entries=1)
|
|
key = make_history_key("p", "t1", "", "c1", "messages")
|
|
entry = {"writes": [("task", "messages", "original")], "seed": "seed"}
|
|
cache.set_many({key: entry})
|
|
looked_up = threading.Event()
|
|
release_reader = threading.Event()
|
|
writer_started = threading.Event()
|
|
writer_finished = threading.Event()
|
|
coordination = Queue()
|
|
removal_reported = threading.Event()
|
|
|
|
def report_removal():
|
|
if not removal_reported.is_set():
|
|
removal_reported.set()
|
|
coordination.put("removed")
|
|
|
|
class PausedLookup(OrderedDict):
|
|
def get(self, lookup_key, default=None):
|
|
value = super().get(lookup_key, default)
|
|
if lookup_key == key:
|
|
looked_up.set()
|
|
assert release_reader.wait(5), "reader was not released"
|
|
return value
|
|
|
|
def __delitem__(self, deleted_key):
|
|
super().__delitem__(deleted_key)
|
|
if deleted_key == key:
|
|
report_removal()
|
|
|
|
def popitem(self, last=True):
|
|
removed_key, value = super().popitem(last=last)
|
|
if removed_key == key:
|
|
report_removal()
|
|
return removed_key, value
|
|
|
|
def clear(self):
|
|
contained_target = key in self
|
|
super().clear()
|
|
if contained_target:
|
|
report_removal()
|
|
|
|
class ObservedLock:
|
|
"""Report a contender before it blocks behind the paused reader."""
|
|
|
|
def __init__(self):
|
|
self._lock = threading.Lock()
|
|
|
|
def __enter__(self):
|
|
if self._lock.locked():
|
|
coordination.put("writer_waiting")
|
|
self._lock.acquire()
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback):
|
|
self._lock.release()
|
|
|
|
# Pause at the actual lookup -> LRU promotion window, not before get_many.
|
|
cache._data = PausedLookup(cache._data)
|
|
cache._lock = ObservedLock()
|
|
|
|
def remove():
|
|
writer_started.set()
|
|
try:
|
|
if mutation == "evict":
|
|
cache.set_many({"other": {"writes": []}})
|
|
elif mutation == "purge":
|
|
cache.delete_thread("p", "t1")
|
|
else:
|
|
asyncio.run(cache.aclose())
|
|
finally:
|
|
writer_finished.set()
|
|
|
|
with ThreadPoolExecutor(max_workers=2) as pool:
|
|
reader = pool.submit(cache.get_many, [key])
|
|
try:
|
|
assert looked_up.wait(5)
|
|
writer = pool.submit(remove)
|
|
assert writer_started.wait(5)
|
|
# A serialized writer reports that it is waiting on the reader's
|
|
# lock. On unfixed main it instead removes the key and reports that,
|
|
# deterministically reproducing the get -> move_to_end race.
|
|
assert coordination.get(timeout=5) == "writer_waiting"
|
|
finally:
|
|
release_reader.set()
|
|
assert reader.result(timeout=5) == {key: entry}
|
|
writer.result(timeout=5)
|
|
assert cache.get_many([key]) == {}
|
|
assert cache.stats().entries == int(mutation == "evict")
|
|
|
|
|
|
def test_parallel_read_write_keeps_capacity_and_accounting():
|
|
cache = MemoryCheckpointHistoryCache(max_entries=4)
|
|
start = threading.Barrier(8)
|
|
|
|
def exercise(worker):
|
|
start.wait(timeout=5)
|
|
for n in range(200):
|
|
key = f"{worker}:{n}"
|
|
cache.set_many({key: {"writes": [("task", "messages", key)]}})
|
|
result = cache.get_many([key])
|
|
if key in result:
|
|
assert result[key]["writes"] == [("task", "messages", key)]
|
|
|
|
with ThreadPoolExecutor(max_workers=8) as pool:
|
|
list(pool.map(exercise, range(8)))
|
|
stats = cache.stats()
|
|
assert stats.entries == 4
|
|
assert stats.evictions == 8 * 200 - 4
|
|
assert stats.hits + stats.misses == 8 * 200
|