diff --git a/app/adapters/network/doh.py b/app/adapters/network/doh.py index 3ad3ebada..1200550ce 100644 --- a/app/adapters/network/doh.py +++ b/app/adapters/network/doh.py @@ -13,7 +13,7 @@ from threading import Lock from typing import Dict, Optional from app.foundation.singleton import Singleton -from app.runtime.execution import OwnedThreadPoolExecutor +from app.runtime.execution import OwnedThreadPoolExecutor, submit_with_context from app.runtime.log import logger from app.runtime.reload import ConfigReloadMixin from app.runtime.settings import get_runtime_setting @@ -74,7 +74,7 @@ def enable_doh(enable: bool) -> bool: executor = _get_executor_locked() # 一次解析的任务必须在同一临界区提交完，避免关闭过程中部分任务落入新线程池 futures = [ - executor.submit(_doh_query, resolver, host) + submit_with_context(executor, _doh_query, resolver, host) for resolver in _doh_setting("DOH_RESOLVERS").split(",") ] for future in as_completed(futures): diff --git a/app/chain/search.py b/app/chain/search.py index 3d3ea26f2..35af22c72 100644 --- a/app/chain/search.py +++ b/app/chain/search.py @@ -11,7 +11,7 @@ from typing import AsyncIterator, Any, Awaitable, Callable, Dict, Iterable, Tupl from typing import List, Optional from unicodedata import normalize -from app.runtime.execution import run_in_threadpool +from app.runtime.execution import run_in_threadpool, submit_with_context from app.chain import ChainBase from app.chain.media import MediaChain from app.runtime.config import global_vars @@ -1291,7 +1291,11 @@ class SearchChain(ChainBase): ) with ThreadPoolExecutor(max_workers=max_workers) as executor: all_tasks = { - executor.submit(__do_site_filter, site_torrent_list): site_key + submit_with_context( + executor, + __do_site_filter, + site_torrent_list, + ): site_key for site_key, site_torrent_list in site_torrents.items() } for future in as_completed(all_tasks): @@ -2365,10 +2369,14 @@ class SearchChain(ChainBase): search_page = search_pages[page_index] # 关键字已按 area 统一解析（imdbid 场景使用 imdb 标识），站点调用无需再分支 search_keyword = mediainfo.imdb_id if area == "imdbid" and mediainfo else keyword - task = executor.submit(self.search_site_torrents, site=site, - keyword=search_keyword, - mtype=mediainfo.type if mediainfo else mtype, - page=search_page) + task = submit_with_context( + executor, + self.search_site_torrents, + site=site, + keyword=search_keyword, + mtype=mediainfo.type if mediainfo else mtype, + page=search_page, + ) pending_tasks[task] = (site, page_index, search_page, search_keyword) for site in indexer_sites: diff --git a/app/runtime/execution.py b/app/runtime/execution.py index bc82deeaa..fa506e23c 100644 --- a/app/runtime/execution.py +++ b/app/runtime/execution.py @@ -2,8 +2,8 @@ import asyncio import inspect import threading import time -from concurrent.futures import Future, ThreadPoolExecutor, wait -from contextvars import copy_context +from concurrent.futures import Executor, Future, ThreadPoolExecutor, wait +from contextvars import Context, copy_context from functools import partial, wraps from typing import Any, Callable, TypeVar, cast @@ -15,6 +15,22 @@ TaskResult = TypeVar("TaskResult") ExecutorResult = TypeVar("ExecutorResult") +def submit_with_context( + executor: Executor, + func: Callable[..., ExecutorResult], + /, + *args: Any, + **kwargs: Any, +) -> Future[ExecutorResult]: + """从空线程上下文提交任务，并在执行时恢复调用方的独立上下文快照。""" + context = copy_context() + + def submit() -> Future[ExecutorResult]: + return executor.submit(context.run, func, *args, **kwargs) + + return Context().run(submit) + + class OwnedThreadPoolExecutor(ThreadPoolExecutor): """ 追踪已接受 Future，并提供可重试的有界关闭合同。 diff --git a/tests/test_doh_helper.py b/tests/test_doh_helper.py index 11a8dd89e..6f8c101da 100644 --- a/tests/test_doh_helper.py +++ b/tests/test_doh_helper.py @@ -3,6 +3,7 @@ import threading import time from app.adapters.network import doh +from app.runtime.correlation import correlation_scope, get_correlation_id from app.runtime.execution import OwnedThreadPoolExecutor @@ -144,3 +145,34 @@ def test_enable_doh_reuses_cached_host_resolution(monkeypatch): assert query_calls == [("resolver.test", "example.com")] assert resolved_hosts == ["203.0.113.7", "203.0.113.7"] + + +def test_doh_queries_use_each_request_context(monkeypatch): + """复用的 DoH worker 应按查询恢复关联 ID，不能丢失或粘住首个请求。""" + original_getaddrinfo = socket.getaddrinfo + observed = [] + helper = object.__new__(doh.DohHelper) + monkeypatch.setattr( + "app.runtime.config.settings.DOH_DOMAINS", + "first.example,second.example", + ) + monkeypatch.setattr("app.runtime.config.settings.DOH_RESOLVERS", "resolver.test") + monkeypatch.setattr( + doh, + "_doh_query", + lambda _resolver, _host: observed.append(get_correlation_id()) or "203.0.113.7", + ) + monkeypatch.setattr(doh, "_orig_getaddrinfo", lambda _host, *_args, **_kwargs: []) + + try: + assert helper.shutdown() is True + assert doh.enable_doh(True) is True + with correlation_scope("doh-first"): + socket.getaddrinfo("first.example", None) + with correlation_scope("doh-second"): + socket.getaddrinfo("second.example", None) + finally: + helper.shutdown() + socket.getaddrinfo = original_getaddrinfo + + assert observed == ["doh-first", "doh-second"] diff --git a/tests/test_runtime_execution.py b/tests/test_runtime_execution.py index 356238249..c4ae468bc 100644 --- a/tests/test_runtime_execution.py +++ b/tests/test_runtime_execution.py @@ -2,6 +2,9 @@ import asyncio import ast +import json +import subprocess +import sys import threading from pathlib import Path @@ -41,6 +44,42 @@ def test_host_uses_canonical_threadpool_boundary() -> None: assert violations == [] +@pytest.mark.parametrize("inherit_context", [0, 1]) +def test_submit_with_context_is_independent_of_thread_inheritance( + inherit_context: int, +) -> None: + """线程继承开关不得改变逐任务快照，worker 也不得保留首个请求状态。""" + script = """ +import json +from concurrent.futures import ThreadPoolExecutor +from contextvars import ContextVar + +from app.runtime.execution import submit_with_context + +request_id = ContextVar("request_id", default=None) +executor = ThreadPoolExecutor(max_workers=1) +observed = [] +for value in ("first", "second"): + token = request_id.set(value) + try: + observed.append(submit_with_context(executor, request_id.get).result()) + finally: + request_id.reset(token) +observed.append(executor.submit(request_id.get).result()) +executor.shutdown() +print(json.dumps(observed)) +""" + + completed = subprocess.run( + [sys.executable, "-X", f"thread_inherit_context={inherit_context}", "-c", script], + check=True, + capture_output=True, + text=True, + ) + + assert json.loads(completed.stdout) == ["first", "second", None] + + def test_plugin_file_adapters_share_runtime_completion_contract() -> None: """市场与插件包适配器不得各自维护另一套线程取消实现。""" assert market_adapter._await_thread_operation is run_in_threadpool_to_completion diff --git a/tests/test_search_plugin_sources.py b/tests/test_search_plugin_sources.py index e45aff69e..fd57ac7ad 100644 --- a/tests/test_search_plugin_sources.py +++ b/tests/test_search_plugin_sources.py @@ -2,9 +2,11 @@ import asyncio from types import SimpleNamespace from unittest.mock import AsyncMock, patch +from app.adapters.network.http import RequestUtils from app.chain.search import SearchChain from app.modules.indexer import IndexerModule from app.runtime.config import settings +from app.runtime.correlation import CORRELATION_ID_HEADER, correlation_scope def make_chain() -> SearchChain: @@ -67,6 +69,46 @@ def test_search_invokes_plugin_once_with_multiple_indexers(): assert len(results) == 3 +def test_sync_site_search_propagates_request_context(): + """同步站点 worker 的出站请求头应保留触发搜索的关联 ID。""" + chain = make_chain() + observed_headers = [] + chain.search_plugin_torrents = lambda **_kwargs: [] + + def search_site_torrents(**_kwargs): + RequestUtils().get_res("https://indexer.example/search") + return [] + + def request(_method, _url, **kwargs): + observed_headers.append(kwargs["headers"]) + return object() + + chain.search_site_torrents = search_site_torrents + + with ( + patch.object(settings, "SEARCH_RESOURCE_PAGES", 1, create=True), + patch("app.adapters.network.http.requests.request", side_effect=request), + patch("app.chain.search.get_configured_system_config") as system_config_oper, + patch("app.chain.search.SitesHelper") as sites_helper, + patch("app.chain.search.ProgressHelper") as progress_helper, + ): + system_config_oper.return_value.get.return_value = [1, 2] + sites_helper.return_value.get_indexers.return_value = [ + {"id": 1, "name": "站点一"}, + {"id": 2, "name": "站点二"}, + ] + progress_helper.return_value = SimpleNamespace( + start=lambda: None, update=lambda **_kwargs: None, end=lambda: None + ) + with correlation_scope("search-request"): + chain._SearchChain__search_all_sites(keyword="keyword") + + assert [headers[CORRELATION_ID_HEADER] for headers in observed_headers] == [ + "search-request", + "search-request", + ] + + def test_async_search_returns_plugin_results_without_indexers(): """异步搜索应支持只有插件资源源的部署方式。""" chain = make_chain() diff --git a/tests/test_search_same_name_disambiguation.py b/tests/test_search_same_name_disambiguation.py index 9d2d6ecac..9823659ab 100644 --- a/tests/test_search_same_name_disambiguation.py +++ b/tests/test_search_same_name_disambiguation.py @@ -4,6 +4,7 @@ from app.chain import search as search_module from app.chain.search import SearchChain from app.domain.context import MediaInfo, TorrentInfo from app.schemas.types import MediaSource, MediaType +from app.runtime.correlation import correlation_scope, get_correlation_id def test_exact_search_rejects_no_year_alias_recognized_as_different_work(monkeypatch): @@ -102,3 +103,52 @@ def test_exact_search_reuses_disambiguation_for_same_parsed_title(monkeypatch): assert len(contexts) == 2 media_chain.recognize_by_meta.assert_called_once() + + +def test_parallel_filter_propagates_request_context(monkeypatch): + """按站点并行过滤时，每个 worker 都应保留触发解析的关联 ID。""" + target = MediaInfo( + media_source=MediaSource.TMDB, + media_id="1", + title="测试电影", + type=MediaType.MOVIE, + year="2024", + ) + torrents = [ + TorrentInfo( + site=index, + site_name=f"测试站点{index}", + title=f"测试电影 2024 1080p GROUP-{index}", + category=MediaType.MOVIE.value, + ) + for index in (1, 2) + ] + observed = [] + + def filter_torrent(_self, _torrent, _filter_params): + observed.append(get_correlation_id()) + return True + + monkeypatch.setattr(search_module.TorrentHelper, "filter_torrent", filter_torrent) + monkeypatch.setattr( + search_module.TorrentHelper, + "match_torrent", + staticmethod(lambda **_kwargs: True), + ) + monkeypatch.setattr( + search_module.TorrentHelper, + "sort_torrents", + staticmethod(lambda contexts: contexts), + ) + chain = object.__new__(SearchChain) + + with correlation_scope("filter-request"): + contexts = chain._SearchChain__parse_result( + torrents=torrents, + mediainfo=target, + rule_groups=[], + filter_params={"free": "true"}, + ) + + assert len(contexts) == 2 + assert observed == ["filter-request", "filter-request"]