mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-09-06 07:56:52 +08:00
fix: bind browser sessions to dedicated worker threads
This commit is contained in:
+259
-32
@@ -2,6 +2,7 @@ import ipaddress
|
|||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
import uuid
|
import uuid
|
||||||
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from typing import Any, Callable, Optional, Protocol
|
from typing import Any, Callable, Optional, Protocol
|
||||||
from urllib.parse import urlparse
|
from urllib.parse import urlparse
|
||||||
@@ -144,8 +145,11 @@ class BrowserSessionHelper:
|
|||||||
PRIVATE_HOST_SUFFIXES = (".localhost", ".local", ".lan", ".home", ".internal")
|
PRIVATE_HOST_SUFFIXES = (".localhost", ".local", ".lan", ".home", ".internal")
|
||||||
PRIVATE_HOSTNAMES = {"localhost", "ip6-localhost", "ip6-loopback"}
|
PRIVATE_HOSTNAMES = {"localhost", "ip6-localhost", "ip6-loopback"}
|
||||||
REF_ATTRIBUTE = "data-moviepilot-agent-ref"
|
REF_ATTRIBUTE = "data-moviepilot-agent-ref"
|
||||||
|
SESSION_WORKER_NAME_PREFIX = "browser-session"
|
||||||
|
|
||||||
_sessions: dict[str, _BrowserSessionState] = {}
|
_sessions: dict[str, _BrowserSessionState] = {}
|
||||||
|
_session_executors: dict[str, ThreadPoolExecutor] = {}
|
||||||
|
_session_thread_ids: dict[str, int] = {}
|
||||||
_sessions_lock = threading.RLock()
|
_sessions_lock = threading.RLock()
|
||||||
|
|
||||||
def __init__(self, headless: bool = True, viewport: Optional[dict[str, int]] = None):
|
def __init__(self, headless: bool = True, viewport: Optional[dict[str, int]] = None):
|
||||||
@@ -211,7 +215,9 @@ class BrowserSessionHelper:
|
|||||||
关闭所有 Agent 浏览器会话。
|
关闭所有 Agent 浏览器会话。
|
||||||
"""
|
"""
|
||||||
with cls._sessions_lock:
|
with cls._sessions_lock:
|
||||||
session_keys = list(cls._sessions.keys())
|
session_keys = list(
|
||||||
|
set(cls._sessions.keys()) | set(cls._session_executors.keys())
|
||||||
|
)
|
||||||
for session_key in session_keys:
|
for session_key in session_keys:
|
||||||
cls.close_session(session_key)
|
cls.close_session(session_key)
|
||||||
|
|
||||||
@@ -223,12 +229,27 @@ class BrowserSessionHelper:
|
|||||||
:param session_key: 会话标识
|
:param session_key: 会话标识
|
||||||
:return: 找到并关闭会话时返回 True
|
:return: 找到并关闭会话时返回 True
|
||||||
"""
|
"""
|
||||||
with cls._sessions_lock:
|
if cls._is_current_session_thread(session_key):
|
||||||
session = cls._sessions.pop(session_key, None)
|
closed = cls._close_session_in_thread(session_key)
|
||||||
if not session:
|
cls._shutdown_session_executor(session_key, wait=False)
|
||||||
return False
|
return closed
|
||||||
cls._close_session_state(session)
|
|
||||||
return True
|
executor = cls._get_existing_session_executor(session_key)
|
||||||
|
if executor:
|
||||||
|
future = executor.submit(
|
||||||
|
cls._run_session_task,
|
||||||
|
session_key,
|
||||||
|
cls._close_session_in_thread,
|
||||||
|
session_key,
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
return future.result()
|
||||||
|
finally:
|
||||||
|
cls._shutdown_session_executor(session_key)
|
||||||
|
|
||||||
|
closed = cls._close_session_in_thread(session_key)
|
||||||
|
cls._shutdown_session_executor(session_key)
|
||||||
|
return closed
|
||||||
|
|
||||||
def with_session(
|
def with_session(
|
||||||
self,
|
self,
|
||||||
@@ -249,6 +270,34 @@ class BrowserSessionHelper:
|
|||||||
:return: 回调函数返回值
|
:return: 回调函数返回值
|
||||||
"""
|
"""
|
||||||
self._prune_sessions()
|
self._prune_sessions()
|
||||||
|
return self._run_in_session_thread(
|
||||||
|
session_key,
|
||||||
|
self._with_session_in_thread,
|
||||||
|
session_key,
|
||||||
|
callback,
|
||||||
|
user_agent=user_agent,
|
||||||
|
cookies=cookies,
|
||||||
|
timeout=timeout,
|
||||||
|
)
|
||||||
|
|
||||||
|
def _with_session_in_thread(
|
||||||
|
self,
|
||||||
|
session_key: str,
|
||||||
|
callback: Callable[[_BrowserSessionState], Any],
|
||||||
|
user_agent: Optional[str] = None,
|
||||||
|
cookies: Optional[str] = None,
|
||||||
|
timeout: Optional[int] = 30,
|
||||||
|
) -> Any:
|
||||||
|
"""
|
||||||
|
在会话专属线程内获取浏览器会话并执行回调。
|
||||||
|
|
||||||
|
:param session_key: 会话标识
|
||||||
|
:param callback: 使用浏览器会话执行操作的回调函数
|
||||||
|
:param user_agent: 新建会话时使用的 User-Agent
|
||||||
|
:param cookies: 本次操作要注入的 Cookie 请求头
|
||||||
|
:param timeout: 默认操作超时时间,单位秒
|
||||||
|
:return: 回调函数返回值
|
||||||
|
"""
|
||||||
session = self._get_or_create_session(
|
session = self._get_or_create_session(
|
||||||
session_key=session_key,
|
session_key=session_key,
|
||||||
user_agent=user_agent,
|
user_agent=user_agent,
|
||||||
@@ -263,6 +312,164 @@ class BrowserSessionHelper:
|
|||||||
session.active_page.set_extra_http_headers({"cookie": cookies})
|
session.active_page.set_extra_http_headers({"cookie": cookies})
|
||||||
return callback(session)
|
return callback(session)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def _run_in_session_thread(
|
||||||
|
cls,
|
||||||
|
session_key: str,
|
||||||
|
callback: Callable[..., Any],
|
||||||
|
*args: Any,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> Any:
|
||||||
|
"""
|
||||||
|
将浏览器同步 API 调用投递到当前会话固定的单线程执行器。
|
||||||
|
|
||||||
|
:param session_key: 会话标识
|
||||||
|
:param callback: 需要在会话线程中运行的回调
|
||||||
|
:param args: 回调位置参数
|
||||||
|
:param kwargs: 回调关键字参数
|
||||||
|
:return: 回调返回值
|
||||||
|
"""
|
||||||
|
if cls._is_current_session_thread(session_key):
|
||||||
|
return callback(*args, **kwargs)
|
||||||
|
|
||||||
|
for _ in range(2):
|
||||||
|
executor = cls._get_session_executor(session_key)
|
||||||
|
try:
|
||||||
|
future = executor.submit(
|
||||||
|
cls._run_session_task,
|
||||||
|
session_key,
|
||||||
|
callback,
|
||||||
|
*args,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
except RuntimeError:
|
||||||
|
cls._discard_session_executor(session_key, executor)
|
||||||
|
continue
|
||||||
|
try:
|
||||||
|
return future.result()
|
||||||
|
except Exception:
|
||||||
|
with cls._sessions_lock:
|
||||||
|
has_session = session_key in cls._sessions
|
||||||
|
if not has_session:
|
||||||
|
cls._shutdown_session_executor(session_key)
|
||||||
|
raise
|
||||||
|
|
||||||
|
raise RuntimeError("浏览器会话线程已关闭")
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def _run_session_task(
|
||||||
|
cls,
|
||||||
|
session_key: str,
|
||||||
|
callback: Callable[..., Any],
|
||||||
|
*args: Any,
|
||||||
|
**kwargs: Any,
|
||||||
|
) -> Any:
|
||||||
|
"""
|
||||||
|
记录当前会话工作线程后执行实际任务。
|
||||||
|
|
||||||
|
:param session_key: 会话标识
|
||||||
|
:param callback: 需要执行的回调
|
||||||
|
:param args: 回调位置参数
|
||||||
|
:param kwargs: 回调关键字参数
|
||||||
|
:return: 回调返回值
|
||||||
|
"""
|
||||||
|
with cls._sessions_lock:
|
||||||
|
cls._session_thread_ids[session_key] = threading.get_ident()
|
||||||
|
return callback(*args, **kwargs)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def _is_current_session_thread(cls, session_key: str) -> bool:
|
||||||
|
"""
|
||||||
|
判断当前代码是否已运行在指定会话的固定线程内。
|
||||||
|
|
||||||
|
:param session_key: 会话标识
|
||||||
|
:return: 当前线程是会话工作线程时返回 True
|
||||||
|
"""
|
||||||
|
with cls._sessions_lock:
|
||||||
|
thread_id = cls._session_thread_ids.get(session_key)
|
||||||
|
return thread_id == threading.get_ident()
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def _get_session_executor(cls, session_key: str) -> ThreadPoolExecutor:
|
||||||
|
"""
|
||||||
|
获取或创建指定会话的单线程执行器。
|
||||||
|
|
||||||
|
:param session_key: 会话标识
|
||||||
|
:return: 会话专属执行器
|
||||||
|
"""
|
||||||
|
with cls._sessions_lock:
|
||||||
|
executor = cls._session_executors.get(session_key)
|
||||||
|
if executor:
|
||||||
|
return executor
|
||||||
|
executor = ThreadPoolExecutor(
|
||||||
|
max_workers=1,
|
||||||
|
thread_name_prefix=cls.SESSION_WORKER_NAME_PREFIX,
|
||||||
|
)
|
||||||
|
cls._session_executors[session_key] = executor
|
||||||
|
return executor
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def _get_existing_session_executor(
|
||||||
|
cls, session_key: str
|
||||||
|
) -> Optional[ThreadPoolExecutor]:
|
||||||
|
"""
|
||||||
|
获取指定会话已存在的执行器。
|
||||||
|
|
||||||
|
:param session_key: 会话标识
|
||||||
|
:return: 已存在的执行器,不存在时返回 None
|
||||||
|
"""
|
||||||
|
with cls._sessions_lock:
|
||||||
|
return cls._session_executors.get(session_key)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def _discard_session_executor(
|
||||||
|
cls,
|
||||||
|
session_key: str,
|
||||||
|
executor: ThreadPoolExecutor,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
丢弃已经关闭的会话执行器。
|
||||||
|
|
||||||
|
:param session_key: 会话标识
|
||||||
|
:param executor: 需要从缓存中移除的执行器
|
||||||
|
"""
|
||||||
|
with cls._sessions_lock:
|
||||||
|
if cls._session_executors.get(session_key) is executor:
|
||||||
|
cls._session_executors.pop(session_key, None)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def _shutdown_session_executor(
|
||||||
|
cls,
|
||||||
|
session_key: str,
|
||||||
|
wait: bool = True,
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
关闭并移除指定会话的执行器。
|
||||||
|
|
||||||
|
:param session_key: 会话标识
|
||||||
|
:param wait: 是否等待工作线程退出
|
||||||
|
"""
|
||||||
|
with cls._sessions_lock:
|
||||||
|
executor = cls._session_executors.pop(session_key, None)
|
||||||
|
cls._session_thread_ids.pop(session_key, None)
|
||||||
|
if executor:
|
||||||
|
executor.shutdown(wait=wait, cancel_futures=True)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def _close_session_in_thread(cls, session_key: str) -> bool:
|
||||||
|
"""
|
||||||
|
在会话固定线程内关闭并移除浏览器会话。
|
||||||
|
|
||||||
|
:param session_key: 会话标识
|
||||||
|
:return: 找到并关闭会话时返回 True
|
||||||
|
"""
|
||||||
|
with cls._sessions_lock:
|
||||||
|
session = cls._sessions.pop(session_key, None)
|
||||||
|
if not session:
|
||||||
|
return False
|
||||||
|
cls._close_session_state(session)
|
||||||
|
return True
|
||||||
|
|
||||||
def open_tab(
|
def open_tab(
|
||||||
self,
|
self,
|
||||||
session: _BrowserSessionState,
|
session: _BrowserSessionState,
|
||||||
@@ -479,24 +686,29 @@ class BrowserSessionHelper:
|
|||||||
if session:
|
if session:
|
||||||
return session
|
return session
|
||||||
|
|
||||||
context = self._launch_context(
|
context = self._launch_context(
|
||||||
headless=self.headless,
|
headless=self.headless,
|
||||||
user_agent=user_agent,
|
user_agent=user_agent,
|
||||||
viewport=self.viewport,
|
viewport=self.viewport,
|
||||||
)
|
)
|
||||||
page = context.new_page()
|
page = context.new_page()
|
||||||
if cookies:
|
if cookies:
|
||||||
page.set_extra_http_headers({"cookie": cookies})
|
page.set_extra_http_headers({"cookie": cookies})
|
||||||
session = _BrowserSessionState(
|
session = _BrowserSessionState(
|
||||||
session_key=session_key,
|
session_key=session_key,
|
||||||
context=context,
|
context=context,
|
||||||
pages=[page],
|
pages=[page],
|
||||||
user_agent=user_agent,
|
user_agent=user_agent,
|
||||||
cookies=cookies,
|
cookies=cookies,
|
||||||
)
|
)
|
||||||
|
with self._sessions_lock:
|
||||||
|
existing_session = self._sessions.get(session_key)
|
||||||
|
if existing_session:
|
||||||
|
self._close_session_state(session)
|
||||||
|
return existing_session
|
||||||
self._sessions[session_key] = session
|
self._sessions[session_key] = session
|
||||||
self._enforce_session_limit()
|
self._enforce_session_limit(protect_session_key=session_key)
|
||||||
return session
|
return session
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def _prune_sessions(cls) -> None:
|
def _prune_sessions(cls) -> None:
|
||||||
@@ -511,14 +723,29 @@ class BrowserSessionHelper:
|
|||||||
cls.close_session(session_key)
|
cls.close_session(session_key)
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def _enforce_session_limit(cls) -> None:
|
def _enforce_session_limit(cls, protect_session_key: Optional[str] = None) -> None:
|
||||||
while len(cls._sessions) > cls.MAX_SESSIONS:
|
"""
|
||||||
oldest_key = min(
|
清理超过数量上限的旧会话。
|
||||||
cls._sessions,
|
|
||||||
key=lambda key: cls._sessions[key].last_used_at,
|
:param protect_session_key: 本次刚创建、需要优先保留的会话标识
|
||||||
)
|
"""
|
||||||
session = cls._sessions.pop(oldest_key)
|
while True:
|
||||||
cls._close_session_state(session)
|
with cls._sessions_lock:
|
||||||
|
if len(cls._sessions) <= cls.MAX_SESSIONS:
|
||||||
|
return
|
||||||
|
candidate_keys = [
|
||||||
|
session_key
|
||||||
|
for session_key in cls._sessions
|
||||||
|
if session_key != protect_session_key
|
||||||
|
]
|
||||||
|
if not candidate_keys:
|
||||||
|
return
|
||||||
|
oldest_key = min(
|
||||||
|
candidate_keys,
|
||||||
|
key=lambda key: cls._sessions[key].last_used_at,
|
||||||
|
)
|
||||||
|
if not cls.close_session(oldest_key):
|
||||||
|
return
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _close_session_state(session: _BrowserSessionState) -> None:
|
def _close_session_state(session: _BrowserSessionState) -> None:
|
||||||
|
|||||||
@@ -1,6 +1,8 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import json
|
import json
|
||||||
|
import threading
|
||||||
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
from unittest.mock import patch
|
from unittest.mock import patch
|
||||||
|
|
||||||
@@ -45,6 +47,7 @@ class _FakePage:
|
|||||||
self.clicks = []
|
self.clicks = []
|
||||||
self.fills = []
|
self.fills = []
|
||||||
self.selects = []
|
self.selects = []
|
||||||
|
self.close_thread_id = None
|
||||||
|
|
||||||
def set_extra_http_headers(self, headers: dict[str, str]) -> None:
|
def set_extra_http_headers(self, headers: dict[str, str]) -> None:
|
||||||
"""记录额外请求头。"""
|
"""记录额外请求头。"""
|
||||||
@@ -124,6 +127,7 @@ class _FakePage:
|
|||||||
|
|
||||||
def close(self) -> None:
|
def close(self) -> None:
|
||||||
"""记录页面关闭状态。"""
|
"""记录页面关闭状态。"""
|
||||||
|
self.close_thread_id = threading.get_ident()
|
||||||
self.closed = True
|
self.closed = True
|
||||||
|
|
||||||
|
|
||||||
@@ -133,6 +137,7 @@ class _FakeContext:
|
|||||||
def __init__(self, pages: Optional[list[_FakePage]] = None) -> None:
|
def __init__(self, pages: Optional[list[_FakePage]] = None) -> None:
|
||||||
self.pages = pages or [_FakePage()]
|
self.pages = pages or [_FakePage()]
|
||||||
self.closed = False
|
self.closed = False
|
||||||
|
self.close_thread_id = None
|
||||||
|
|
||||||
def new_page(self) -> _FakePage:
|
def new_page(self) -> _FakePage:
|
||||||
"""返回或创建模拟页面。"""
|
"""返回或创建模拟页面。"""
|
||||||
@@ -146,6 +151,7 @@ class _FakeContext:
|
|||||||
|
|
||||||
def close(self) -> None:
|
def close(self) -> None:
|
||||||
"""记录上下文关闭状态。"""
|
"""记录上下文关闭状态。"""
|
||||||
|
self.close_thread_id = threading.get_ident()
|
||||||
self.closed = True
|
self.closed = True
|
||||||
|
|
||||||
|
|
||||||
@@ -250,6 +256,55 @@ def test_browser_session_helper_reuses_page_within_session():
|
|||||||
assert not context.closed
|
assert not context.closed
|
||||||
|
|
||||||
|
|
||||||
|
def test_browser_session_helper_runs_same_session_on_one_worker_thread():
|
||||||
|
"""同一 session_key 的浏览器操作应固定在同一个工作线程。"""
|
||||||
|
page = _FakePage()
|
||||||
|
context = _FakeContext([page])
|
||||||
|
helper = BrowserSessionHelper()
|
||||||
|
caller_thread_ids = set()
|
||||||
|
session_thread_ids = []
|
||||||
|
barrier = threading.Barrier(2)
|
||||||
|
|
||||||
|
def _run_from_caller_thread() -> int:
|
||||||
|
"""从外部调用线程进入同一个浏览器会话。"""
|
||||||
|
caller_thread_ids.add(threading.get_ident())
|
||||||
|
barrier.wait(timeout=1)
|
||||||
|
return helper.with_session(
|
||||||
|
"session-1",
|
||||||
|
lambda _session: threading.get_ident(),
|
||||||
|
)
|
||||||
|
|
||||||
|
with patch.object(BrowserSessionHelper, "_launch_context", return_value=context):
|
||||||
|
with ThreadPoolExecutor(max_workers=2) as executor:
|
||||||
|
futures = [
|
||||||
|
executor.submit(_run_from_caller_thread),
|
||||||
|
executor.submit(_run_from_caller_thread),
|
||||||
|
]
|
||||||
|
session_thread_ids = [future.result(timeout=1) for future in futures]
|
||||||
|
|
||||||
|
assert len(caller_thread_ids) == 2
|
||||||
|
assert len(set(session_thread_ids)) == 1
|
||||||
|
assert session_thread_ids[0] not in caller_thread_ids
|
||||||
|
|
||||||
|
|
||||||
|
def test_browser_session_helper_closes_session_on_worker_thread():
|
||||||
|
"""关闭会话时应在创建浏览器对象的工作线程内释放资源。"""
|
||||||
|
page = _FakePage()
|
||||||
|
context = _FakeContext([page])
|
||||||
|
helper = BrowserSessionHelper()
|
||||||
|
|
||||||
|
with patch.object(BrowserSessionHelper, "_launch_context", return_value=context):
|
||||||
|
session_thread_id = helper.with_session(
|
||||||
|
"session-1",
|
||||||
|
lambda _session: threading.get_ident(),
|
||||||
|
)
|
||||||
|
closed = BrowserSessionHelper.close_session("session-1")
|
||||||
|
|
||||||
|
assert closed is True
|
||||||
|
assert page.close_thread_id == session_thread_id
|
||||||
|
assert context.close_thread_id == session_thread_id
|
||||||
|
|
||||||
|
|
||||||
def test_browse_webpage_returns_snapshot_with_refs_after_goto():
|
def test_browse_webpage_returns_snapshot_with_refs_after_goto():
|
||||||
"""goto 后应返回包含可交互元素 ref 的页面快照。"""
|
"""goto 后应返回包含可交互元素 ref 的页面快照。"""
|
||||||
page = _FakePage()
|
page = _FakePage()
|
||||||
|
|||||||
Reference in New Issue
Block a user