Source code for zict.async_buffer

from __future__ import annotations

import asyncio
import contextvars
from import Callable, Collection
from concurrent.futures import Executor, ThreadPoolExecutor
from functools import wraps
from itertools import chain
from typing import Any, Literal

from zict.buffer import Buffer
from zict.common import KT, VT, T, locked

[docs] class AsyncBuffer(Buffer[KT, VT]): """Extension of :class:`~zict.Buffer` that allows offloading all reads and writes from/to slow to a separate worker thread. This requires ``fast`` to be fully thread-safe (e.g. a plain dict). ``slow.__setitem__`` and ``slow.__getitem__`` will be called from the offloaded thread, while all of its other methods (including, notably for the purpose of thread-safety consideration, ``__contains__`` and ``__delitem__``) will be called from the main thread. See Also -------- Buffer Parameters ---------- Same as in Buffer, plus: executor: concurrent.futures.Executor, optional An Executor instance to use for offloading. It must not pickle/unpickle. Defaults to an internal ThreadPoolExecutor. nthreads: int, optional Number of offloaded threads to run in parallel. Defaults to 1. Mutually exclusive with executor parameter. """ executor: Executor | None nthreads: int | None futures: set[asyncio.Future] evicting: dict[asyncio.Future, float] @wraps(Buffer.__init__) def __init__( self, *args: Any, executor: Executor | None = None, nthreads: int = 1, **kwargs: Any, ) -> None: super().__init__(*args, **kwargs) self.executor = executor self.nthreads = None if executor else nthreads self._internal_executor = executor is None self.futures = set() self.evicting = {}
[docs] def close(self) -> None: # Call LRU.close(), which stops LRU.evict_until_below_target() halfway through super().close() for future in self.futures: future.cancel() if self.executor is not None and self.nthreads is not None: self.executor.shutdown(wait=True) self.executor = None
def _offload(self, func: Callable[..., T], *args: Any) -> asyncio.Future[T]: if self.executor is None: assert self.nthreads self.executor = ThreadPoolExecutor( self.nthreads, thread_name_prefix="zict.AsyncBuffer offloader" ) loop = asyncio.get_running_loop() context = contextvars.copy_context() future = loop.run_in_executor(self.executor,, func, *args) self.futures.add(future) future.add_done_callback(self.futures.remove) return future # type: ignore[return-value] # Return an asyncio.Future, instead of just writing it as an async function, to make # it easier for overriders to tell apart the use case when all keys were already # in fast
[docs] @locked def async_get( self, keys: Collection[KT], missing: Literal["raise", "omit"] = "raise" ) -> asyncio.Future[dict[KT, VT]]: """Fetch one or more key/value pairs. If not all keys are available in fast, offload to a worker thread moving keys from slow to fast, as well as possibly moving older keys from fast to slow. Parameters ---------- keys: collection of zero or more keys to get missing: raise or omit, optional raise (default) If any key is missing, raise KeyError. omit If a key is missing, return a dict with less keys than those requested. Notes ----- All keys may be present when you call ``async_get``, but ``__delitem__`` may be called on one of them before the actual data is fetched. ``__setitem__`` also internally calls ``__delitem__`` in a non-atomic way, so you may get ``KeyError`` when updating a value too. """ # This block avoids spawning a thread if keys are missing from both fast and # slow. It is otherwise just a performance optimization. if missing == "omit": keys = [key for key in keys if key in self] elif missing == "raise": for key in keys: if key not in self: raise KeyError(key) else: raise ValueError(f"missing: expected raise or omit; got {missing}") # End performance optimization try: # Do not pull keys towards the top of the LRU unless they are all available. # This matters when there is a very long queue of async_get futures. d = except KeyError: pass else: f: asyncio.Future[dict[KT, VT]] = asyncio.Future() f.set_result(d) return f def _async_get() -> dict[KT, VT]: d = {} for k in keys: if raise asyncio.CancelledError() try: # This can cause keys to be restored and older keys to be evicted d[k] = self[k] except KeyError: # Race condition: key was there when async_get was called, but got # deleted afterwards. if missing == "raise": raise return d return self._offload(_async_get)
def __setitem__(self, key: KT, value: VT) -> None: """Immediately set a key in fast. If this causes the total weight to exceed n, asynchronously start moving keys from fast to slow in a worker thread. """ self.set_noevict(key, value) self.async_evict_until_below_target()
[docs] @locked def async_evict_until_below_target(self, n: float | None = None) -> None: """If the total weight exceeds n, asynchronously start moving keys from fast to slow in a worker thread. """ if n is None: n = self.n n = max(0.0, n) weight = min(chain([], self.evicting.values())) if weight <= n: return # Note: this can get cancelled by LRU.close(), which in turn is # triggered by Buffer.close() future = self._offload(self.evict_until_below_target, n) self.evicting[future] = n future.add_done_callback(self.evicting.__delitem__)