mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-09-07 00:16:57 +08:00
fix: 保持搜索线程任务的请求上下文 (#6451)
This commit is contained in:
@@ -13,7 +13,7 @@ from threading import Lock
|
|||||||
from typing import Dict, Optional
|
from typing import Dict, Optional
|
||||||
|
|
||||||
from app.foundation.singleton import Singleton
|
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.log import logger
|
||||||
from app.runtime.reload import ConfigReloadMixin
|
from app.runtime.reload import ConfigReloadMixin
|
||||||
from app.runtime.settings import get_runtime_setting
|
from app.runtime.settings import get_runtime_setting
|
||||||
@@ -74,7 +74,7 @@ def enable_doh(enable: bool) -> bool:
|
|||||||
executor = _get_executor_locked()
|
executor = _get_executor_locked()
|
||||||
# 一次解析的任务必须在同一临界区提交完,避免关闭过程中部分任务落入新线程池
|
# 一次解析的任务必须在同一临界区提交完,避免关闭过程中部分任务落入新线程池
|
||||||
futures = [
|
futures = [
|
||||||
executor.submit(_doh_query, resolver, host)
|
submit_with_context(executor, _doh_query, resolver, host)
|
||||||
for resolver in _doh_setting("DOH_RESOLVERS").split(",")
|
for resolver in _doh_setting("DOH_RESOLVERS").split(",")
|
||||||
]
|
]
|
||||||
for future in as_completed(futures):
|
for future in as_completed(futures):
|
||||||
|
|||||||
+14
-6
@@ -11,7 +11,7 @@ from typing import AsyncIterator, Any, Awaitable, Callable, Dict, Iterable, Tupl
|
|||||||
from typing import List, Optional
|
from typing import List, Optional
|
||||||
from unicodedata import normalize
|
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 import ChainBase
|
||||||
from app.chain.media import MediaChain
|
from app.chain.media import MediaChain
|
||||||
from app.runtime.config import global_vars
|
from app.runtime.config import global_vars
|
||||||
@@ -1291,7 +1291,11 @@ class SearchChain(ChainBase):
|
|||||||
)
|
)
|
||||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||||
all_tasks = {
|
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 site_key, site_torrent_list in site_torrents.items()
|
||||||
}
|
}
|
||||||
for future in as_completed(all_tasks):
|
for future in as_completed(all_tasks):
|
||||||
@@ -2365,10 +2369,14 @@ class SearchChain(ChainBase):
|
|||||||
search_page = search_pages[page_index]
|
search_page = search_pages[page_index]
|
||||||
# 关键字已按 area 统一解析(imdbid 场景使用 imdb 标识),站点调用无需再分支
|
# 关键字已按 area 统一解析(imdbid 场景使用 imdb 标识),站点调用无需再分支
|
||||||
search_keyword = mediainfo.imdb_id if area == "imdbid" and mediainfo else keyword
|
search_keyword = mediainfo.imdb_id if area == "imdbid" and mediainfo else keyword
|
||||||
task = executor.submit(self.search_site_torrents, site=site,
|
task = submit_with_context(
|
||||||
keyword=search_keyword,
|
executor,
|
||||||
mtype=mediainfo.type if mediainfo else mtype,
|
self.search_site_torrents,
|
||||||
page=search_page)
|
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)
|
pending_tasks[task] = (site, page_index, search_page, search_keyword)
|
||||||
|
|
||||||
for site in indexer_sites:
|
for site in indexer_sites:
|
||||||
|
|||||||
@@ -2,8 +2,8 @@ import asyncio
|
|||||||
import inspect
|
import inspect
|
||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
from concurrent.futures import Future, ThreadPoolExecutor, wait
|
from concurrent.futures import Executor, Future, ThreadPoolExecutor, wait
|
||||||
from contextvars import copy_context
|
from contextvars import Context, copy_context
|
||||||
from functools import partial, wraps
|
from functools import partial, wraps
|
||||||
from typing import Any, Callable, TypeVar, cast
|
from typing import Any, Callable, TypeVar, cast
|
||||||
|
|
||||||
@@ -15,6 +15,22 @@ TaskResult = TypeVar("TaskResult")
|
|||||||
ExecutorResult = TypeVar("ExecutorResult")
|
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):
|
class OwnedThreadPoolExecutor(ThreadPoolExecutor):
|
||||||
"""
|
"""
|
||||||
追踪已接受 Future,并提供可重试的有界关闭合同。
|
追踪已接受 Future,并提供可重试的有界关闭合同。
|
||||||
|
|||||||
@@ -3,6 +3,7 @@ import threading
|
|||||||
import time
|
import time
|
||||||
|
|
||||||
from app.adapters.network import doh
|
from app.adapters.network import doh
|
||||||
|
from app.runtime.correlation import correlation_scope, get_correlation_id
|
||||||
from app.runtime.execution import OwnedThreadPoolExecutor
|
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 query_calls == [("resolver.test", "example.com")]
|
||||||
assert resolved_hosts == ["203.0.113.7", "203.0.113.7"]
|
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"]
|
||||||
|
|||||||
@@ -2,6 +2,9 @@
|
|||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import ast
|
import ast
|
||||||
|
import json
|
||||||
|
import subprocess
|
||||||
|
import sys
|
||||||
import threading
|
import threading
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
@@ -41,6 +44,42 @@ def test_host_uses_canonical_threadpool_boundary() -> None:
|
|||||||
assert violations == []
|
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:
|
def test_plugin_file_adapters_share_runtime_completion_contract() -> None:
|
||||||
"""市场与插件包适配器不得各自维护另一套线程取消实现。"""
|
"""市场与插件包适配器不得各自维护另一套线程取消实现。"""
|
||||||
assert market_adapter._await_thread_operation is run_in_threadpool_to_completion
|
assert market_adapter._await_thread_operation is run_in_threadpool_to_completion
|
||||||
|
|||||||
@@ -2,9 +2,11 @@ import asyncio
|
|||||||
from types import SimpleNamespace
|
from types import SimpleNamespace
|
||||||
from unittest.mock import AsyncMock, patch
|
from unittest.mock import AsyncMock, patch
|
||||||
|
|
||||||
|
from app.adapters.network.http import RequestUtils
|
||||||
from app.chain.search import SearchChain
|
from app.chain.search import SearchChain
|
||||||
from app.modules.indexer import IndexerModule
|
from app.modules.indexer import IndexerModule
|
||||||
from app.runtime.config import settings
|
from app.runtime.config import settings
|
||||||
|
from app.runtime.correlation import CORRELATION_ID_HEADER, correlation_scope
|
||||||
|
|
||||||
|
|
||||||
def make_chain() -> SearchChain:
|
def make_chain() -> SearchChain:
|
||||||
@@ -67,6 +69,46 @@ def test_search_invokes_plugin_once_with_multiple_indexers():
|
|||||||
assert len(results) == 3
|
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():
|
def test_async_search_returns_plugin_results_without_indexers():
|
||||||
"""异步搜索应支持只有插件资源源的部署方式。"""
|
"""异步搜索应支持只有插件资源源的部署方式。"""
|
||||||
chain = make_chain()
|
chain = make_chain()
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ from app.chain import search as search_module
|
|||||||
from app.chain.search import SearchChain
|
from app.chain.search import SearchChain
|
||||||
from app.domain.context import MediaInfo, TorrentInfo
|
from app.domain.context import MediaInfo, TorrentInfo
|
||||||
from app.schemas.types import MediaSource, MediaType
|
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):
|
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
|
assert len(contexts) == 2
|
||||||
media_chain.recognize_by_meta.assert_called_once()
|
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"]
|
||||||
|
|||||||
Reference in New Issue
Block a user