feat: SSH连接池串行化 + Polymarket/数据源统一走连接池

- 新增ssh_proxy_pool.py: 全局单例连接池,队列串行化SSH请求
- stockstats_utils.py: _fetch_via_ssh改用连接池
- polymarket.py: 跳过本地直接走SSH连接池
- 消除并发SSH资源竞争问题
This commit is contained in:
小钰 committed 2026-07-08 14:34:41 +08:00
1 parent 5207a2bf18
commit 709ea0fd78
3 files changed
+131 -45

No files matched your search

+123
View File
@@ -0,0 +1,123 @@
"""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)