songzhuoyuan
78 lines
2.8 KiB
Python
78 lines
2.8 KiB
Python
"""
|
||
统一的外部请求超时控制工具
|
||
|
||
- install_default_requests_timeout: 给所有基于 requests 的外部请求
|
||
（akshare、tushare、pywencai 等）注入默认超时，调用方未指定 timeout 时生效。
|
||
- call_with_timeout: 在独立线程中执行阻塞调用，超过时限立即放弃等待，
|
||
防止个别环节长时间卡死整个程序。
|
||
"""
|
||
|
||
import os
|
||
import threading
|
||
|
||
# 默认连接超时 / 读取超时（单位：秒），可通过环境变量覆盖
|
||
DEFAULT_CONNECT_TIMEOUT = float(os.getenv("HTTP_CONNECT_TIMEOUT", "6"))
|
||
DEFAULT_READ_TIMEOUT = float(os.getenv("HTTP_READ_TIMEOUT", "20"))
|
||
DEFAULT_CALL_TIMEOUT = float(os.getenv("HTTP_CALL_TIMEOUT", "30"))
|
||
|
||
_installed = False
|
||
_install_lock = threading.Lock()
|
||
|
||
|
||
def install_default_requests_timeout(connect_timeout=None, read_timeout=None):
|
||
"""给 requests 库注入默认超时（仅当调用方未显式指定 timeout 时生效）。"""
|
||
global _installed
|
||
with _install_lock:
|
||
if _installed:
|
||
return
|
||
|
||
connect = DEFAULT_CONNECT_TIMEOUT if connect_timeout is None else connect_timeout
|
||
read = DEFAULT_READ_TIMEOUT if read_timeout is None else read_timeout
|
||
|
||
try:
|
||
import requests.sessions
|
||
original_request = requests.sessions.Session.request
|
||
|
||
def request_with_timeout(self, method, url, **kwargs):
|
||
if kwargs.get("timeout") is None:
|
||
kwargs["timeout"] = (connect, read)
|
||
return original_request(self, method, url, **kwargs)
|
||
|
||
requests.sessions.Session.request = request_with_timeout
|
||
_installed = True
|
||
print(f"✅ 已为外部请求注入默认超时（连接 {connect}s / 读取 {read}s）")
|
||
except Exception as e:
|
||
print(f"⚠️ 注入 requests 默认超时失败: {e}")
|
||
|
||
|
||
def call_with_timeout(func, timeout=None, *args, **kwargs):
|
||
"""在线程中执行函数，超过 timeout 秒则放弃等待并抛出 TimeoutError。
|
||
|
||
注意：超时后线程会继续在后台运行直至结束（守护线程），不会影响主程序。
|
||
"""
|
||
if timeout is None:
|
||
timeout = DEFAULT_CALL_TIMEOUT
|
||
if timeout <= 0:
|
||
return func(*args, **kwargs)
|
||
|
||
result_box = {}
|
||
|
||
def _run():
|
||
try:
|
||
result_box["ok"] = True
|
||
result_box["value"] = func(*args, **kwargs)
|
||
except BaseException as e: # noqa: BLE001
|
||
result_box["ok"] = False
|
||
result_box["error"] = e
|
||
|
||
worker = threading.Thread(target=_run, daemon=True)
|
||
worker.start()
|
||
worker.join(timeout)
|
||
|
||
if worker.is_alive():
|
||
func_name = getattr(func, "__name__", "function")
|
||
raise TimeoutError(f"调用 {func_name} 超过 {timeout}s 未返回，已放弃等待")
|
||
if result_box.get("ok"):
|
||
return result_box["value"]
|
||
raise result_box.get("error", RuntimeError("未知错误"))
|