fix(media): don't hang flush() after a repeated shutdown() (#2033)

Author: DavidTraina

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,
         *,

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