mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-11 14:38:38 +00:00
bb75f8d736
* feat(sandbox): share sandbox identity derivation and acquire serialization (#4741) Remote providers (AIO, E2B, BoxLite, Tenki, OpenSandbox) each inlined the same sha256(user:thread)[:16] sandbox-id expression and kept per-scope lock dicts that grew unboundedly until shutdown. This extracts both mechanisms into shared components without changing provider lifecycle, ids, capacity semantics, or public tool behavior: - sandbox/identity.py: keyword-only derive_sandbox_scope_token (byte-pinned compatibility contract) + is_sandbox_scope_token; per-provider golden vectors pin current behavior including BoxLite's raw-None quirk and each provider's private user_id resolution. - sandbox/acquire_serialization.py: AcquireSerializer — per-key lock table with holder/waiter refcount reclamation, bounded dedicated executor (async waits off both the event loop and the default executor), worker-owned cancellation cleanup (no event-loop callback dependency), idempotent close(). - Each provider adopts both components; AIO/E2B key by (user_id, thread_id) with acquire and (E2B) release serialized; BoxLite/Tenki/OpenSandbox key by derived sandbox id and offload the whole sync acquire to the serializer's executor so a cancelled awaiter cannot overlap a retried same-scope body (leaked-remote-VM regression caught in review). - thread_id=None acquires stay unserialized; provider shutdown()/reset() close the serializer; E2B capacity/ledger/reconciliation and AIO ownership/flock machinery untouched. - blocking-IO anchor proves contended OpenSandbox acquire_async stays off the event loop (teeth verified red/green); AGENTS.md documents the shared components. * refactor(sandbox): address review on acquire serialization (#5089) - Replace unreachable checkin branch with an assertion: run() returns False only after abandon(), which the except handler always re-raises; the old _checkin would have double-decremented the refcount. - Document the task.cancelling() == 0 assumption in hold_async. - Drop unused thread_id/user_id kwargs from BoxLite and Tenki _acquire_scope_locked (OpenSandbox still forwards them). * fix(sandbox): preserve request ContextVars in acquire executor bridge (#5089) loop.run_in_executor() does not copy contextvars, unlike the inherited SandboxProvider.acquire_async() which used asyncio.to_thread(). The BoxLite/OpenSandbox/Tenki acquire_async bridges introduced in this PR therefore dropped the request trace id (logged as trace_id=-). Add AcquireSerializer.run_on_executor(), which copies the calling context and runs the callable through ctx.run, and route all three providers through it. Add regression tests binding request_trace_context and verifying the worker thread observes it.
271 lines
9.9 KiB
Python
271 lines
9.9 KiB
Python
"""Contract tests for AcquireSerializer (RFC #4741 §5)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import threading
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
import pytest
|
|
|
|
from deerflow.sandbox.acquire_serialization import AcquireSerializer
|
|
from deerflow.trace_context import get_current_trace_id, request_trace_context
|
|
|
|
|
|
class TestSyncMutualExclusion:
|
|
def test_same_key_never_overlaps(self):
|
|
serializer = AcquireSerializer()
|
|
active = 0
|
|
max_active = 0
|
|
counter_lock = threading.Lock()
|
|
|
|
def worker():
|
|
nonlocal active, max_active
|
|
with serializer.hold("k"):
|
|
with counter_lock:
|
|
active += 1
|
|
max_active = max(max_active, active)
|
|
time.sleep(0.02)
|
|
with counter_lock:
|
|
active -= 1
|
|
|
|
threads = [threading.Thread(target=worker) for _ in range(8)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
assert max_active == 1
|
|
|
|
def test_different_keys_proceed_concurrently(self):
|
|
serializer = AcquireSerializer()
|
|
barrier = threading.Barrier(2, timeout=2)
|
|
|
|
def worker(key):
|
|
with serializer.hold(key):
|
|
barrier.wait() # raises BrokenBarrierError if serialized
|
|
|
|
t1 = threading.Thread(target=worker, args=("a",))
|
|
t2 = threading.Thread(target=worker, args=("b",))
|
|
t1.start()
|
|
t2.start()
|
|
t1.join(timeout=5)
|
|
t2.join(timeout=5)
|
|
assert not t1.is_alive() and not t2.is_alive()
|
|
|
|
def test_lock_table_reclaims_entries(self):
|
|
serializer = AcquireSerializer()
|
|
for i in range(50):
|
|
with serializer.hold(f"scope-{i}"):
|
|
pass
|
|
assert len(serializer._table) == 0
|
|
|
|
|
|
class TestAsyncContract:
|
|
@pytest.mark.asyncio
|
|
async def test_cancellation_before_acquire_leaks_nothing(self):
|
|
serializer = AcquireSerializer()
|
|
with serializer.hold("k"): # uncontended acquire, safe on the loop
|
|
waiter = asyncio.create_task(_hold_once(serializer, "k"))
|
|
await asyncio.sleep(0.05) # let the waiter start blocking in the worker
|
|
waiter.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await waiter
|
|
# The worker acquires the now-free lock in the background; its done
|
|
# callback must release it and reclaim the entry.
|
|
async with serializer.hold_async("k"):
|
|
pass
|
|
await _wait_for_table_empty(serializer)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_wait_releases_without_event_loop_callback(self):
|
|
serializer = AcquireSerializer()
|
|
holder_ready = threading.Event()
|
|
release_holder = threading.Event()
|
|
|
|
def hold_in_thread():
|
|
with serializer.hold("k"):
|
|
holder_ready.set()
|
|
release_holder.wait(2)
|
|
|
|
holder = threading.Thread(target=hold_in_thread)
|
|
holder.start()
|
|
assert holder_ready.wait(2)
|
|
waiter = asyncio.create_task(_hold_once(serializer, "k"))
|
|
await asyncio.sleep(0.05)
|
|
waiter.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await waiter
|
|
|
|
# Once cancellation has been observed, cleanup must be completed by
|
|
# the acquire worker itself. Blocking this loop keeps asyncio Future
|
|
# callbacks from hiding a lock leak here.
|
|
release_holder.set()
|
|
deadline = time.monotonic() + 2
|
|
while serializer._table and time.monotonic() < deadline:
|
|
time.sleep(0.01)
|
|
holder.join(timeout=2)
|
|
assert not holder.is_alive()
|
|
assert len(serializer._table) == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancellation_after_acquire_releases(self):
|
|
serializer = AcquireSerializer()
|
|
|
|
async def run():
|
|
async with serializer.hold_async("k"):
|
|
await asyncio.sleep(10)
|
|
|
|
task = asyncio.create_task(run())
|
|
await asyncio.sleep(0.05)
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
async with serializer.hold_async("k"): # must not hang
|
|
pass
|
|
assert len(serializer._table) == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_wait_keeps_loop_alive_and_off_default_executor(self):
|
|
serializer = AcquireSerializer()
|
|
loop = asyncio.get_running_loop()
|
|
|
|
class ExplodingExecutor(ThreadPoolExecutor):
|
|
def submit(self, *args, **kwargs):
|
|
raise AssertionError("serializer wait used the default executor")
|
|
|
|
previous = loop._default_executor
|
|
loop.set_default_executor(ExplodingExecutor())
|
|
try:
|
|
entered = asyncio.Event()
|
|
|
|
async def holder():
|
|
async with serializer.hold_async("k"):
|
|
entered.set()
|
|
await asyncio.sleep(0.2)
|
|
|
|
h = asyncio.create_task(holder())
|
|
await entered.wait()
|
|
w = asyncio.create_task(_hold_once(serializer, "k"))
|
|
for _ in range(5): # loop stays responsive while w waits
|
|
await asyncio.sleep(0.02)
|
|
assert not w.done()
|
|
await h
|
|
await asyncio.wait_for(w, 2)
|
|
finally:
|
|
if previous is not None: # fresh loops have no default executor yet
|
|
loop.set_default_executor(previous)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_async_waiters_serialize(self):
|
|
serializer = AcquireSerializer()
|
|
active = 0
|
|
max_active = 0
|
|
|
|
async def worker():
|
|
nonlocal active, max_active
|
|
async with serializer.hold_async("k"):
|
|
active += 1
|
|
max_active = max(max_active, active)
|
|
await asyncio.sleep(0.01)
|
|
active -= 1
|
|
|
|
await asyncio.gather(*[worker() for _ in range(8)])
|
|
assert max_active == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_on_executor_preserves_request_trace_context(self):
|
|
"""run_on_executor must carry request ContextVars into the worker.
|
|
|
|
Regression test (#5089): raw ``loop.run_in_executor`` does not copy
|
|
contextvars, so without an explicit ``copy_context`` the worker thread
|
|
reads the trace id bound by ``request_trace_context()`` as unset.
|
|
"""
|
|
serializer = AcquireSerializer()
|
|
try:
|
|
with request_trace_context("trace-5089"):
|
|
seen = await serializer.run_on_executor(get_current_trace_id)
|
|
assert seen == "trace-5089"
|
|
finally:
|
|
serializer.close()
|
|
|
|
|
|
async def _hold_once(serializer, key):
|
|
async with serializer.hold_async(key):
|
|
pass
|
|
|
|
|
|
async def _wait_for_table_empty(serializer, timeout=2.0):
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
if len(serializer._table) == 0:
|
|
return
|
|
await asyncio.sleep(0.01)
|
|
raise AssertionError("serializer lock table did not drain")
|
|
|
|
|
|
class TestClose:
|
|
def test_close_rejects_new_holders_and_is_idempotent(self):
|
|
serializer = AcquireSerializer()
|
|
serializer.close()
|
|
serializer.close()
|
|
with pytest.raises(RuntimeError, match="closed"):
|
|
with serializer.hold("k"):
|
|
pass
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_close_does_not_invalidate_active_critical_section(self):
|
|
serializer = AcquireSerializer()
|
|
async with serializer.hold_async("k"):
|
|
serializer.close() # must not raise, must not break exit
|
|
# New holders rejected afterwards:
|
|
with pytest.raises(RuntimeError, match="closed"):
|
|
async with serializer.hold_async("k"):
|
|
pass
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_executor_shutdown_between_checkout_and_schedule_reclaims_entry(self):
|
|
serializer = AcquireSerializer()
|
|
serializer._executor.shutdown(wait=False, cancel_futures=True) # simulates close() landing mid-window
|
|
with pytest.raises(RuntimeError, match="closed"):
|
|
async with serializer.hold_async("k"):
|
|
pass
|
|
assert len(serializer._table) == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_close_reports_closed_to_queued_waiter(self):
|
|
"""A waiter whose queued acquire is cancelled by close() must see
|
|
RuntimeError("closed"), not a spurious CancelledError (#4741)."""
|
|
serializer = AcquireSerializer(max_workers=1)
|
|
blocker = threading.Event()
|
|
with serializer.hold("k"): # uncontended acquire, safe on the loop
|
|
# Occupy the single worker so the waiter's acquire future stays
|
|
# queued (pending) and is cancelled by close()'s cancel_futures.
|
|
occupied = asyncio.get_running_loop().run_in_executor(serializer.executor, blocker.wait, 5)
|
|
waiter = asyncio.create_task(_hold_once(serializer, "k"))
|
|
await asyncio.sleep(0.1) # let the waiter check out and queue its acquire
|
|
serializer.close()
|
|
with pytest.raises(RuntimeError, match="closed"):
|
|
await waiter
|
|
blocker.set()
|
|
await occupied
|
|
assert len(serializer._table) == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_close_reclaims_caller_cancelled_queued_waiter(self):
|
|
serializer = AcquireSerializer(max_workers=1)
|
|
blocker = threading.Event()
|
|
occupied = asyncio.get_running_loop().run_in_executor(serializer.executor, blocker.wait, 5)
|
|
waiter = asyncio.create_task(_hold_once(serializer, "k"))
|
|
await asyncio.sleep(0.05)
|
|
waiter.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await waiter
|
|
|
|
# The acquire was abandoned by its caller but has not started. close()
|
|
# must not strand its checkout when cancelling the executor queue.
|
|
serializer.close()
|
|
blocker.set()
|
|
await occupied
|
|
await _wait_for_table_empty(serializer)
|