From 271867f6fa418db293f9baaf9b2d3b0c5ab8738f Mon Sep 17 00:00:00 2001 From: HenryG <314296805+HenryGHJ@users.noreply.github.com> Date: Sat, 8 Aug 2026 01:51:28 +0800 Subject: [PATCH] fix(gateway): bound media history workers --- gateway/platforms/base.py | 86 ++++++++++++++++--- .../gateway/test_73771_media_resend_dedup.py | 71 +++++++++++++++ 2 files changed, 143 insertions(+), 14 deletions(-) diff --git a/gateway/platforms/base.py b/gateway/platforms/base.py index 61376f4de9b10..b1473e1350d31 100644 --- a/gateway/platforms/base.py +++ b/gateway/platforms/base.py @@ -16,6 +16,7 @@ import socket as _socket import subprocess import sys import tempfile +import threading import time import uuid import weakref @@ -49,6 +50,13 @@ _POST_DELIVERY_CALLBACK_TIMEOUT_SECONDS = 30.0 # Keep this comfortably below the Discord heartbeat watchdog window and fail # open rather than withholding a legitimate attachment. _HISTORY_MEDIA_LOOKUP_TIMEOUT_SECONDS = 5.0 +# Timed-out reads cannot be cancelled while SQLite/Python code is already +# running. Isolate and cap them so wedged best-effort dedup work cannot consume +# the shared asyncio executor or create an unbounded number of worker threads. +_HISTORY_MEDIA_LOOKUP_MAX_WORKERS = 2 +_HISTORY_MEDIA_LOOKUP_ADMISSION = threading.BoundedSemaphore( + _HISTORY_MEDIA_LOOKUP_MAX_WORKERS +) def _platform_name(platform) -> str: @@ -3480,6 +3488,66 @@ class BasePlatformAdapter(ABC): from gateway.run import _collect_history_media_paths return _collect_history_media_paths(history) + async def _bounded_history_media_paths_for_session( + self, session_key: str + ) -> Optional[set]: + """Run best-effort history lookup in a bounded isolated daemon thread.""" + admission = _HISTORY_MEDIA_LOOKUP_ADMISSION + if not admission.acquire(blocking=False): + logger.warning( + "[%s] Media-delivery history lookup capacity exhausted for %s; " + "delivering bare local file path(s) without history dedup", + self.name, + session_key, + ) + return None + + loop = asyncio.get_running_loop() + result_future = loop.create_future() + + def _publish_result(result=None, error=None): + if result_future.done(): + return + if error is not None: + result_future.set_exception(error) + else: + result_future.set_result(result) + + def _worker(): + try: + result = self._history_media_paths_for_session(session_key) + except BaseException as exc: + try: + loop.call_soon_threadsafe(_publish_result, None, exc) + except RuntimeError: + pass # Event loop already closed during gateway shutdown. + else: + try: + loop.call_soon_threadsafe(_publish_result, result, None) + except RuntimeError: + pass # Event loop already closed during gateway shutdown. + finally: + admission.release() + + threading.Thread( + target=_worker, + name="media-history-lookup", + daemon=True, + ).start() + try: + return await asyncio.wait_for( + result_future, + timeout=_HISTORY_MEDIA_LOOKUP_TIMEOUT_SECONDS, + ) + except asyncio.TimeoutError: + logger.warning( + "[%s] Timed out loading media-delivery history for %s; " + "delivering bare local file path(s) without history dedup", + self.name, + session_key, + ) + return None + @abstractmethod async def connect(self, *, is_reconnect: bool = False) -> bool: """ @@ -5923,21 +5991,11 @@ class BasePlatformAdapter(ABC): # fail open by delivering the candidate file. _history_media_paths = None if local_files: - try: - _history_media_paths = await asyncio.wait_for( - asyncio.to_thread( - self._history_media_paths_for_session, - session_key, - ), - timeout=_HISTORY_MEDIA_LOOKUP_TIMEOUT_SECONDS, - ) - except asyncio.TimeoutError: - logger.warning( - "[%s] Timed out loading media-delivery history for %s; " - "delivering bare local file path(s) without history dedup", - self.name, - session_key, + _history_media_paths = ( + await self._bounded_history_media_paths_for_session( + session_key ) + ) if _history_media_paths: _suppressed = [p for p in local_files if p in _history_media_paths] if _suppressed: diff --git a/tests/gateway/test_73771_media_resend_dedup.py b/tests/gateway/test_73771_media_resend_dedup.py index dbed4ed554a81..8eb1781d3a74d 100644 --- a/tests/gateway/test_73771_media_resend_dedup.py +++ b/tests/gateway/test_73771_media_resend_dedup.py @@ -23,6 +23,7 @@ sibling): import asyncio import logging +import threading import time from types import SimpleNamespace from unittest.mock import AsyncMock @@ -246,6 +247,34 @@ async def test_plain_text_response_does_not_load_transcript(): assert any("Plain response" in item["content"] for item in adapter.sent) +@pytest.mark.asyncio +async def test_explicit_media_response_does_not_load_transcript(tmp_path, monkeypatch): + """Explicit MEDIA delivery must not touch SQLite-backed history.""" + pdf = _allowed_file(tmp_path, monkeypatch, "explicit-no-history.pdf") + adapter = _DummyAdapter() + adapter._keep_typing = _hold_typing + + class _ExplodingStore: + calls = 0 + + def peek_session_id(self, _session_key): + self.calls += 1 + raise AssertionError("explicit MEDIA delivery loaded session history") + + store = _ExplodingStore() + adapter.set_session_store(store) + + async def handler(_event): + return f"Here is the file.\nMEDIA:{pdf}" + + adapter.set_message_handler(handler) + event = _make_event() + await adapter._process_message_background(event, build_session_key(event.source)) + + assert store.calls == 0 + assert adapter.documents == [str(pdf)] + + @pytest.mark.asyncio async def test_bare_path_history_lookup_does_not_block_event_loop(tmp_path, monkeypatch): """A slow transcript read must run outside the platform event loop.""" @@ -318,6 +347,48 @@ async def test_bare_path_history_lookup_timeout_fails_open(tmp_path, monkeypatch assert adapter.documents == [str(pdf)] +@pytest.mark.asyncio +async def test_history_lookup_saturation_fails_open_without_new_worker(monkeypatch): + """Wedged lookups are bounded and cannot consume unbounded worker threads.""" + monkeypatch.setattr("gateway.platforms.base._HISTORY_MEDIA_LOOKUP_TIMEOUT_SECONDS", 1.0) + monkeypatch.setattr( + "gateway.platforms.base._HISTORY_MEDIA_LOOKUP_ADMISSION", + threading.BoundedSemaphore(2), + ) + adapter = _DummyAdapter() + release = threading.Event() + two_started = threading.Event() + calls = 0 + calls_lock = threading.Lock() + + def blocked_lookup(_session_key): + nonlocal calls + with calls_lock: + calls += 1 + if calls == 2: + two_started.set() + release.wait(timeout=1) + return None + + monkeypatch.setattr(adapter, "_history_media_paths_for_session", blocked_lookup) + first = asyncio.create_task(adapter._bounded_history_media_paths_for_session("one")) + second = asyncio.create_task(adapter._bounded_history_media_paths_for_session("two")) + deadline = time.monotonic() + 1 + while not two_started.is_set() and time.monotonic() < deadline: + await asyncio.sleep(0.005) + assert two_started.is_set() + + began = time.monotonic() + third = await adapter._bounded_history_media_paths_for_session("three") + elapsed = time.monotonic() - began + + assert third is None + assert elapsed < 0.1 + assert calls == 2 + release.set() + await asyncio.gather(first, second) + + # --------------------------------------------------------------------------- # Streaming sibling (run.py _deliver_media_from_response) # ---------------------------------------------------------------------------