"""Source fallback execution with timeout and normalized errors.""" from __future__ import annotations import concurrent.futures import contextlib import io import os import random import threading import time from dataclasses import dataclass from typing import Any, Callable import pandas as pd from app.core.config import settings class SourceUnavailable(Exception): """Raised when one upstream source cannot provide usable data.""" class SourceDataInsufficient(SourceUnavailable): """Raised when a source returned data but it did not meet row-count/quality requirements.""" class AllSourcesFailed(Exception): """Raised when every upstream source failed and no stale cache exists.""" def __init__(self, endpoint: str, attempts: list[dict[str, Any]]) -> None: super().__init__(f"all sources failed for {endpoint}") self.endpoint = endpoint self.attempts = attempts @dataclass(frozen=True) class SourceResult: source: str data: Any row_count: int | None = None SourceCallable = Callable[[], Any] class GlobalSourcePool: """全局源调用池，控制并发数，防止 worker pool 饱和""" def __init__(self, max_concurrent: int = 2, max_workers: int = 4): self._max_concurrent = max_concurrent self._semaphore = threading.Semaphore(max_concurrent) self._executor = concurrent.futures.ThreadPoolExecutor( max_workers=max_workers, thread_name_prefix="source_pool" ) self._stats = { "active": 0, "queued": 0, "completed": 0, "failed": 0, "timeout": 0, } self._lock = threading.Lock() def get_stats(self) -> dict[str, Any]: with self._lock: return dict(self._stats) def increment_stat(self, key: str) -> None: with self._lock: if key in self._stats: self._stats[key] += 1 def call_with_timeout( self, func: SourceCallable, timeout_seconds: int, source_name: str | None = None, ) -> Any: """带全局并发控制的调用""" self.increment_stat("queued") # 获取信号量，限制并发数 with self._semaphore: self.increment_stat("queued") self.increment_stat("active") try: future = self._executor.submit(_invoke_source, func, source_name) try: result = future.result(timeout=timeout_seconds) self.increment_stat("completed") return result except concurrent.futures.TimeoutError as exc: future.cancel() self.increment_stat("timeout") raise SourceUnavailable(f"timeout after {timeout_seconds}s") from exc except Exception: self.increment_stat("failed") raise finally: self.increment_stat("active") except Exception: self.increment_stat("active") raise # 全局源调用池（懒初始化） _global_pool: GlobalSourcePool | None = None _pool_lock = threading.Lock() def _get_global_pool() -> GlobalSourcePool: """获取或创建全局源调用池""" global _global_pool if _global_pool is None: with _pool_lock: if _global_pool is None: # 根据 HF Space 的 worker 限制调整并发数 # Free tier: 1-2 workers, Pro: 10+ workers max_concurrent = int(os.getenv("MAX_CONCURRENT_SOURCES", "2")) max_workers = int(os.getenv("SOURCE_POOL_WORKERS", "4")) _global_pool = GlobalSourcePool( max_concurrent=max_concurrent, max_workers=max_workers ) return _global_pool def _count_rows(data: Any) -> int | None: if isinstance(data, (list, tuple)): return len(data) if isinstance(data, pd.DataFrame): return len(data) if isinstance(data, dict): for key in ("daily", "records", "reports", "notices", "items", "stocks", "list", "data"): value = data.get(key) if isinstance(value, (list, tuple)): return len(value) if hasattr(data, "__len__") and not isinstance(data, (str, bytes, dict)): return len(data) return None def _is_transient_error(exc: Exception) -> bool: if isinstance(exc, SourceDataInsufficient): return False if isinstance(exc, (SourceUnavailable, TimeoutError, OSError)): return True try: from requests.exceptions import ConnectionError, HTTPError, Timeout as RequestsTimeout if isinstance(exc, (RequestsTimeout, ConnectionError)): return True if isinstance(exc, HTTPError) and exc.response is not None: code = exc.response.status_code return code >= 500 or code == 429 except Exception: pass code = getattr(exc, "code", None) if isinstance(code, int) and (code >= 500 or code == 429): return True return False def _backoff_seconds(base: float, jitter: float, max_backoff: float, attempt_index: int) -> float: delay = min(base * (2 ** attempt_index), max_backoff) if jitter: delay += random.uniform(0, jitter) return delay def call_with_timeout(func: SourceCallable, timeout_seconds: int, source_name: str | None = None) -> Any: """带全局并发控制的调用""" pool = _get_global_pool() return pool.call_with_timeout(func, timeout_seconds, source_name) def get_source_pool_stats() -> dict[str, Any]: """获取源调用池统计信息""" pool = _get_global_pool() return pool.get_stats() def _invoke_source(func: SourceCallable, source_name: str | None = None) -> Any: # For AKShare sources we usually want to bypass any HTTP proxy, because # several Eastmoney/THS endpoints reject or distort proxied requests. disabled_proxy = ( settings.disable_proxy_for_akshare and source_name is not None and source_name.startswith("akshare.") ) proxy_keys = ("HTTP_PROXY", "HTTPS_PROXY", "http_proxy", "https_proxy") saved_proxy: dict[str, str] = {} if disabled_proxy: for key in proxy_keys: if key in os.environ: saved_proxy[key] = os.environ.pop(key) try: if not settings.suppress_source_output: return func() with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()): return func() finally: for key, value in saved_proxy.items(): os.environ[key] = value def run_sources( endpoint: str, sources: list[tuple[str, SourceCallable]], timeout_seconds: int, retry_attempts: int = 1, retry_backoff_seconds: float = 0.0, retry_max_backoff_seconds: float = 5.0, retry_jitter_seconds: float = 0.0, min_rows: int | None = None, circuit_retries: int = 0, ) -> tuple[SourceResult, list[dict[str, Any]]]: attempts: list[dict[str, Any]] = [] retry_attempts = max(1, min(int(retry_attempts or settings.source_retry_attempts), 3)) retry_backoff_seconds = float(retry_backoff_seconds or settings.source_retry_backoff_seconds) retry_max_backoff_seconds = float(retry_max_backoff_seconds or settings.source_retry_max_backoff_seconds) retry_jitter_seconds = float(retry_jitter_seconds or settings.source_retry_jitter_seconds) circuit_retries = max(0, int(circuit_retries or 0)) min_rows = int(min_rows) if min_rows is not None else None for circuit_pass in range(circuit_retries + 1): if circuit_pass > 0: time.sleep( _backoff_seconds( retry_backoff_seconds, retry_jitter_seconds, retry_max_backoff_seconds, circuit_pass + 1, ) ) pass_had_terminal_failure = False for source_name, func in sources: for attempt_number in range(1, retry_attempts + 1): started = time.perf_counter() try: data = call_with_timeout(func, timeout_seconds, source_name) elapsed_ms = int((time.perf_counter() - started) * 1000) if data is None: raise SourceUnavailable("empty result") row_count = _count_rows(data) if min_rows is not None and row_count is not None and row_count < min_rows: raise SourceDataInsufficient(f"source returned {row_count} rows, required {min_rows}") attempts.append({"source": source_name, "ok": True, "attempt": attempt_number, "elapsed_ms": elapsed_ms}) return SourceResult(source=source_name, data=data, row_count=row_count), attempts except SourceDataInsufficient as exc: elapsed_ms = int((time.perf_counter() - started) * 1000) attempts.append( { "source": source_name, "ok": False, "attempt": attempt_number, "elapsed_ms": elapsed_ms, "error": f"{type(exc).__name__}: {exc}", } ) pass_had_terminal_failure = True break except Exception as exc: elapsed_ms = int((time.perf_counter() - started) * 1000) attempts.append( { "source": source_name, "ok": False, "attempt": attempt_number, "elapsed_ms": elapsed_ms, "error": f"{type(exc).__name__}: {exc}", } ) if not _is_transient_error(exc): pass_had_terminal_failure = True break if attempt_number < retry_attempts: time.sleep( _backoff_seconds( retry_backoff_seconds, retry_jitter_seconds, retry_max_backoff_seconds, attempt_number, ) ) continue if not pass_had_terminal_failure: break raise AllSourcesFailed(endpoint=endpoint, attempts=attempts)