deer-flow/backend/packages/harness/deerflow/sandbox/acquire_serialization.py

"""Per-key acquire serialization with bounded lock-table growth (RFC #4741).

Shared component for sandbox providers. Each provider selects the key that
preserves its current collision and lifecycle semantics (``(user_id,
thread_id)`` for AIO/E2B, the derived sandbox id for BoxLite/Tenki/
OpenSandbox); the serializer never interprets it.

The async path runs the blocking ``threading.Lock.acquire`` on a bounded
dedicated executor (never the default executor, never the event loop). A small
handoff state makes the acquire worker release an abandoned lock itself, so
cleanup does not depend on the cancelling event loop running a done callback.
"""

from __future__ import annotations

import asyncio
import contextvars
import os
import threading
from collections.abc import AsyncIterator, Callable, Hashable, Iterator
from concurrent.futures import Future, ThreadPoolExecutor
from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass, field
from functools import partial

DEFAULT_MAX_WORKERS = min(32, (os.cpu_count() or 1) + 4)


@dataclass
class _Entry:
    lock: threading.Lock = field(default_factory=threading.Lock)
    refs: int = 0  # holders + waiters


class _AsyncAcquire:
    """Race-safe ownership handoff between an event loop and acquire worker."""

    def __init__(self, entry: _Entry, cleanup: Callable[[], None]) -> None:
        self._entry = entry
        self._cleanup = cleanup
        self._state_lock = threading.Lock()
        self._abandoned = False
        self._acquired = False
        self._cleaned = False

    def run(self) -> bool:
        self._entry.lock.acquire()
        with self._state_lock:
            if self._abandoned:
                cleanup = self._mark_cleaned()
            else:
                self._acquired = True
                cleanup = False
        if cleanup:
            self._entry.lock.release()
            self._cleanup()
        return not cleanup

    def abandon(self) -> None:
        with self._state_lock:
            self._abandoned = True
            cleanup = self._acquired and self._mark_cleaned()
            if cleanup:
                self._acquired = False
        if cleanup:
            self._entry.lock.release()
            self._cleanup()

    def cancel_queued(self) -> None:
        """Clean up an executor job cancelled before its worker started."""
        with self._state_lock:
            cleanup = self._mark_cleaned()
        if cleanup:
            self._cleanup()

    def worker_done(self, future: Future[bool]) -> None:
        """Reclaim a checkout if executor shutdown cancelled queued work."""
        if future.cancelled():
            self.cancel_queued()

    def release(self) -> None:
        with self._state_lock:
            cleanup = self._acquired and self._mark_cleaned()
            if cleanup:
                self._acquired = False
        if not cleanup:  # pragma: no cover - indicates an internal ownership bug
            raise RuntimeError("Serializer lock is not held")
        self._entry.lock.release()
        self._cleanup()

    def _mark_cleaned(self) -> bool:
        if self._cleaned:
            return False
        self._cleaned = True
        return True


class AcquireSerializer[KeyT: Hashable]:
    """Serialize provider-selected lifecycle transitions per key."""

    def __init__(
        self,
        *,
        max_workers: int = DEFAULT_MAX_WORKERS,
        thread_name_prefix: str = "sandbox-acquire-wait",
    ) -> None:
        self._table: dict[KeyT, _Entry] = {}
        self._table_lock = threading.Lock()
        self._closed = False
        self._executor = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix=thread_name_prefix)

    @property
    def executor(self) -> ThreadPoolExecutor:
        """The dedicated bounded executor backing async waits.

        Providers may offload an entire synchronous acquire (with hold()
        inside) here so a cancelled awaiter abandons the worker, not the
        lock: the hold then follows the body to completion, matching the
        pre-serializer to_thread bridge semantics.
        """
        return self._executor

    async def run_on_executor[**P, T](self, func: Callable[P, T], /, *args: P.args, **kwargs: P.kwargs) -> T:
        """Run a blocking callable on the serializer's dedicated executor.

        ``asyncio.to_thread`` copies ``ContextVar`` values automatically; raw
        ``loop.run_in_executor`` does not. The providers whose ``acquire_async``
        offloads the whole synchronous acquire here replaced the inherited
        ``to_thread`` bridge, so copy the calling context explicitly: without
        it, request-scoped ContextVars (e.g. the trace id bound by
        ``request_trace_context``) read as unset inside the worker thread.
        """
        loop = asyncio.get_running_loop()
        ctx = contextvars.copy_context()
        call = partial(func, *args, **kwargs)
        return await loop.run_in_executor(self._executor, ctx.run, call)

    @contextmanager
    def hold(self, key: KeyT) -> Iterator[None]:
        entry = self._checkout(key)
        entry.lock.acquire()
        try:
            yield
        finally:
            entry.lock.release()
            self._checkin(key, entry)

    @asynccontextmanager
    async def hold_async(self, key: KeyT) -> AsyncIterator[None]:
        entry = self._checkout(key)
        loop = asyncio.get_running_loop()
        acquire = _AsyncAcquire(entry, lambda: self._checkin(key, entry))
        try:
            worker_future = self._executor.submit(acquire.run)
        except RuntimeError:  # executor shut down between checkout and scheduling
            self._checkin(key, entry)
            raise RuntimeError("AcquireSerializer is closed") from None
        # concurrent.futures callbacks run in the cancelling/worker thread, so
        # executor-shutdown cleanup does not depend on this loop progressing.
        worker_future.add_done_callback(acquire.worker_done)
        acquire_future = asyncio.wrap_future(worker_future, loop=loop)
        try:
            acquired = await asyncio.shield(acquire_future)
        except asyncio.CancelledError:
            task = asyncio.current_task()
            if acquire_future.cancelled() and task is not None and task.cancelling() == 0:
                # close() cancelled the queued wait via executor shutdown, not
                # the caller: check in exactly once (no done callback — the
                # cancelled future is already done) and report the documented
                # closure error instead of a spurious CancelledError. The
                # `cancelling() == 0` guard assumes this await is the task's
                # only pending cancellation (no stale count left by an outer
                # swallowed CancelledError).
                acquire.cancel_queued()
                raise RuntimeError("AcquireSerializer is closed") from None
            if acquire_future.cancelled():
                acquire.cancel_queued()
            else:
                acquire.abandon()
            raise
        # run() returns False only when the worker observed _abandoned, and
        # abandon() is called solely from the except handler above, which
        # always re-raises — so a normal completion implies the lock is held.
        # If this were reachable, the worker would already have checked in;
        # an extra _checkin here would underflow the refcount, so assert.
        assert acquired, "acquire worker returned without holding the lock"
        try:
            yield
        finally:
            acquire.release()

    def close(self) -> None:
        """Reject new holders and release executor resources. Idempotent.

        A currently executing critical section is not invalidated: its exit
        still releases its lock and reclaims its table entry.
        """
        with self._table_lock:
            if self._closed:
                return
            self._closed = True
        self._executor.shutdown(wait=False, cancel_futures=True)

    def _checkout(self, key: KeyT) -> _Entry:
        with self._table_lock:
            if self._closed:
                raise RuntimeError("AcquireSerializer is closed")
            entry = self._table.get(key)
            if entry is None:
                entry = _Entry()
                self._table[key] = entry
            entry.refs += 1
            return entry

    def _checkin(self, key: KeyT, entry: _Entry) -> None:
        with self._table_lock:
            entry.refs -= 1
            if entry.refs == 0 and not entry.lock.locked():
                self._table.pop(key, None)