diff --git a/langfuse/_task_manager/media_manager.py b/langfuse/_task_manager/media_manager.py index 9a66ecd63..3ad9ac62f 100644 --- a/langfuse/_task_manager/media_manager.py +++ b/langfuse/_task_manager/media_manager.py @@ -62,7 +62,7 @@ def process_next_media_upload(self) -> None: upload_job = self._queue.get(block=True, timeout=1) if upload_job is _SHUTDOWN_SENTINEL: - self._queue.task_done() + # Already marked done by signal_shutdown(). return logger.debug( @@ -92,6 +92,10 @@ def signal_shutdown(self, *, count: int = 1) -> None: # observe the paused flag on the next loop iteration. break + # Mark it done now, since no consumer may take it (e.g. after a + # repeated shutdown()) and queue.join() would then never return. + self._queue.task_done() + def _find_and_process_media( self, *, diff --git a/tests/unit/test_resource_manager.py b/tests/unit/test_resource_manager.py index 232ec4f2f..27e600cdc 100644 --- a/tests/unit/test_resource_manager.py +++ b/tests/unit/test_resource_manager.py @@ -1,11 +1,13 @@ """Test the LangfuseResourceManager and get_client() function.""" import logging +import threading from queue import Queue from types import SimpleNamespace from typing import Sequence from unittest.mock import Mock +import pytest from opentelemetry.sdk.trace import ReadableSpan from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult @@ -471,3 +473,36 @@ def signal_shutdown(self, *, count): ("join", 0), ("join", 1), ] + + +@pytest.mark.filterwarnings("error::pytest.PytestUnhandledThreadExceptionWarning") +@pytest.mark.parametrize("concurrent", [False, True], ids=["sequential", "concurrent"]) +def test_flush_returns_after_repeated_shutdown(concurrent): + client = Langfuse( + public_key="pk-repeated-shutdown", + secret_key="sk-repeated-shutdown", + span_exporter=NoOpSpanExporter(), + ) + media_queue = client._resources._media_upload_queue + shutdowns = [ + threading.Thread(target=client.shutdown, daemon=True) for _ in range(2) + ] + + for shutdown in shutdowns: + shutdown.start() + if not concurrent: + shutdown.join(timeout=5) + for shutdown in shutdowns: + shutdown.join(timeout=5) + + flush = threading.Thread(target=client.flush, daemon=True) + flush.start() + flush.join(timeout=2) + hung = [thread for thread in (*shutdowns, flush) if thread.is_alive()] + # Acknowledge shutdown signals that no consumer took, so a hung thread and + # the reset() in the autouse fixture can finish. + while media_queue.unfinished_tasks: + media_queue.get_nowait() + media_queue.task_done() + + assert not hung