スレッドセーフな Single-Flight TTL/LRU キャッシュを実装する

import threading
import time
from collections import OrderedDict
from typing import Callable, TypeVar, Generic, Hashable, Optional, Any

K = TypeVar("K", bound=Hashable)
V = TypeVar("V")


class _InFlight(Generic[V]):
    __slots__ = ("event", "value", "exc", "owner", "done", "detached")

    def __init__(self, owner: int) -> None:
        self.event = threading.Event()
        self.value: Optional[V] = None
        self.exc: Optional[BaseException] = None
        self.owner = owner
        self.done = False
        self.detached = False


class _Entry(Generic[V]):
    __slots__ = ("value", "expire_at")

    def __init__(self, value: V, expire_at: float) -> None:
        self.value = value
        self.expire_at = expire_at


class SingleFlightTTLCache(Generic[K, V]):
    def __init__(
        self,
        capacity: int,
        ttl: float,
        clock: Callable[[], float] = time.monotonic,
    ) -> None:
        if capacity < 0:
            raise ValueError("capacity must be non-negative")
        if ttl <= 0:
            raise ValueError("ttl must be positive")
        self._capacity = capacity
        self._ttl = ttl
        self._clock = clock
        self._lock = threading.Lock()
        self._store: "OrderedDict[K, _Entry[V]]" = OrderedDict()
        self._inflight: dict[K, _InFlight[V]] = {}
        # per-thread stack of keys currently being computed by this thread,
        # used to detect same-thread recursion into an in-flight key.
        self._thread_local = threading.local()

    def _owned_keys(self) -> set:
        stack = getattr(self._thread_local, "stack", None)
        if stack is None:
            stack = set()
            self._thread_local.stack = stack
        return stack

    def _evict_if_needed(self) -> None:
        while len(self._store) > self._capacity:
            self._store.popitem(last=False)

    def _purge_expired_locked(self) -> None:
        now = self._clock()
        expired_keys = [k for k, e in self._store.items() if now >= e.expire_at]
        for k in expired_keys:
            del self._store[k]

    def get_or_compute(self, key: K, compute: Callable[[], V]) -> V:
        owned = self._owned_keys()

        while True:
            with self._lock:
                # Check completed, unexpired entry first.
                entry = self._store.get(key)
                if entry is not None:
                    if self._clock() >= entry.expire_at:
                        del self._store[key]
                    else:
                        self._store.move_to_end(key)
                        return entry.value

                inflight = self._inflight.get(key)
                if inflight is not None:
                    if inflight.owner == threading.get_ident() and key in owned:
                        raise RuntimeError(
                            "Detected same-thread recursive computation for key: "
                            f"{key!r}"
                        )
                    # Wait for existing computation outside the lock.
                    wait_target = inflight
                else:
                    wait_target = None
                    my_flight = _InFlight(owner=threading.get_ident())
                    self._inflight[key] = my_flight

            if wait_target is not None:
                wait_target.event.wait()
                if wait_target.exc is not None:
                    raise wait_target.exc
                return wait_target.value  # type: ignore[return-value]

            # We own the computation for this key now.
            owned.add(key)
            try:
                result = compute()
            except BaseException as e:
                with self._lock:
                    current = self._inflight.get(key)
                    if current is my_flight:
                        del self._inflight[key]
                my_flight.exc = e
                my_flight.done = True
                my_flight.event.set()
                raise
            else:
                with self._lock:
                    current = self._inflight.get(key)
                    still_current = current is my_flight
                    if still_current:
                        del self._inflight[key]
                    if still_current and not my_flight.detached:
                        expire_at = self._clock() + self._ttl
                        if self._capacity > 0:
                            self._store[key] = _Entry(result, expire_at)
                            self._store.move_to_end(key)
                            self._evict_if_needed()
                        # capacity == 0: never store
                    # If detached or superseded, do not store the result,
                    # but still release waiters below with the value.
                my_flight.value = result
                my_flight.done = True
                my_flight.event.set()
                return result
            finally:
                owned.discard(key)

    def invalidate(self, key: K) -> None:
        with self._lock:
            if key in self._store:
                del self._store[key]
            inflight = self._inflight.get(key)
            if inflight is not None:
                inflight.detached = True
                del self._inflight[key]

    def clear(self) -> None:
        with self._lock:
            self._store.clear()
            for inflight in self._inflight.values():
                inflight.detached = True
            self._inflight.clear()

    def __len__(self) -> int:
        with self._lock:
            self._purge_expired_locked()
            return len(self._store)

import collections
import heapq
import threading
import time
from collections.abc import Callable
from typing import Any, Dict, Optional, Tuple, TypeVar

K = TypeVar("K")
V = TypeVar("V")


