mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-24 05:26:20 +00:00
fix(events): keep JSONL lock generation stable across deletion (#5455)
* fix(events): preserve JSONL lock generation across deletion * test(events): cover JSONL lock generation after delete * chore(events): clarify JSONL delete lock lifecycle * style(events): restore JSONL trailing newline
This commit is contained in:
parent
f7f4a022e6
commit
fc4e0c32ba
@ -30,6 +30,7 @@ import asyncio
|
|||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
import re
|
import re
|
||||||
|
import weakref
|
||||||
from collections.abc import Callable, Coroutine
|
from collections.abc import Callable, Coroutine
|
||||||
from datetime import UTC, datetime
|
from datetime import UTC, datetime
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@ -49,11 +50,16 @@ class JsonlRunEventStore(RunEventStore):
|
|||||||
def __init__(self, base_dir: str | Path | None = None):
|
def __init__(self, base_dir: str | Path | None = None):
|
||||||
self._base_dir = Path(base_dir) if base_dir else Path(".deer-flow")
|
self._base_dir = Path(base_dir) if base_dir else Path(".deer-flow")
|
||||||
self._seq_counters: dict[str, int] = {} # thread_id -> current max seq
|
self._seq_counters: dict[str, int] = {} # thread_id -> current max seq
|
||||||
# Per-thread asyncio.Lock — serialises concurrent writes within one process.
|
# Weak ownership avoids leaking one lock per historical thread without
|
||||||
self._write_locks: dict[str, asyncio.Lock] = {}
|
# splitting a live lock generation while a holder/waiter still owns it.
|
||||||
|
self._write_locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary()
|
||||||
|
|
||||||
def _get_write_lock(self, thread_id: str) -> asyncio.Lock:
|
def _get_write_lock(self, thread_id: str) -> asyncio.Lock:
|
||||||
return self._write_locks.setdefault(thread_id, asyncio.Lock())
|
lock = self._write_locks.get(thread_id)
|
||||||
|
if lock is None:
|
||||||
|
lock = asyncio.Lock()
|
||||||
|
self._write_locks[thread_id] = lock
|
||||||
|
return lock
|
||||||
|
|
||||||
async def _run_mutation[T](self, thread_id: str, operation: Callable[[], Coroutine[Any, Any, T]]) -> T:
|
async def _run_mutation[T](self, thread_id: str, operation: Callable[[], Coroutine[Any, Any, T]]) -> T:
|
||||||
"""Drain an admitted mutation before propagating caller cancellation.
|
"""Drain an admitted mutation before propagating caller cancellation.
|
||||||
@ -418,11 +424,8 @@ class JsonlRunEventStore(RunEventStore):
|
|||||||
count = len(all_events)
|
count = len(all_events)
|
||||||
await asyncio.to_thread(self._delete_thread_files, thread_id)
|
await asyncio.to_thread(self._delete_thread_files, thread_id)
|
||||||
self._seq_counters.pop(thread_id, None)
|
self._seq_counters.pop(thread_id, None)
|
||||||
# Pop the lock inside the held mutation to minimise the window where a new caller
|
# Mutations already queued on this lock resume after deletion; with
|
||||||
# could obtain a fresh lock while a waiting coroutine still holds the old one.
|
# files and the counter cleared, they recreate the thread at seq 1.
|
||||||
# Note: coroutines that already acquired a reference to this lock before the
|
|
||||||
# delete will still proceed after we release — this is an accepted narrow race.
|
|
||||||
self._write_locks.pop(thread_id, None)
|
|
||||||
return count
|
return count
|
||||||
|
|
||||||
return await self._run_mutation(thread_id, mutate)
|
return await self._run_mutation(thread_id, mutate)
|
||||||
|
|||||||
106
backend/tests/test_jsonl_event_store_lock_lifecycle.py
Normal file
106
backend/tests/test_jsonl_event_store_lock_lifecycle.py
Normal file
@ -0,0 +1,106 @@
|
|||||||
|
"""Lock-registry lifecycle regressions for the JSONL run-event store."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import threading
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from deerflow.runtime.events.store.jsonl import JsonlRunEventStore
|
||||||
|
|
||||||
|
|
||||||
|
def _event(content: str = "message") -> dict:
|
||||||
|
return {
|
||||||
|
"thread_id": "t1",
|
||||||
|
"run_id": "r1",
|
||||||
|
"event_type": "message",
|
||||||
|
"category": "message",
|
||||||
|
"content": content,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class _PausedDelete:
|
||||||
|
def __init__(self, operation):
|
||||||
|
self.operation = operation
|
||||||
|
self.loop = asyncio.get_running_loop()
|
||||||
|
self.entered = asyncio.Event()
|
||||||
|
self.finished = asyncio.Event()
|
||||||
|
self.release = threading.Event()
|
||||||
|
|
||||||
|
def __call__(self, *args):
|
||||||
|
self.loop.call_soon_threadsafe(self.entered.set)
|
||||||
|
try:
|
||||||
|
if not self.release.wait(10):
|
||||||
|
raise TimeoutError("test did not release paused delete")
|
||||||
|
return self.operation(*args)
|
||||||
|
finally:
|
||||||
|
self.loop.call_soon_threadsafe(self.finished.set)
|
||||||
|
|
||||||
|
|
||||||
|
async def _checkpoint() -> None:
|
||||||
|
await asyncio.sleep(0)
|
||||||
|
await asyncio.sleep(0)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_delete_keeps_one_lock_generation_while_waiter_exists(tmp_path, monkeypatch):
|
||||||
|
store = JsonlRunEventStore(tmp_path)
|
||||||
|
await store.put(**_event("baseline"))
|
||||||
|
|
||||||
|
paused_delete = _PausedDelete(store._delete_thread_files)
|
||||||
|
monkeypatch.setattr(store, "_delete_thread_files", paused_delete)
|
||||||
|
|
||||||
|
original_get_lock = store._get_write_lock
|
||||||
|
second_lookup = asyncio.Event()
|
||||||
|
lookup_count = 0
|
||||||
|
|
||||||
|
def observed_get_lock(thread_id: str):
|
||||||
|
nonlocal lookup_count
|
||||||
|
lock = original_get_lock(thread_id)
|
||||||
|
lookup_count += 1
|
||||||
|
if lookup_count == 2:
|
||||||
|
second_lookup.set()
|
||||||
|
return lock
|
||||||
|
|
||||||
|
monkeypatch.setattr(store, "_get_write_lock", observed_get_lock)
|
||||||
|
|
||||||
|
release_queued = asyncio.Event()
|
||||||
|
queued_entered = asyncio.Event()
|
||||||
|
later_entered = asyncio.Event()
|
||||||
|
deletion = asyncio.create_task(store.delete_by_thread("t1"))
|
||||||
|
queued = None
|
||||||
|
later = None
|
||||||
|
try:
|
||||||
|
await asyncio.wait_for(paused_delete.entered.wait(), 5)
|
||||||
|
|
||||||
|
async def queued_operation():
|
||||||
|
queued_entered.set()
|
||||||
|
await release_queued.wait()
|
||||||
|
|
||||||
|
queued = asyncio.create_task(store._run_mutation("t1", queued_operation))
|
||||||
|
await asyncio.wait_for(second_lookup.wait(), 5)
|
||||||
|
await _checkpoint()
|
||||||
|
|
||||||
|
paused_delete.release.set()
|
||||||
|
assert await deletion == 1
|
||||||
|
await asyncio.wait_for(queued_entered.wait(), 5)
|
||||||
|
|
||||||
|
async def later_operation():
|
||||||
|
later_entered.set()
|
||||||
|
|
||||||
|
later = asyncio.create_task(store._run_mutation("t1", later_operation))
|
||||||
|
await _checkpoint()
|
||||||
|
assert not later_entered.is_set(), "a later mutation acquired a new lock generation while the old waiter still owned the thread"
|
||||||
|
|
||||||
|
release_queued.set()
|
||||||
|
await queued
|
||||||
|
await later
|
||||||
|
assert later_entered.is_set()
|
||||||
|
await _checkpoint()
|
||||||
|
assert "t1" not in store._write_locks
|
||||||
|
finally:
|
||||||
|
paused_delete.release.set()
|
||||||
|
release_queued.set()
|
||||||
|
await asyncio.gather(deletion, *([queued] if queued is not None else []), *([later] if later is not None else []), return_exceptions=True)
|
||||||
|
await asyncio.wait_for(paused_delete.finished.wait(), 5)
|
||||||
Loading…
x
Reference in New Issue
Block a user