refactor: unify doh executor shutdown

app/runtime/execution.py

@@ -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(