"""外部系统进程内运行态缓存与同步 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 mcp_catalog: Any = None mcp_catalog_expires_at: float = 0.0 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 get_mcp_catalog(self, identity: str) -> Any: now = time.monotonic() with self._lock: entry = self._entries.get(identity) if entry is None or entry.mcp_catalog_expires_at <= now: if entry is not None: entry.mcp_catalog = None entry.mcp_catalog_expires_at = 0.0 return None self._entries.move_to_end(identity) return entry.mcp_catalog def set_mcp_catalog( self, identity: str, catalog: Any, *, ttl_seconds: float ) -> None: with self._lock: entry = self._entry_locked(identity) entry.mcp_catalog = catalog entry.mcp_catalog_expires_at = time.monotonic() + max(0.0, ttl_seconds) def invalidate_mcp_catalog(self, identity: str) -> None: with self._lock: entry = self._entries.get(identity) if entry is not None: entry.mcp_catalog = None entry.mcp_catalog_expires_at = 0.0 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)