207 lines
7.0 KiB
Python
207 lines
7.0 KiB
Python
"""外部系统进程内运行态缓存与同步 single-flight。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import atexit
|
||
import time
|
||
from collections import OrderedDict
|
||
from concurrent.futures import Future
|
||
from contextlib import contextmanager
|
||
from dataclasses import dataclass
|
||
from threading import RLock
|
||
from typing import Any, Callable, Iterator, TypeVar
|
||
|
||
|
||
T = TypeVar("T")
|
||
|
||
|
||
@dataclass
|
||
class _RuntimeEntry:
|
||
client: Any = None
|
||
active_client_leases: int = 0
|
||
evicted: bool = False
|
||
auth_headers: dict[str, str] | None = None
|
||
auth_expires_at: float = 0.0
|
||
auth_generation: int = 0
|
||
spec: dict[str, Any] | None = None
|
||
spec_expires_at: float = 0.0
|
||
catalog: Any = None
|
||
catalog_spec: dict[str, Any] | None = None
|
||
|
||
|
||
class ExternalRuntimeCache:
|
||
"""按连接身份隔离的有界 LRU；凭据、Token 和规格均只驻留当前进程。"""
|
||
|
||
def __init__(self, *, max_entries: int = 256):
|
||
self.max_entries = max_entries
|
||
self._entries: OrderedDict[str, _RuntimeEntry] = OrderedDict()
|
||
self._inflight: dict[tuple[str, str], Future[Any]] = {}
|
||
self._lock = RLock()
|
||
|
||
@staticmethod
|
||
def _close_client(client: Any) -> None:
|
||
close = getattr(client, "close", None)
|
||
if callable(close):
|
||
try:
|
||
close()
|
||
except Exception:
|
||
pass
|
||
|
||
def _entry_locked(self, identity: str) -> _RuntimeEntry:
|
||
entry = self._entries.get(identity)
|
||
if entry is None:
|
||
entry = _RuntimeEntry()
|
||
self._entries[identity] = entry
|
||
self._entries.move_to_end(identity)
|
||
while len(self._entries) > self.max_entries:
|
||
_, evicted = self._entries.popitem(last=False)
|
||
evicted.evicted = True
|
||
if evicted.active_client_leases == 0:
|
||
self._close_client(evicted.client)
|
||
return entry
|
||
|
||
@contextmanager
|
||
def client(self, identity: str, factory: Callable[[], T]) -> Iterator[T]:
|
||
with self._lock:
|
||
entry = self._entry_locked(identity)
|
||
if entry.client is None:
|
||
entry.client = factory()
|
||
entry.active_client_leases += 1
|
||
client = entry.client
|
||
try:
|
||
yield client
|
||
finally:
|
||
close_client = None
|
||
with self._lock:
|
||
entry.active_client_leases -= 1
|
||
if entry.evicted and entry.active_client_leases == 0:
|
||
close_client = entry.client
|
||
entry.client = None
|
||
if close_client is not None:
|
||
self._close_client(close_client)
|
||
|
||
def get_auth(self, identity: str) -> dict[str, str] | None:
|
||
now = time.monotonic()
|
||
with self._lock:
|
||
entry = self._entries.get(identity)
|
||
if entry is None or entry.auth_expires_at <= now:
|
||
if entry is not None:
|
||
entry.auth_headers = None
|
||
entry.auth_expires_at = 0.0
|
||
return None
|
||
self._entries.move_to_end(identity)
|
||
return dict(entry.auth_headers or {})
|
||
|
||
def set_auth(
|
||
self, identity: str, headers: dict[str, str], *, ttl_seconds: float
|
||
) -> None:
|
||
with self._lock:
|
||
entry = self._entry_locked(identity)
|
||
entry.auth_headers = dict(headers)
|
||
entry.auth_expires_at = time.monotonic() + max(0.0, ttl_seconds)
|
||
entry.auth_generation += 1
|
||
|
||
def auth_state(self, identity: str) -> tuple[dict[str, str] | None, int]:
|
||
headers = self.get_auth(identity)
|
||
with self._lock:
|
||
entry = self._entries.get(identity)
|
||
return headers, entry.auth_generation if entry is not None else 0
|
||
|
||
def invalidate_auth(self, identity: str) -> None:
|
||
with self._lock:
|
||
entry = self._entries.get(identity)
|
||
if entry is not None:
|
||
entry.auth_headers = None
|
||
entry.auth_expires_at = 0.0
|
||
|
||
def get_spec(self, identity: str) -> dict[str, Any] | None:
|
||
now = time.monotonic()
|
||
with self._lock:
|
||
entry = self._entries.get(identity)
|
||
if entry is None or entry.spec_expires_at <= now:
|
||
if entry is not None:
|
||
entry.spec = None
|
||
entry.spec_expires_at = 0.0
|
||
entry.catalog = None
|
||
entry.catalog_spec = None
|
||
return None
|
||
self._entries.move_to_end(identity)
|
||
return entry.spec
|
||
|
||
def set_spec(
|
||
self, identity: str, spec: dict[str, Any], *, ttl_seconds: float
|
||
) -> None:
|
||
with self._lock:
|
||
entry = self._entry_locked(identity)
|
||
entry.spec = spec
|
||
entry.spec_expires_at = time.monotonic() + max(0.0, ttl_seconds)
|
||
entry.catalog = None
|
||
entry.catalog_spec = None
|
||
|
||
def invalidate_spec(self, identity: str) -> None:
|
||
with self._lock:
|
||
entry = self._entries.get(identity)
|
||
if entry is not None:
|
||
entry.spec = None
|
||
entry.spec_expires_at = 0.0
|
||
entry.catalog = None
|
||
entry.catalog_spec = None
|
||
|
||
def get_catalog(self, identity: str, spec: dict[str, Any]) -> Any:
|
||
with self._lock:
|
||
entry = self._entries.get(identity)
|
||
if entry is None or entry.catalog_spec is not spec:
|
||
return None
|
||
self._entries.move_to_end(identity)
|
||
return entry.catalog
|
||
|
||
def set_catalog(self, identity: str, spec: dict[str, Any], catalog: Any) -> None:
|
||
with self._lock:
|
||
entry = self._entry_locked(identity)
|
||
entry.catalog_spec = spec
|
||
entry.catalog = catalog
|
||
|
||
def singleflight(self, namespace: str, key: str, compute: Callable[[], T]) -> T:
|
||
flight_key = (namespace, key)
|
||
with self._lock:
|
||
future = self._inflight.get(flight_key)
|
||
leader = future is None
|
||
if leader:
|
||
future = Future()
|
||
self._inflight[flight_key] = future
|
||
assert future is not None
|
||
if not leader:
|
||
return future.result()
|
||
try:
|
||
result = compute()
|
||
except BaseException as exc:
|
||
future.set_exception(exc)
|
||
raise
|
||
else:
|
||
future.set_result(result)
|
||
return result
|
||
finally:
|
||
with self._lock:
|
||
if self._inflight.get(flight_key) is future:
|
||
self._inflight.pop(flight_key, None)
|
||
|
||
def clear(self) -> None:
|
||
with self._lock:
|
||
entries = list(self._entries.values())
|
||
self._entries.clear()
|
||
for entry in entries:
|
||
self._close_client(entry.client)
|
||
|
||
def spec_count(self) -> int:
|
||
now = time.monotonic()
|
||
with self._lock:
|
||
return sum(
|
||
1
|
||
for entry in self._entries.values()
|
||
if entry.spec is not None and entry.spec_expires_at > now
|
||
)
|
||
|
||
|
||
RUNTIME_CACHE = ExternalRuntimeCache()
|
||
atexit.register(RUNTIME_CACHE.clear)
|