fix(gateway): bound media history workers
This commit is contained in:
HenryG
kshitij
parent
e52acf76a1
commit
271867f6fa
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
Loading…
Reference in New Issue