スレッドセーフな 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