content-find-agent-new/content_agent/integrations/bounded_pool.py

"""有界并发执行 + daemon 线程池(修永久卡死).

两件事:
1. `DaemonThreadPoolExecutor`:worker 线程 daemon 化。即使某 worker 卡在 socket(在 read
   超时触发前),也不会阻止解释器退出——配合各外部调用的 read 短超时,卡死 worker 必在有限
   时间内自终,绝不成为永久僵尸。标准库 ThreadPoolExecutor 不暴露 daemon 选项,这里复刻
   `_adjust_thread_count` 在 start 前置 daemon;若标准库内部结构变动则回退标准行为。
2. `run_bounded`:逐 future `result(timeout=)` 收割,**单条超时/异常 → on_timeout 占位、跳过、
   不抛、不中止整批**;收尾 `shutdown(wait=False, cancel_futures=True)`,绝不隐式 `wait=True`
   死等。结果按 offset 归位,与完成顺序无关(并发=串行产物一致)。
"""

from __future__ import annotations

import threading
import weakref
from concurrent.futures import ThreadPoolExecutor
from concurrent.futures import thread as _cf_thread
from concurrent.futures import TimeoutError as FutureTimeoutError
from typing import Any, Callable, TypeVar

T = TypeVar("T")


class DaemonThreadPoolExecutor(ThreadPoolExecutor):
    def _adjust_thread_count(self) -> None:  # noqa: D401 - 复刻标准库,仅令线程 daemon
        try:
            if self._idle_semaphore.acquire(timeout=0):  # type: ignore[attr-defined]
                return

            def weakref_cb(_: Any, q: Any = self._work_queue) -> None:  # type: ignore[attr-defined]
                q.put(None)

            num_threads = len(self._threads)  # type: ignore[attr-defined]
            if num_threads >= self._max_workers:  # type: ignore[attr-defined]
                return
            thread_name = "%s_%d" % (self._thread_name_prefix or self, num_threads)  # type: ignore[attr-defined]
            t = threading.Thread(
                name=thread_name,
                target=_cf_thread._worker,  # type: ignore[attr-defined]
                args=(
                    weakref.ref(self, weakref_cb),
                    self._work_queue,  # type: ignore[attr-defined]
                    self._initializer,  # type: ignore[attr-defined]
                    self._initargs,  # type: ignore[attr-defined]
                ),
                daemon=True,
            )
            t.start()
            self._threads.add(t)  # type: ignore[attr-defined]
            _cf_thread._threads_queues[t] = self._work_queue  # type: ignore[attr-defined]
        except Exception:
            # 标准库内部结构与预期不符 → 退回标准行为(read 短超时仍保证 worker 有限时间自终)。
            super()._adjust_thread_count()


def run_bounded(
    items: list[T],
    work_fn: Callable[[T], Any],
    *,
    max_workers: int,
    per_future_timeout: float,
    on_timeout: Callable[[T, int], Any],
    thread_name_prefix: str = "bounded",
) -> list[Any]:
    """并发执行 work_fn(item),逐条 result(timeout=per_future_timeout)。

    单条超时/未兜住的异常 → on_timeout(item, offset) 占位,不抛、不中止整批。
    返回与 items 同序、同长的结果 list。
    """
    results: list[Any] = [None] * len(items)
    if not items:
        return results
    workers = max(1, min(int(max_workers), len(items)))
    executor = DaemonThreadPoolExecutor(max_workers=workers, thread_name_prefix=thread_name_prefix)
    try:
        future_to_offset = {
            executor.submit(work_fn, item): offset for offset, item in enumerate(items)
        }
        for future, offset in future_to_offset.items():
            try:
                results[offset] = future.result(timeout=per_future_timeout)
            except FutureTimeoutError:
                results[offset] = on_timeout(items[offset], offset)
            except Exception:
                # worker 内部未兜住的意外 → 也记占位,绝不让 result() 炸主线程。
                results[offset] = on_timeout(items[offset], offset)
    finally:
        executor.shutdown(wait=False, cancel_futures=True)
    return results