"""Separate caller cancellation from backend task and executor-future results.""" from __future__ import annotations __lazy_modules__ = ["asyncio", "collections", "concurrent.futures", "filelock._error"] import asyncio import contextlib import time from collections import deque from concurrent.futures import Future as ConcurrentFuture from dataclasses import dataclass from threading import Lock from typing import TYPE_CHECKING, Final, Generic, Literal, NoReturn, TypeVar, cast from weakref import WeakKeyDictionary from filelock._api import ( _append_exception_context, _fork_transition, _raise_chained_errors, _resolve_read_write_timeout, ) from filelock._error import Timeout if TYPE_CHECKING: from collections.abc import AsyncGenerator, Awaitable, Callable from concurrent.futures import Executor from filelock._read_write import ReadWriteLock from filelock._soft_rw import SoftReadWriteLock _T = TypeVar("_T") #: The holder an async wrapper records its backend level under. No thread has this identifier, so sync code on the #: executor thread that took the level cannot re-enter or release it. _TASKS_THREAD_ID: Final[int] = 0 class _AsyncTransitionUnavailableError(Exception): pass @dataclass(frozen=True) class _BackendOutcome(Generic[_T]): value: _T | None = None error: BaseException | None = None class _AsyncTransitionGate: """ First-in, first-out admission for the calls that transition one lock context. Each call queues a ticket, and the head of the queue holds the gate. A call that gives up removes its own ticket, so abandoned calls do not pile up behind a long holder. """ def __init__(self) -> None: self._queue_lock: Final[Lock] = Lock() self._queue: Final[deque[ConcurrentFuture[None]]] = deque() @contextlib.asynccontextmanager async def hold(self) -> AsyncGenerator[None]: ticket: ConcurrentFuture[None] = ConcurrentFuture() try: self._join(ticket) if not ticket.done(): await self._wait_for_turn(ticket, blocking=True, cancel_check=None, deadline=None, poll_interval=0.05) yield finally: self._leave(ticket) @contextlib.asynccontextmanager async def hold_for_acquire( self, *, blocking: bool, cancel_check: Callable[[], bool] | None, deadline: float | None, poll_interval: float, ) -> AsyncGenerator[None]: ticket: ConcurrentFuture[None] = ConcurrentFuture() try: self._join(ticket) if not ticket.done(): await self._wait_for_turn( ticket, blocking=blocking, cancel_check=cancel_check, deadline=deadline, poll_interval=poll_interval, ) yield finally: self._leave(ticket) async def _wait_for_turn( self, ticket: ConcurrentFuture[None], *, blocking: bool, cancel_check: Callable[[], bool] | None, deadline: float | None, poll_interval: float, ) -> None: if not blocking: raise _AsyncTransitionUnavailableError waiter = asyncio.wrap_future(ticket) while not ticket.done(): if cancel_check is not None and cancel_check(): raise _AsyncTransitionUnavailableError wait_interval = poll_interval if deadline is not None: if (remaining := deadline - time.perf_counter()) <= 0: raise _AsyncTransitionUnavailableError wait_interval = min(wait_interval, remaining) await asyncio.wait((waiter,), timeout=wait_interval) # A hand-off that a signal handler's exception cut short leaves the head undone, so the head settles itself. self._settle_head() def _join(self, ticket: ConcurrentFuture[None]) -> None: with self._queue_lock: self._queue.append(ticket) if len(self._queue) == 1: ticket.set_result(None) def _leave(self, ticket: ConcurrentFuture[None]) -> None: with self._queue_lock: try: self._queue.remove(ticket) except ValueError: # The join never completed, or this ticket already left. return self._settle_head_locked() def _settle_head(self) -> None: with self._queue_lock: self._settle_head_locked() def _settle_head_locked(self) -> None: if self._queue and not (head := self._queue[0]).done(): head.set_result(None) @dataclass class _TaskHold: """ One task's hold, as the nesting count per wrapper it entered through. A task's hold ends when its last wrapper level is released or closed. Closing one wrapper ends only the levels taken through it. """ depths: dict[object, int] class _TaskOwners: """ Per-task ownership in front of a sync reader/writer lock that tracks threads. An async wrapper runs each call on an executor thread, and that thread does not identify the calling task. Waiting writers block new readers so a stream of readers cannot starve them. A thread lock guards the table because tasks on different event loops share it. The table stands for one backend level, taken once for all tasks. Sync code can end that level directly, with ``close()`` or a forced release on a shared singleton. The next admission then forgets the holds instead of letting more tasks join them. """ def __init__(self) -> None: self._lock: Final[Lock] = Lock() self._mode: Literal["read", "write"] | None = None self._holds: dict[asyncio.Task[object], _TaskHold] = {} self._transitioning = False self._waiting_writers = 0 #: One ticket per waiting call, settled on the next state change. A call that stops waiting takes its own out, #: so one that times out or is canceled leaves nothing behind. self._waiters: dict[ConcurrentFuture[None], None] = {} async def acquire( # ruff:ignore[too-many-arguments] # one keyword per piece of wrapper state the table consults self, mode: Literal["read", "write"], *, timeout: float, blocking: bool, lock_file: str, via: object, enter: Callable[[float], Awaitable[None]], joinable: Callable[[], bool], ) -> None: task = _current_task() timeout = _resolve_read_write_timeout(timeout, blocking=blocking) deadline = None if not blocking or timeout < 0 else time.perf_counter() + timeout waiting_writer = False try: while isinstance( admission := self._admit( task, mode, lock_file, via=via, joinable=joinable, waiting_writer=waiting_writer ), ConcurrentFuture, ): try: if not blocking or (deadline is not None and time.perf_counter() >= deadline): raise Timeout(lock_file) if mode == "write" and not waiting_writer: waiting_writer = True with self._lock: self._waiting_writers += 1 wait = None if deadline is None else deadline - time.perf_counter() await asyncio.wait((asyncio.wrap_future(admission),), timeout=wait) finally: self._forget(admission) finally: if waiting_writer: with self._lock: self._waiting_writers -= 1 self._notify() if admission == "held": return try: await enter(-1 if deadline is None else max(0.0, deadline - time.perf_counter())) except BaseException: self._finish_transition(mode=None) raise self._finish_transition(mode=mode, holds={task: _TaskHold(depths={via: 1})}) def _admit( # ruff:ignore[too-many-arguments] # the acquire's state, passed through under the table lock self, task: asyncio.Task[object], mode: Literal["read", "write"], lock_file: str, *, via: object, joinable: Callable[[], bool], waiting_writer: bool, ) -> Literal["held", "enter"] | ConcurrentFuture[None]: with self._lock: # A "held" answer trusts the backend level. Once the backend lost it, the holds are gone and the next task # has to take the lock again. if self._mode is not None and not self._transitioning and not joinable(): self._holds.clear() self._mode = None if (hold := self._holds.get(task)) is not None: if self._mode != mode: msg = ( f"Cannot acquire {mode} lock on {lock_file}: already holding a " f"{'write' if mode == 'read' else 'read'} lock " f"({'downgrade' if mode == 'read' else 'upgrade'} not allowed)" ) raise RuntimeError(msg) hold.depths[via] = hold.depths.get(via, 0) + 1 return "held" writers_ahead = self._waiting_writers - waiting_writer if not self._transitioning: if mode == "read" and self._mode == "read" and not writers_ahead: self._holds[task] = _TaskHold(depths={via: 1}) return "held" if not self._holds and (mode == "write" or not writers_ahead): self._transitioning, self._mode = True, mode return "enter" return self._ticket() async def release( self, *, force: bool, lock_file: str, via: object, leave: Callable[[], asyncio.Future[_BackendOutcome[None]]], still_held: Callable[[], bool], ) -> None: task = _current_task() with self._lock: if (hold := self._holds.get(task)) is None or via not in hold.depths: if force: return msg = f"Cannot release a lock on {lock_file} that is not held by this task through this wrapper" raise RuntimeError(msg) if not force and hold.depths[via] > 1: hold.depths[via] -= 1 return removed = hold.depths.pop(via) if hold.depths: return del self._holds[task] if self._holds: return self._transitioning = True mode = self._mode if not still_held(): self._finish_transition(mode=None) msg = f"Cannot release a lock on {lock_file} that is not held by this task" raise RuntimeError(msg) try: cancellation = await _drain_cancellation(leave()) except BaseException: hold.depths[via] = removed # Keep the hold for a retried release only while the sync lock still has the level. Restored after the sync # lock let go, it would stand for a lock nobody holds. self._finish_transition(mode=mode if still_held() else None, holds={task: hold}) raise self._finish_transition(mode=None) if cancellation is not None: # The release succeeded, so a restored hold would be a phantom no backend call can clear. raise cancellation async def retire( self, via: object | None, *, leave: Callable[[], asyncio.Future[_BackendOutcome[None]]], still_held: Callable[[], bool], ) -> None: """ End the holds taken through *via*, or every hold when it is ``None``, as one transition. Leaves the backend through *leave* when that ends the last hold, and always when *via* is ``None``. Arrivals wait while it runs, so none joins a hold the backend is letting go of. A transition already under way is waited out first: a task entering through *via* must not get a hold after the wrapper is gone. A cancellation is raised once the holds are settled. """ cancellation: asyncio.CancelledError | None = None while isinstance(settled := self._drop_holds_via(via), ConcurrentFuture): try: await _wait_until_done(asyncio.wrap_future(settled)) except asyncio.CancelledError as error: cancellation = error finally: self._forget(settled) mode, dropped = settled if via is None or (mode is not None and still_held()): try: cancellation = await _drain_cancellation(leave()) or cancellation except BaseException: # A close that failed with the backend still holding stays retryable, as a failed release does. self._finish_transition(mode=mode if still_held() else None, holds=dropped) raise self._finish_transition(mode=None) elif mode is not None: # Sync code already ended the level taken for tasks, so there is nothing to give back. self._finish_transition(mode=None) if cancellation is not None: raise cancellation def _drop_holds_via( self, via: object | None ) -> tuple[Literal["read", "write"] | None, dict[asyncio.Task[object], _TaskHold]] | ConcurrentFuture[None]: # Returns the future a running transition settles, or else the mode to leave (None when other holds keep the # backend level) and the dropped holds. The table stays transitioning while that mode is left. with self._lock: if self._transitioning: return self._ticket() if via is None: dropped = dict(self._holds) self._holds.clear() else: dropped = {} for task, hold in list(self._holds.items()): if (count := hold.depths.pop(via, None)) is not None: dropped[task] = _TaskHold(depths={via: count}) if not hold.depths: del self._holds[task] if self._holds and via is not None: return None, dropped mode = self._mode self._transitioning = mode is not None or via is None return mode, dropped def _finish_transition( self, *, mode: Literal["read", "write"] | None, holds: dict[asyncio.Task[object], _TaskHold] | None = None ) -> None: with self._lock: self._transitioning, self._mode = False, mode if mode is not None and holds: for task, hold in holds.items(): self._holds.setdefault(task, _TaskHold(depths={})).depths.update(hold.depths) self._notify() def _ticket(self) -> ConcurrentFuture[None]: # Called under the table lock, in the section that decided to wait, so no state change slips in between. ticket: ConcurrentFuture[None] = ConcurrentFuture() self._waiters[ticket] = None return ticket def _forget(self, ticket: ConcurrentFuture[None]) -> None: with self._lock: self._waiters.pop(ticket, None) def _notify(self) -> None: waiters, self._waiters = self._waiters, {} for ticket in waiters: ticket.set_result(None) _TASK_OWNERS: Final[WeakKeyDictionary[ReadWriteLock | SoftReadWriteLock, _TaskOwners]] = WeakKeyDictionary() _TASK_OWNERS_LOCK: Final[Lock] = Lock() def _task_owners_for(sync_lock: ReadWriteLock | SoftReadWriteLock) -> _TaskOwners: """Singleton wrappers share one sync lock, so they must share its task holds too.""" with _fork_transition(), _TASK_OWNERS_LOCK: if (owners := _TASK_OWNERS.get(sync_lock)) is None: owners = _TASK_OWNERS[sync_lock] = _TaskOwners() return owners def _current_task() -> asyncio.Task[object]: if (task := asyncio.current_task()) is None: msg = "an async reader/writer lock must be used from inside an asyncio task" raise RuntimeError(msg) return task async def _drain_cancellation(future: asyncio.Future[_BackendOutcome[None]]) -> asyncio.CancelledError | None: # Returned, not raised, so the caller can tell a call that completed under cancellation from one that failed. try: await _wait_until_done(future) except asyncio.CancelledError as cancellation: try: await _drain_future(future) except BaseException as error: # ruff:ignore[blind-except] # reported with the cancellation below _raise_cancelled_error(cancellation, error) return cancellation _future_result(future) return None async def _drain_future(future: asyncio.Future[_BackendOutcome[_T]]) -> _T: while not future.done(): with contextlib.suppress(asyncio.CancelledError): await _wait_until_done(future) return _future_result(future) async def _wait_until_done(future: asyncio.Future[_T]) -> None: if not future.done(): await asyncio.wait((future,)) def _future_result(future: asyncio.Future[_BackendOutcome[_T]]) -> _T: outcome = future.result() if (error := outcome.error) is None: return cast("_T", outcome.value) context = error.__context__ try: raise error # ruff:ignore[raise-within-try] # the handler restores context changed across the async boundary except BaseException: error.__context__ = context raise def _run_in_executor( loop: asyncio.AbstractEventLoop, executor: Executor | None, func: Callable[[], _T] ) -> asyncio.Future[_BackendOutcome[_T]]: outcome: asyncio.Future[_BackendOutcome[_T]] = loop.create_future() def settle(call: asyncio.Future[_BackendOutcome[_T]]) -> None: try: result = call.result() except asyncio.CancelledError: # shutdown(cancel_futures=True) dropped the call unrun: a CancelledError would cancel a task nobody canceled # and hide that the lock state never changed. msg = "the executor shut down before running the lock backend call" result = _BackendOutcome(error=RuntimeError(msg)) outcome.set_result(result) loop.run_in_executor(executor, _capture_call, func).add_done_callback(settle) return outcome def _capture_call(func: Callable[[], _T]) -> _BackendOutcome[_T]: try: return _BackendOutcome(value=func()) except BaseException as error: # ruff:ignore[blind-except] # backend control-flow exceptions are operation results return _BackendOutcome(error=error) def _raise_cancelled_error(cancellation: asyncio.CancelledError, error: BaseException) -> NoReturn: # A reconciliation step failed while unwinding a cancellation, so keep both exception chains. Splice the error's # existing context onto the cancellation, then make the cancellation the error's context, so both the failure and # the cancellation that triggered it survive. Shared by the async wrappers so cancellations report the same way. if (context := error.__context__) is not None and context is not cancellation: if (cancellation_context := cancellation.__context__) is not None: _append_exception_context(context, cancellation_context) cancellation.__context__ = context error.__context__ = cancellation _raise_chained_errors(error) async def _capture_awaitable(awaitable: Awaitable[_T]) -> _BackendOutcome[_T]: try: return _BackendOutcome(value=await awaitable) except BaseException as error: # ruff:ignore[blind-except] # backend cancellation must remain distinct from caller cancellation return _BackendOutcome(error=error) __all__ = [ "_TASKS_THREAD_ID", "_AsyncTransitionGate", "_AsyncTransitionUnavailableError", "_BackendOutcome", "_capture_awaitable", "_capture_call", "_drain_future", "_future_result", "_raise_cancelled_error", "_run_in_executor", "_task_owners_for", "_wait_until_done", ]