class SingleFlightTTLCache[K, V]:
    def __init__(
        self,
        capacity: int,
        ttl: float,
        clock: Callable[[], float] = time.monotonic,
    ) -> None:
        if capacity < 0:
            raise ValueError("capacity は非負でなければなりません")
        if ttl <= 0:
            raise ValueError("ttl は正の値でなければなりません")

        self.capacity = capacity
        self.ttl = ttl
        self.clock = clock

        self._cache: Dict[K, Tuple[float, V]] = {}
        self._in_flight: Dict[K, _InFlight[V]] = {}
        self._lru: collections.deque[K] = collections.deque()
        self._lock = threading.Lock()
        self._thread_recursion_guard: Dict[int, K] = {}

    def get_or_compute(self, key: K, compute: Callable[[], V]) -> V:
        with self._lock:
            if key in self._thread_recursion_guard:
                raise RuntimeError("実行中のキーへの再帰呼び出し")

            now = self.clock()
            if key in self._cache:
                expiration_time, value = self._cache[key]
                if now < expiration_time:
                    self._lru.remove(key)
                    self._lru.append(key)
                    return value
                else:
                    del self._cache[key]
                    self._lru.remove(key)

            if key in self._in_flight:
                in_flight = self._in_flight[key]
                in_flight.waiters.append(threading.current_thread())
                self._thread_recursion_guard[threading.get_ident()] = key
                try:
                    return in_flight.result.result()
                finally:
                    self._thread_recursion_guard.pop(threading.get_ident(), None)

            if self.capacity == 0:
                # キャッシュは行わず、単に計算して返す
                try:
                    return compute()
                except BaseException as e:
                    raise e

            in_flight = _InFlight[V]()
            self._in_flight[key] = in_flight
            in_flight.waiters.append(threading.current_thread())
            self._thread_recursion_guard[threading.get_ident()] = key

        try:
            result = compute()
            expiration_time = self.clock() + self.ttl
            self.put(key, expiration_time, result)
            in_flight.result.set(result)
            return result
        except BaseException as e:
            self.release_in_flight(key, e)
            raise e
        finally:
            self.release_in_flight_waiters(key)
            self._thread_recursion_guard.pop(threading.get_ident(), None)

    def put(self, key: K, expiration_time: float, value: V) -> None:
        with self._lock:
            if key in self._in_flight:
                # キーがまだ実行中の場合、それは切り離されたことを意味する
                # ため、より新しい値を上書きすべきではありません。
                # これは、compute の実行中に invalidate が呼び出された場合に起こり得ます。
                return

            self._cache[key] = (expiration_time, value)
            self._lru.append(key)
            self.evict_if_needed()

    def evict_if_needed(self) -> None:
        while len(self._cache) > self.capacity:
            lru_key = self._lru.popleft()
            if lru_key in self._cache:
                del self._cache[lru_key]

    def invalidate(self, key: K) -> None:
        with self._lock:
            if key in self._cache:
                del self._cache[key]
                self._lru.remove(key)
            if key in self._in_flight:
                # 実行中の計算をキャッシュから切り離します。
                # すでに待機している呼び出し元は引き続き結果を受け取りますが、
                # それはキャッシュされません。
                self._in_flight[key].detached = True

    def clear(self) -> None:
        with self._lock:
            self._cache.clear()
            self._lru.clear()
            for in_flight in self._in_flight.values():
                in_flight.detached = True

    def __len__(self) -> int:
        self.remove_expired()
        return len(self._cache)

    def remove_expired(self) -> None:
        with self._lock:
            now = self.clock()
            keys_to_remove = [
                key
                for key, (expiration_time, _) in self._cache.items()
                if now >= expiration_time
            ]
            for key in keys_to_remove:
                del self._cache[key]
                self._lru.remove(key)

    def release_in_flight(self, key: K, exception: BaseException) -> None:
        with self._lock:
            if key not in self._in_flight:
                return
            in_flight = self._in_flight[key]
            if not in_flight.detached:
                in_flight.result.set_exception(exception)
            else:
                # 切り離されている場合、結果はキャッシュされるべきではありません。
                # それでも待機中のものに通知する必要があります。
                in_flight.result.set_exception(exception)
            del self._in_flight[key]

    def release_in_flight_waiters(self, key: K) -> None:
        with self._lock:
            if key not in self._in_flight:
                return
            in_flight = self._in_flight[key]
            # すべての待機中のものを起こす
            for waiter in in_flight.waiters:
                waiter.unblock()
            in_flight.waiters.clear()


class _InFlight[V]:
    def __init__(self) -> None:
        self.result: threading.Future[V] = threading.Future()
        self.waiters: list[threading.Thread] = []
        self.detached: bool = False