app/services/source_runner.py

"""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)