hug/lib/python3.12/site-packages/filelock/_async_read_write.py

"""Async wrapper around :class:`ReadWriteLock` for use with ``asyncio``."""

from __future__ import annotations

__lazy_modules__ = ["asyncio", "concurrent.futures", "filelock._async", "filelock._error", "functools", "threading"]

import asyncio
import functools
import os
import sqlite3
import threading
from concurrent.futures import ThreadPoolExecutor
from concurrent.futures import thread as executor_threads
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING, Final, Literal, ParamSpec, TypeVar, cast

from filelock._api import (
    _append_exception_context,
    _ensure_current_process,
    _fork_transition,
    _register_fork_object,
)
from filelock._async import (
    _TASKS_THREAD_ID,
    _BackendOutcome,
    _capture_call,
    _drain_future,
    _future_result,
    _raise_cancelled_error,
    _task_owners_for,
    _wait_until_done,
)
from filelock._error import Timeout
from filelock._read_write import ReadWriteLock

if TYPE_CHECKING:
    from collections.abc import AsyncGenerator, Callable, MutableMapping
    from concurrent import futures
    from types import TracebackType

    from filelock._api import AcquireReturnProxy


_P = ParamSpec("_P")
_R = TypeVar("_R")


class AsyncReadWriteLock:
    """
    Async wrapper around :class:`ReadWriteLock` for use in ``asyncio`` applications.

    This wrapper dispatches every blocking SQLite operation to a thread pool via ``loop.run_in_executor()`` because
    Python's :mod:`sqlite3` module has no async API. Each ``asyncio`` task owns its own hold. Tasks share a read lock,
    and a task asking for the write lock waits until the other holders release, with waiting writers ahead of new
    readers. As with :class:`ReadWriteLock`, upgrading or downgrading a held lock raises :class:`RuntimeError`.
    Singleton wrappers over one path share their task holds, and closing one ends only the holds taken through it.

    :param lock_file: path to the SQLite database file used as the lock
    :param timeout: maximum wait time in seconds; ``-1`` waits up to SQLite's busy-timeout cap of about 23 days
    :param blocking: if ``False``, raise :class:`~filelock.Timeout` immediately when the lock is unavailable
    :param is_singleton: if ``True``, reuse existing :class:`ReadWriteLock` instances for the same resolved path
    :param loop: event loop for ``run_in_executor``; ``None`` uses the running loop
    :param executor: executor for ``run_in_executor``. When ``None`` this lock creates and owns a dedicated
        single-thread executor so every operation runs on the same thread (SQLite affinity requires this) and shuts it
        down in :meth:`close`. This lock uses a caller-supplied executor as-is and never shuts it down, so after passing
        no executor call :meth:`close` to release the owned one.

    .. versionadded:: 3.21.0

    """

    def __init__(  # ruff:ignore[too-many-arguments]  # public constructor: one parameter per documented lock option
        self,
        lock_file: str | os.PathLike[str],
        timeout: float = -1,
        *,
        blocking: bool = True,
        is_singleton: bool = True,
        loop: asyncio.AbstractEventLoop | None = None,
        executor: futures.Executor | None = None,
    ) -> None:
        creator_pid = os.getpid()
        self._creator_pid = creator_pid
        self._fork_invalidated = False
        self._closed = False
        self._owns_backend = not is_singleton
        #: Marks the task holds taken through this wrapper, which its close() ends.
        self._via: Final = object()
        _register_fork_object(self)
        with _fork_transition():
            self._lock = ReadWriteLock(lock_file, timeout, blocking=blocking, is_singleton=is_singleton)
            self._owners = _task_owners_for(self._lock)
            self._loop = loop
            self._owns_executor = executor is None
            self._executor = executor or ThreadPoolExecutor(max_workers=1)
            if os.getpid() != creator_pid:  # pragma: forked child
                msg = "AsyncReadWriteLock construction cannot continue after fork"
                raise RuntimeError(msg)

    @property
    def lock_file(self) -> str:
        """The path to the lock file."""
        return self._lock.lock_file

    @property
    def timeout(self) -> float:
        """The default timeout."""
        return self._lock.timeout

    @property
    def blocking(self) -> bool:
        """Whether blocking is enabled by default."""
        return self._lock.blocking

    @property
    def loop(self) -> asyncio.AbstractEventLoop | None:
        """The event loop (or ``None`` for the running loop)."""
        return self._loop

    @property
    def executor(self) -> futures.Executor:
        """The executor used for ``run_in_executor`` (a dedicated single-thread one if none was supplied)."""
        return self._executor

    @asynccontextmanager
    async def read_lock(self, timeout: float | None = None, *, blocking: bool | None = None) -> AsyncGenerator[None]:
        """
        Async context manager that acquires and releases a shared read lock.

        Falls back to instance defaults for *timeout* and *blocking* when ``None``.

        :param timeout: maximum wait time in seconds, or ``None`` to use the instance default
        :param blocking: if ``False``, raise :class:`~filelock.Timeout` immediately; ``None`` uses the instance default

        """
        await self.acquire_read(timeout, blocking=blocking)
        body_error: BaseException | None = None
        try:
            yield
        except BaseException as error:
            body_error = error
            raise
        finally:
            await self._release_in_context(body_error)

    @asynccontextmanager
    async def write_lock(self, timeout: float | None = None, *, blocking: bool | None = None) -> AsyncGenerator[None]:
        """
        Async context manager that acquires and releases an exclusive write lock.

        Falls back to instance defaults for *timeout* and *blocking* when ``None``.

        :param timeout: maximum wait time in seconds, or ``None`` to use the instance default
        :param blocking: if ``False``, raise :class:`~filelock.Timeout` immediately; ``None`` uses the instance default

        """
        await self.acquire_write(timeout, blocking=blocking)
        body_error: BaseException | None = None
        try:
            yield
        except BaseException as error:
            body_error = error
            raise
        finally:
            await self._release_in_context(body_error)

    async def _release_in_context(self, body_error: BaseException | None) -> None:
        try:
            await self.release()
        except BaseException as release_error:
            if body_error is not None:
                _append_exception_context(release_error, body_error)
            raise

    async def acquire_read(
        self, timeout: float | None = None, *, blocking: bool | None = None
    ) -> AsyncAcquireReadWriteReturnProxy:
        """
        Acquire a shared read lock.

        See :meth:`ReadWriteLock.acquire_read` for full semantics.

        :param timeout: seconds to wait; ``None`` uses the instance setting; ``-1`` waits up to SQLite's
            busy-timeout cap of about 23 days
        :param blocking: if ``False``, raise :class:`~filelock.Timeout` on contention;
            ``None`` uses the instance setting

        :returns: a proxy that can be used as an async context manager to release the lock

        :raises RuntimeError: if the calling task already holds the write lock
        :raises Timeout: if the lock cannot be acquired within *timeout* seconds
        :raises TypeError: if *timeout* is a :class:`bool` or not a real number
        :raises ValueError: if a blocking call gets a ``nan`` *timeout* or a negative one other than ``-1``

        """
        self._raise_if_unusable()
        await self._acquire("read", timeout, blocking=blocking)
        return AsyncAcquireReadWriteReturnProxy(lock=self)

    async def acquire_write(
        self, timeout: float | None = None, *, blocking: bool | None = None
    ) -> AsyncAcquireReadWriteReturnProxy:
        """
        Acquire an exclusive write lock.

        See :meth:`ReadWriteLock.acquire_write` for full semantics.

        :param timeout: seconds to wait; ``None`` uses the instance setting; ``-1`` waits up to SQLite's
            busy-timeout cap of about 23 days
        :param blocking: if ``False``, raise :class:`~filelock.Timeout` on contention;
            ``None`` uses the instance setting

        :returns: a proxy that can be used as an async context manager to release the lock

        :raises RuntimeError: if the calling task already holds the read lock
        :raises Timeout: if the lock cannot be acquired within *timeout* seconds
        :raises TypeError: if *timeout* is a :class:`bool` or not a real number
        :raises ValueError: if a blocking call gets a ``nan`` *timeout* or a negative one other than ``-1``

        """
        self._raise_if_unusable()
        await self._acquire("write", timeout, blocking=blocking)
        return AsyncAcquireReadWriteReturnProxy(lock=self)

    async def _acquire(self, mode: Literal["read", "write"], timeout: float | None, *, blocking: bool | None) -> None:
        timeout = self.timeout if timeout is None else timeout
        blocking = self.blocking if blocking is None else blocking

        async def enter(remaining: float) -> None:
            # A close() may have landed while this task waited for the table.
            self._raise_if_unusable()
            await self._run_acquire(
                functools.partial(
                    self._lock._acquire_level_of,  # ruff:ignore[private-member-access]  # the wrapper owns this backend hold
                    _TASKS_THREAD_ID,
                    mode,
                    remaining,
                    blocking=blocking,
                )
            )

        def joinable() -> bool:
            self._raise_if_unusable()
            return self._backend_holds()

        await self._owners.acquire(
            mode,
            timeout=timeout,
            blocking=blocking,
            lock_file=self.lock_file,
            via=self._via,
            enter=enter,
            joinable=joinable,
        )

    async def release(self, *, force: bool = False) -> None:
        """
        Release one level of the calling task's hold.

        The wrapper releases the database lock once no task holds it.

        :param force: if ``True``, end every level the calling task took through this wrapper. Levels taken through
            other wrappers on the same path stay. A wrapper that holds no level for the task releases nothing.

        :raises RuntimeError: if the calling task holds no level through this wrapper and *force* is ``False``

        """
        _ensure_current_process()
        if self._inherited:  # pragma: needs fork
            return
        await self._owners.release(
            force=force,
            lock_file=self.lock_file,
            via=self._via,
            leave=functools.partial(self._submit, self._leave_backend),
            still_held=self._backend_holds,
        )

    async def close(self) -> None:
        """
        End the holds of the tasks that acquired through this instance and release the lock once no task holds it.

        After calling this method, the lock instance is no longer usable. With ``is_singleton=False`` it also closes the
        underlying SQLite connection. A singleton wrapper leaves the shared :class:`ReadWriteLock` open, so tasks that
        acquired through another wrapper and synchronous users of the same path keep their holds. The next acquire
        notices a synchronous ``close()`` or forced release of that shared lock. One that lands in the few bytecodes
        between that check and the admission still admits the task into the ended hold.

        """
        _ensure_current_process()
        if self._inherited:  # pragma: needs fork
            return
        if self._closed:
            return
        # Refuse acquires from here on, so none is admitted into a hold the backend is letting go of.
        self._closed = True
        try:
            await self._owners.retire(
                None if self._owns_backend else self._via,
                leave=functools.partial(self._submit, self._lock.close if self._owns_backend else self._leave_backend),
                still_held=self._backend_holds,
            )
        except BaseException as error:
            # Only a failed backend call can leave the level held, and then a retried close() has to find it.
            if not isinstance(error, asyncio.CancelledError) and self._backend_holds():
                self._closed = False
            else:
                self._shutdown_owned_executor()
            raise
        # Wait for the worker to exit rather than letting it drain in the background: a caller that forks right
        # after closing deserves a single-threaded process, and os.fork warns about any surviving thread.
        if self._owns_executor:
            await asyncio.to_thread(functools.partial(self._executor.shutdown, wait=True))

    async def _run_acquire(self, acquire: Callable[..., AcquireReturnProxy]) -> None:
        abandoned = threading.Event()
        acquire_future = self._submit(acquire, cancel_check=abandoned.is_set)
        try:
            await _wait_until_done(acquire_future)
        except asyncio.CancelledError as cancellation:
            # Stops a backend wait on another process's hold, which would otherwise outlast the caller by as long.
            abandoned.set()
            try:
                await _drain_future(acquire_future)
            except Timeout:
                # Nothing was taken, so the caller's cancellation is the whole outcome.
                raise cancellation from None
            except BaseException as error:  # ruff:ignore[blind-except]  # reported with the cancellation below
                _raise_cancelled_error(cancellation, error)
            try:
                await _drain_future(self._submit(self._leave_backend))
            except BaseException as error:  # ruff:ignore[blind-except]  # reported with the cancellation below
                _raise_cancelled_error(cancellation, error)
            raise
        _future_result(acquire_future)

    def _backend_holds(self) -> bool:
        # A sync close() or forced release on a shared singleton ends the tasks' level without the table taking part.
        return self._lock._holds_level_of(_TASKS_THREAD_ID)  # ruff:ignore[private-member-access]  # the wrapper owns this backend hold

    def _leave_backend(self) -> None:
        # Only the tasks' level: a sync user of the same singleton keeps its own.
        self._lock._release_level_of(_TASKS_THREAD_ID)  # ruff:ignore[private-member-access]  # the wrapper owns this backend hold

    def _submit(
        self, func: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs
    ) -> asyncio.Future[_BackendOutcome[_R]]:
        return (self._loop or asyncio.get_running_loop()).run_in_executor(
            self._executor,
            _capture_call,
            functools.partial(func, *args, **kwargs),
        )

    def _shutdown_owned_executor(self) -> None:
        if self._owns_executor:
            self._executor.shutdown(wait=False)

    @property
    def _inherited(self) -> bool:
        return self._fork_invalidated or os.getpid() != self._creator_pid

    def _raise_if_unusable(self) -> None:
        _ensure_current_process()
        if self._inherited:  # pragma: needs fork
            msg = f"AsyncReadWriteLock on {self.lock_file} was invalidated by fork(); construct a new instance"
            raise RuntimeError(msg)
        if self._closed:
            msg = "Cannot operate on a closed database."
            raise sqlite3.ProgrammingError(msg)

    def _reset_after_fork_in_child(self) -> None:  # pragma: forked child
        self._fork_invalidated = True
        if self._owns_executor:
            # The worker did not survive the fork, yet before 3.12 the child's exit still wakes it through its queue. On
            # PyPy that queue is pure Python, and its lock may have been held by the worker at fork time, so exit hangs.
            # CPython 3.12+ forgets every inherited worker in the child; do the same for the one this lock owns.
            # typeshed declares the registry read-only; it is a WeakKeyDictionary the exit hook iterates.
            registry = cast(
                "MutableMapping[threading.Thread, object]",
                executor_threads._threads_queues,  # ruff:ignore[private-member-access]  # the exit hook's registry
            )
            for worker in cast("ThreadPoolExecutor", self._executor)._threads:  # ruff:ignore[private-member-access]  # only the executor lists its workers
                registry.pop(worker, None)

    def __del__(self) -> None:
        # Safety net when close() was never called: shut down the executor we own so its worker thread does not
        # outlive the lock. shutdown(wait=False) never blocks.
        if os.getpid() == getattr(self, "_creator_pid", None) and getattr(self, "_owns_executor", False):
            self._executor.shutdown(wait=False)


class AsyncAcquireReadWriteReturnProxy:
    """Context-aware object that releases the async read/write lock on exit."""

    def __init__(self, lock: AsyncReadWriteLock) -> None:
        self.lock = lock

    async def __aenter__(self) -> AsyncReadWriteLock:
        return self.lock

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> None:
        await self.lock.release()


__all__ = [
    "AsyncAcquireReadWriteReturnProxy",
    "AsyncReadWriteLock",
]