feat: SSH连接池串行化 + Polymarket/数据源统一走连接池
- 新增ssh_proxy_pool.py: 全局单例连接池,队列串行化SSH请求 - stockstats_utils.py: _fetch_via_ssh改用连接池 - polymarket.py: 跳过本地直接走SSH连接池 - 消除并发SSH资源竞争问题
This commit is contained in:
1 parent
5207a2bf18
commit
709ea0fd78
3 files changed
+131
-45
No files matched your search
@@ -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)
|
||||
Reference in new issue
Block a user