小钰
709ea0fd78
- 新增ssh_proxy_pool.py: 全局单例连接池，队列串行化SSH请求 - stockstats_utils.py: _fetch_via_ssh改用连接池 - polymarket.py: 跳过本地直接走SSH连接池 - 消除并发SSH资源竞争问题
124 lines
3.9 KiB
Python
124 lines
3.9 KiB
Python
"""SSH Proxy Connection Pool — 串行化所有SSH中转请求，避免并发资源竞争。
|
||
|
||
所有需要走中转的数据源统一调用 ssh_proxy_fetch(url)，
|
||
内部通过队列串行化SSH连接，避免多线程/多进程并发导致的连接冲突。
|
||
"""
|
||
import json
|
||
import logging
|
||
import subprocess
|
||
import threading
|
||
import time
|
||
from queue import Queue, Empty
|
||
from typing import Optional
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# SSH配置缓存
|
||
_ssh_config: Optional[tuple] = None
|
||
_config_lock = threading.Lock()
|
||
|
||
|
||
def _get_ssh_config() -> tuple:
|
||
"""获取SSH配置（host, port, password），带缓存。"""
|
||
global _ssh_config
|
||
if _ssh_config is not None:
|
||
return _ssh_config
|
||
with _config_lock:
|
||
if _ssh_config is not None:
|
||
return _ssh_config
|
||
from .stockstats_utils import _get_ssh_config as _load_config
|
||
_ssh_config = _load_config()
|
||
return _ssh_config
|
||
|
||
|
||
class SSHProxyPool:
|
||
"""SSH代理连接池 — 串行化所有SSH请求。"""
|
||
|
||
def __init__(self):
|
||
self._queue = Queue()
|
||
self._lock = threading.Lock()
|
||
self._initialized = False
|
||
|
||
def _ensure_initialized(self):
|
||
"""确保连接池已初始化。"""
|
||
if self._initialized:
|
||
return
|
||
with self._lock:
|
||
if self._initialized:
|
||
return
|
||
# 放一个占位符，表示可以发起请求
|
||
self._queue.put(True)
|
||
self._initialized = True
|
||
|
||
def fetch(self, url: str, timeout: int = 30) -> str:
|
||
"""通过SSH代理获取URL内容，串行化执行。
|
||
|
||
Args:
|
||
url: 要获取的URL
|
||
timeout: 超时时间（秒）
|
||
|
||
Returns:
|
||
URL响应内容字符串
|
||
|
||
Raises:
|
||
RuntimeError: SSH代理失败时抛出
|
||
"""
|
||
self._ensure_initialized()
|
||
|
||
# 等待轮到自己（最多等timeout秒）
|
||
try:
|
||
self._queue.get(timeout=timeout)
|
||
except Empty:
|
||
raise RuntimeError(f"SSH proxy queue timeout after {timeout}s")
|
||
|
||
try:
|
||
host, port, password = _get_ssh_config()
|
||
cmd = [
|
||
"sshpass", "-p", password,
|
||
"ssh", "-o", "StrictHostKeyChecking=no",
|
||
"-o", "ConnectTimeout=15",
|
||
"-o", "ServerAliveInterval=10",
|
||
"-p", str(port), f"root@{host}",
|
||
f'curl -sL --max-time {timeout} "{url}"'
|
||
]
|
||
result = subprocess.run(
|
||
cmd, capture_output=True, text=True,
|
||
timeout=timeout + 10 # 给SSH额外时间
|
||
)
|
||
|
||
if result.returncode != 0:
|
||
raise RuntimeError(f"SSH proxy failed (rc={result.returncode}): {result.stderr[:200]}")
|
||
|
||
if not result.stdout or not result.stdout.strip():
|
||
raise RuntimeError("SSH proxy returned empty response")
|
||
|
||
return result.stdout
|
||
|
||
except subprocess.TimeoutExpired:
|
||
raise RuntimeError(f"SSH proxy subprocess timeout after {timeout}s")
|
||
finally:
|
||
# 无论成功失败，都放回占位符，让下一个请求执行
|
||
self._queue.put(True)
|
||
|
||
def fetch_json(self, url: str, timeout: int = 30) -> dict:
|
||
"""通过SSH代理获取JSON响应。"""
|
||
text = self.fetch(url, timeout)
|
||
try:
|
||
return json.loads(text)
|
||
except json.JSONDecodeError as e:
|
||
raise RuntimeError(f"SSH proxy returned invalid JSON: {e}")
|
||
|
||
|
||
# 全局单例
|
||
_pool = SSHProxyPool()
|
||
|
||
|
||
def ssh_proxy_fetch(url: str, timeout: int = 30) -> str:
|
||
"""通过SSH代理获取URL内容（串行化）。"""
|
||
return _pool.fetch(url, timeout)
|
||
|
||
|
||
def ssh_proxy_fetch_json(url: str, timeout: int = 30) -> dict:
|
||
"""通过SSH代理获取JSON响应（串行化）。"""
|
||
return _pool.fetch_json(url, timeout)
|