mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-09-05 07:27:15 +08:00
refactor: unify doh executor shutdown
This commit is contained in:
jxxghp
+47
-20
@@ -3,25 +3,26 @@ doh函数的实现。
|
||||
author: https://github.com/C5H12O5/syno-videoinfo-plugin
|
||||
"""
|
||||
import base64
|
||||
import concurrent
|
||||
import concurrent.futures
|
||||
import json
|
||||
import socket
|
||||
import struct
|
||||
import urllib
|
||||
import urllib.request
|
||||
from concurrent.futures import as_completed
|
||||
from threading import Lock
|
||||
from typing import Dict, Optional
|
||||
|
||||
from app.foundation.singleton import Singleton
|
||||
from app.runtime.execution import OwnedThreadPoolExecutor
|
||||
from app.runtime.log import logger
|
||||
from app.runtime.reload import ConfigReloadMixin
|
||||
from app.runtime.settings import get_runtime_setting
|
||||
from app.foundation.singleton import Singleton
|
||||
|
||||
# DoH 关闭时需要释放线程池；保持惰性创建可避免未启用 DoH 时占用进程级资源
|
||||
_executor: Optional[concurrent.futures.ThreadPoolExecutor] = None
|
||||
_executor: Optional[OwnedThreadPoolExecutor] = None
|
||||
_executor_lock = Lock()
|
||||
_doh_enabled = False
|
||||
_DOH_EXECUTOR_STOP_TIMEOUT_SECONDS = 10.0
|
||||
|
||||
# 定义默认的DoH配置
|
||||
_doh_timeout = 5
|
||||
@@ -36,20 +37,23 @@ def _doh_setting(key: str):
|
||||
return get_runtime_setting(key)
|
||||
|
||||
|
||||
def _get_executor_locked() -> concurrent.futures.ThreadPoolExecutor:
|
||||
"""在持有执行器锁时按需获取 DoH 查询线程池"""
|
||||
def _get_executor_locked() -> OwnedThreadPoolExecutor:
|
||||
"""在持有执行器锁时按需获取 DoH 查询线程池。"""
|
||||
global _executor
|
||||
if _executor is None:
|
||||
_executor = concurrent.futures.ThreadPoolExecutor()
|
||||
_executor = OwnedThreadPoolExecutor()
|
||||
return _executor
|
||||
|
||||
|
||||
def enable_doh(enable: bool) -> None:
|
||||
def enable_doh(enable: bool) -> bool:
|
||||
"""
|
||||
对 socket.getaddrinfo 进行补丁
|
||||
对 socket.getaddrinfo 进行补丁。
|
||||
|
||||
:param enable: 是否启用 DoH 解析
|
||||
:return: 状态切换成功时返回 True；旧 executor 未收敛时返回 False
|
||||
"""
|
||||
|
||||
global _doh_enabled
|
||||
global _doh_enabled, _executor
|
||||
|
||||
def _patched_getaddrinfo(host: str, *args, **kwargs):
|
||||
"""
|
||||
@@ -73,19 +77,31 @@ def enable_doh(enable: bool) -> None:
|
||||
executor.submit(_doh_query, resolver, host)
|
||||
for resolver in _doh_setting("DOH_RESOLVERS").split(",")
|
||||
]
|
||||
for future in concurrent.futures.as_completed(futures):
|
||||
for future in as_completed(futures):
|
||||
ip = future.result()
|
||||
if ip is not None:
|
||||
logger.info(f"已解析 [{host}] 为 [{ip}]")
|
||||
with _doh_lock:
|
||||
_doh_cache[host] = ip
|
||||
# 关闭可能在查询等待期间恢复系统 DNS；关闭后的结果不得回填到下一轮配置。
|
||||
with _executor_lock:
|
||||
cache_allowed = _doh_enabled
|
||||
if cache_allowed:
|
||||
with _doh_lock:
|
||||
_doh_cache[host] = ip
|
||||
host = ip
|
||||
break
|
||||
return _orig_getaddrinfo(host, *args, **kwargs)
|
||||
|
||||
with _executor_lock:
|
||||
if enable and _executor is not None and not _executor.accepting:
|
||||
# 上一轮 shutdown 超时后必须继续持有原 owner；只有真实收敛才能替换执行器。
|
||||
if not _executor.shutdown_bounded(timeout=0.0):
|
||||
_doh_enabled = False
|
||||
socket.getaddrinfo = _orig_getaddrinfo
|
||||
return False
|
||||
_executor = None
|
||||
_doh_enabled = enable
|
||||
socket.getaddrinfo = _patched_getaddrinfo if enable else _orig_getaddrinfo
|
||||
return True
|
||||
|
||||
|
||||
class DohHelper(ConfigReloadMixin, metaclass=Singleton):
|
||||
@@ -101,29 +117,40 @@ class DohHelper(ConfigReloadMixin, metaclass=Singleton):
|
||||
def on_config_changed(self) -> None:
|
||||
"""配置变化时清理缓存并重新应用 DoH 状态。"""
|
||||
if not _doh_setting("DOH_ENABLE"):
|
||||
self.shutdown()
|
||||
if not self.shutdown():
|
||||
logger.error("DoH配置关闭后查询线程池未在预算内收敛")
|
||||
return
|
||||
with _doh_lock:
|
||||
# DOH配置有变动的情况下，清空缓存
|
||||
_doh_cache.clear()
|
||||
enable_doh(True)
|
||||
if not enable_doh(True):
|
||||
logger.error("DoH查询线程池尚未收敛，暂不重新启用")
|
||||
|
||||
def get_reload_name(self) -> str:
|
||||
"""返回 DoH 配置重载名称。"""
|
||||
return 'DoH'
|
||||
|
||||
def shutdown(self) -> None:
|
||||
"""恢复系统 DNS 并释放 DoH 查询线程池"""
|
||||
def shutdown(
|
||||
self,
|
||||
timeout: float = _DOH_EXECUTOR_STOP_TIMEOUT_SECONDS,
|
||||
) -> bool:
|
||||
"""
|
||||
恢复系统 DNS 并有限等待 DoH 查询线程池。
|
||||
|
||||
:param timeout: 等待已接受查询和 worker 终止的最长秒数
|
||||
:return: 查询线程池真实终止时返回 True，否则返回 False
|
||||
"""
|
||||
global _executor, _doh_enabled
|
||||
with _executor_lock:
|
||||
_doh_enabled = False
|
||||
socket.getaddrinfo = _orig_getaddrinfo
|
||||
executor = _executor
|
||||
_executor = None
|
||||
converged = executor is None or executor.shutdown_bounded(timeout=timeout)
|
||||
if converged and _executor is executor:
|
||||
_executor = None
|
||||
with _doh_lock:
|
||||
_doh_cache.clear()
|
||||
if executor:
|
||||
executor.shutdown(wait=True)
|
||||
return converged
|
||||
|
||||
|
||||
def _doh_query(resolver: str, host: str) -> Optional[str]:
|
||||
|
||||
@@ -1,15 +1,105 @@
|
||||
import asyncio
|
||||
import inspect
|
||||
import threading
|
||||
import time
|
||||
from concurrent.futures import Future, ThreadPoolExecutor, wait
|
||||
from contextvars import copy_context
|
||||
from functools import partial, wraps
|
||||
from typing import Any, Callable, TypeVar
|
||||
from typing import Any, Callable, TypeVar, cast
|
||||
|
||||
from app.schemas.exception import ImmediateException
|
||||
from anyio.to_thread import run_sync
|
||||
|
||||
|
||||
TaskResult = TypeVar("TaskResult")
|
||||
ExecutorResult = TypeVar("ExecutorResult")
|
||||
|
||||
|
||||
class OwnedThreadPoolExecutor(ThreadPoolExecutor):
|
||||
"""
|
||||
追踪已接受 Future，并提供可重试的有界关闭合同。
|
||||
|
||||
该 owner 不取消排队任务，保持 ``ThreadPoolExecutor.shutdown(wait=True)``
|
||||
的历史完成语义；区别仅在于调用方可以在预算耗尽后保留同一实例继续收敛。
|
||||
"""
|
||||
|
||||
def __init__(self, max_workers: int | None = None) -> None:
|
||||
"""初始化线程池、提交准入状态和 Future owner 集合。"""
|
||||
super().__init__(max_workers=max_workers)
|
||||
self._ownership_lock = threading.RLock()
|
||||
self._accepting = True
|
||||
self._owned_futures: set[Future[Any]] = set()
|
||||
|
||||
@property
|
||||
def accepting(self) -> bool:
|
||||
"""返回执行器是否仍允许提交新任务。"""
|
||||
with self._ownership_lock:
|
||||
return self._accepting
|
||||
|
||||
def submit(
|
||||
self,
|
||||
fn: Callable[..., ExecutorResult],
|
||||
/,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> Future[ExecutorResult]:
|
||||
"""提交任务并在其达到终态前保留 owner。"""
|
||||
with self._ownership_lock:
|
||||
future = super().submit(fn, *args, **kwargs)
|
||||
self._owned_futures.add(future)
|
||||
future.add_done_callback(self._discard_future)
|
||||
return future
|
||||
|
||||
def shutdown(
|
||||
self,
|
||||
wait: bool = True,
|
||||
*,
|
||||
cancel_futures: bool = False,
|
||||
) -> None:
|
||||
"""封口提交准入并保留标准库 shutdown 的调用语义。"""
|
||||
with self._ownership_lock:
|
||||
self._accepting = False
|
||||
# 先在锁内封口；真正等待必须在锁外进行，否则 worker 的完成回调无法释放 owner。
|
||||
super().shutdown(wait=False, cancel_futures=cancel_futures)
|
||||
if wait:
|
||||
super().shutdown(wait=True, cancel_futures=cancel_futures)
|
||||
|
||||
def _discard_future(self, future: Future[Any]) -> None:
|
||||
"""任务达到终态后释放 owner 记录。"""
|
||||
with self._ownership_lock:
|
||||
self._owned_futures.discard(future)
|
||||
|
||||
def shutdown_bounded(self, timeout: float) -> bool:
|
||||
"""
|
||||
封口新提交并有限等待全部已接受任务。
|
||||
|
||||
:param timeout: 等待 Future 达到终态的最长秒数
|
||||
:return: 所有任务与 worker 均已终止时返回 True，否则返回 False
|
||||
"""
|
||||
deadline = time.monotonic() + max(0.0, timeout)
|
||||
with self._ownership_lock:
|
||||
self._accepting = False
|
||||
# 不取消排队工作，保持历史 shutdown(wait=True) 的完成语义。
|
||||
super().shutdown(wait=False)
|
||||
owned_futures = tuple(self._owned_futures)
|
||||
if owned_futures:
|
||||
_, pending_futures = wait(
|
||||
owned_futures,
|
||||
timeout=max(0.0, deadline - time.monotonic()),
|
||||
)
|
||||
if pending_futures:
|
||||
return False
|
||||
# 标准库只提供无界 wait=True；Future 又会先标记完成再执行 done callback，
|
||||
# 因此封口后读取稳定 worker 集合，复用同一 deadline 做有限 join。
|
||||
worker_threads = tuple(cast(set[threading.Thread], self._threads))
|
||||
current_thread = threading.current_thread()
|
||||
for worker_thread in worker_threads:
|
||||
if worker_thread is current_thread:
|
||||
continue
|
||||
worker_thread.join(
|
||||
timeout=max(0.0, deadline - time.monotonic()),
|
||||
)
|
||||
return all(not worker_thread.is_alive() for worker_thread in worker_threads)
|
||||
|
||||
|
||||
async def await_task_to_terminal(
|
||||
|
||||
+8
-67
@@ -1,10 +1,9 @@
|
||||
import threading
|
||||
import time
|
||||
from concurrent.futures import Future, ThreadPoolExecutor, wait
|
||||
from concurrent.futures import Future
|
||||
from contextvars import copy_context
|
||||
from typing import Any, Callable, TypeVar, cast
|
||||
|
||||
from app.foundation.singleton import Singleton
|
||||
from app.runtime.execution import OwnedThreadPoolExecutor
|
||||
from app.runtime.settings import RuntimeSettingsCompat
|
||||
|
||||
settings = RuntimeSettingsCompat()
|
||||
@@ -12,67 +11,8 @@ settings = RuntimeSettingsCompat()
|
||||
_Result = TypeVar("_Result")
|
||||
_THREAD_POOL_STOP_TIMEOUT_SECONDS = 10.0
|
||||
|
||||
|
||||
class _OwnedThreadPoolExecutor(ThreadPoolExecutor):
|
||||
"""
|
||||
追踪所有提交入口的共享执行器，包括旧调用方直接使用的 ``pool.submit``。
|
||||
"""
|
||||
|
||||
def __init__(self, max_workers: int) -> None:
|
||||
"""初始化线程池及其 Future owner 集合。"""
|
||||
super().__init__(max_workers=max_workers)
|
||||
self._ownership_lock = threading.RLock()
|
||||
self._owned_futures: set[Future[Any]] = set()
|
||||
|
||||
def submit(
|
||||
self,
|
||||
fn: Callable[..., _Result],
|
||||
/,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> Future[_Result]:
|
||||
"""提交任务并在其达到终态前保留 owner。"""
|
||||
with self._ownership_lock:
|
||||
future = super().submit(fn, *args, **kwargs)
|
||||
self._owned_futures.add(future)
|
||||
future.add_done_callback(self._discard_future)
|
||||
return future
|
||||
|
||||
def _discard_future(self, future: Future[Any]) -> None:
|
||||
"""任务达到终态后释放 owner 记录。"""
|
||||
with self._ownership_lock:
|
||||
self._owned_futures.discard(future)
|
||||
|
||||
def shutdown_bounded(self, timeout: float) -> bool:
|
||||
"""
|
||||
封口新提交并有限等待全部已接受任务。
|
||||
|
||||
:param timeout: 等待 Future 达到终态的最长秒数
|
||||
:return: 所有任务与 worker 均已终止时返回 True，否则返回 False
|
||||
"""
|
||||
deadline = time.monotonic() + max(0.0, timeout)
|
||||
with self._ownership_lock:
|
||||
# 不取消排队工作，保持历史 shutdown(wait=True) 的完成语义。
|
||||
super().shutdown(wait=False)
|
||||
owned_futures = tuple(self._owned_futures)
|
||||
if owned_futures:
|
||||
_, pending_futures = wait(
|
||||
owned_futures,
|
||||
timeout=max(0.0, deadline - time.monotonic()),
|
||||
)
|
||||
if pending_futures:
|
||||
return False
|
||||
# 标准库只提供无界 wait=True；Future 又会先标记完成再执行 done callback，
|
||||
# 因此封口后读取稳定 worker 集合，复用同一 deadline 做有限 join。
|
||||
worker_threads = tuple(cast(set[threading.Thread], self._threads))
|
||||
current_thread = threading.current_thread()
|
||||
for worker_thread in worker_threads:
|
||||
if worker_thread is current_thread:
|
||||
continue
|
||||
worker_thread.join(
|
||||
timeout=max(0.0, deadline - time.monotonic()),
|
||||
)
|
||||
return all(not worker_thread.is_alive() for worker_thread in worker_threads)
|
||||
# 阶段 64 曾在本模块引入该私有名；保留精确别名，避免测试或外部诊断代码失效。
|
||||
_OwnedThreadPoolExecutor = OwnedThreadPoolExecutor
|
||||
|
||||
|
||||
# strict mypy 跳过 foundation 实现导入，因此无法在本文件解析既有 Singleton 元类类型。
|
||||
@@ -83,7 +23,7 @@ class ThreadHelper(metaclass=Singleton): # type: ignore[metaclass]
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""按系统配置创建共享后台线程池。"""
|
||||
self.pool = _OwnedThreadPoolExecutor(max_workers=settings.CONF.threadpool)
|
||||
self.pool = OwnedThreadPoolExecutor(max_workers=settings.CONF.threadpool)
|
||||
|
||||
def submit(
|
||||
self,
|
||||
@@ -99,7 +39,8 @@ class ThreadHelper(metaclass=Singleton): # type: ignore[metaclass]
|
||||
:return: future
|
||||
"""
|
||||
context = copy_context()
|
||||
return self.pool.submit(context.run, func, *args, **kwargs)
|
||||
# strict mypy 跳过 execution 实现导入，需要在兼容门面恢复 Future 泛型。
|
||||
return cast(Future[_Result], self.pool.submit(context.run, func, *args, **kwargs))
|
||||
|
||||
def shutdown(
|
||||
self,
|
||||
@@ -111,4 +52,4 @@ class ThreadHelper(metaclass=Singleton): # type: ignore[metaclass]
|
||||
:param timeout: 等待已接受任务达到终态的最长秒数
|
||||
:return: 全部任务和 worker 均已终止时返回 True，否则返回 False
|
||||
"""
|
||||
return self.pool.shutdown_bounded(timeout=timeout)
|
||||
return bool(self.pool.shutdown_bounded(timeout=timeout))
|
||||
|
||||
Reference in New Issue
Block a user