diff --git a/langfuse/_client/resource_manager.py b/langfuse/_client/resource_manager.py index 67c44920a..520b4ceba 100644 --- a/langfuse/_client/resource_manager.py +++ b/langfuse/_client/resource_manager.py @@ -134,10 +134,10 @@ def __new__( id_generator: Optional[IdGenerator] = None, span_exporter: Optional[SpanExporter] = None, ) -> "LangfuseResourceManager": - if public_key in cls._instances: - return cls._instances[public_key] - with cls._lock: + if public_key in cls._instances: + return cls._instances[public_key] + if public_key not in cls._instances: instance = super(LangfuseResourceManager, cls).__new__(cls) @@ -209,6 +209,8 @@ def _initialize_instance( self.mask_otel_spans = mask_otel_spans self.environment = environment self._shutdown = False + self._shutdown_event = threading.Event() + self._admission_lock = threading.Lock() # Store additional client settings for get_client() to use self.timeout = timeout @@ -420,6 +422,18 @@ def _at_fork_reinit(self) -> None: # the lock is class-level state needed by the child (e.g. to create a new client) # even if this particular instance was already shut down. LangfuseResourceManager._lock = threading.RLock() + # Recreate shutdown event and admission lock for child process + if hasattr(self, "_shutdown_event"): + self._shutdown_event = threading.Event() + if self._shutdown: + self._shutdown_event.set() + if hasattr(self, "_admission_lock"): + self._admission_lock = threading.Lock() + # Recreate media manager state lock if held at fork + if hasattr(self, "_media_manager") and hasattr( + self._media_manager, "_state_lock" + ): + self._media_manager._state_lock = threading.Lock() if self._shutdown: return @@ -508,10 +522,16 @@ def add_score_task(self, event: dict, *, force_sample: bool = False) -> None: ) if should_sample: - langfuse_logger.debug( - f"Score: Enqueuing event type={event['type']} for trace_id={event['body'].trace_id} name={event['body'].name} value={event['body'].value}" - ) - self._score_ingestion_queue.put(event, block=False) + with self._admission_lock: + if self._shutdown: + langfuse_logger.warning( + "Score: Dropping event because the Langfuse client has already been shut down." + ) + return + langfuse_logger.debug( + f"Score: Enqueuing event type={event['type']} for trace_id={event['body'].trace_id} name={event['body'].name} value={event['body'].value}" + ) + self._score_ingestion_queue.put(event, block=False) except Full: langfuse_logger.warning( @@ -531,10 +551,16 @@ def add_trace_task( event: dict, ) -> None: try: - langfuse_logger.debug( - f"Trace: Enqueuing event type={event['type']} for trace_id={event['body'].id}" - ) - self._score_ingestion_queue.put(event, block=False) + with self._admission_lock: + if self._shutdown: + langfuse_logger.warning( + "Trace: Dropping event because the Langfuse client has already been shut down." + ) + return + langfuse_logger.debug( + f"Trace: Enqueuing event type={event['type']} for trace_id={event['body'].id}" + ) + self._score_ingestion_queue.put(event, block=False) except Full: langfuse_logger.warning( @@ -612,13 +638,40 @@ def flush(self) -> None: langfuse_logger.debug("Successfully flushed media upload queue") def shutdown(self) -> None: - self._shutdown = True + is_first = False + with self._admission_lock: + if self._shutdown: + # Another thread already started shutdown; wait for it + pass + else: + self._shutdown = True + is_first = True + + if not is_first: + # Wait indefinitely for the first shutdown to complete. + # A timeout would hide an incomplete shutdown and violate + # the invariant that shutdown() only returns after resources + # are drained and consumers stopped. + self._shutdown_event.wait() + return # Unregister the atexit handler first - atexit.unregister(self.shutdown) + try: + atexit.unregister(self.shutdown) + except Exception: + pass - self.flush() - self._stop_and_join_consumer_threads() + try: + self.flush() + finally: + try: + # Block further media enqueues that could be stranded + # after flush but before consumers are stopped. Media discovered + # during the preceding force_flush() was already enqueued. + self._media_manager.begin_shutdown() + self._stop_and_join_consumer_threads() + finally: + self._shutdown_event.set() def _init_tracer_provider( diff --git a/langfuse/_task_manager/media_manager.py b/langfuse/_task_manager/media_manager.py index 14dceec19..d38429121 100644 --- a/langfuse/_task_manager/media_manager.py +++ b/langfuse/_task_manager/media_manager.py @@ -1,4 +1,5 @@ import os +import threading import time from queue import Empty, Full, Queue from typing import Any, Callable, Optional, TypeVar, cast @@ -45,6 +46,8 @@ def __init__( self._enabled = os.environ.get( LANGFUSE_MEDIA_UPLOAD_ENABLED, "True" ).lower() not in ("false", "0") + self._shutdown = False + self._state_lock = threading.Lock() def reinitialize( self, @@ -53,9 +56,15 @@ def reinitialize( httpx_client: httpx.Client, media_upload_queue: Queue, ) -> None: - self._api_client = api_client - self._httpx_client = httpx_client - self._queue = media_upload_queue + with self._state_lock: + self._api_client = api_client + self._httpx_client = httpx_client + self._queue = media_upload_queue + self._shutdown = False + + def begin_shutdown(self) -> None: + with self._state_lock: + self._shutdown = True def process_next_media_upload(self) -> None: try: @@ -98,6 +107,12 @@ def _find_and_process_media( ) -> Any: if not self._enabled: return data + with self._state_lock: + if self._shutdown: + logger.warning( + "Media: Skipping upload because the Langfuse client has already been shut down." + ) + return data seen = set() max_levels = 10 @@ -279,10 +294,16 @@ def _process_media( field=field, ) - self._queue.put( - item=upload_media_job, - block=False, - ) + with self._state_lock: + if self._shutdown: + logger.warning( + f"Media: Skipping upload for media_id={media._media_id} because the Langfuse client has already been shut down." + ) + return + self._queue.put( + item=upload_media_job, + block=False, + ) logger.debug( f"Queue: Enqueued media ID {media._media_id} for upload processing | trace_id={trace_id} | field={field}" ) diff --git a/tests/unit/test_media.py b/tests/unit/test_media.py index 387eae745..8300a6985 100644 --- a/tests/unit/test_media.py +++ b/tests/unit/test_media.py @@ -263,6 +263,7 @@ def test_resolve_media_references_uses_configured_httpx_client(): "https://example.com/test.jpg", timeout=fetch_timeout_seconds ) + def test_init_with_urlsafe_base64_data_uri(): original_bytes = b"\xfb\xff" urlsafe_base64 = base64.urlsafe_b64encode(original_bytes).decode() @@ -275,3 +276,78 @@ def test_init_with_urlsafe_base64_data_uri(): assert media._content_type == "application/octet-stream" assert media._content_bytes == original_bytes + +def test_1799_media_admission_toctou_and_begin_shutdown(): + """Media check+put must be atomic – no job stranded after begin_shutdown.""" + import threading + from queue import Queue + from unittest.mock import Mock + + from langfuse._task_manager.media_manager import MediaManager + + q: Queue = Queue() + mm = MediaManager( + api_client=Mock(), + httpx_client=Mock(), + media_upload_queue=q, + ) + + # Stub upload so queue drain is instant + mm._process_upload_media_job = Mock(return_value=None) # type: ignore[attr-defined] + + # Deterministic barrier: block inside queue.put while holding _state_lock + orig_put = q.put + block_in_put = threading.Event() + release_put = threading.Event() + + def blocking_put(*args, **kwargs): + block_in_put.set() + assert release_put.wait(timeout=2) + return orig_put(*args, **kwargs) + + q.put = blocking_put # type: ignore[assignment] + + # Create a valid media object + media = LangfuseMedia(content_bytes=b"hello", content_type="text/plain") + assert media._media_id is not None + + producer = threading.Thread( + target=lambda: mm._process_media( + media=media, trace_id="0" * 32, observation_id="0" * 16, field="input" + ) + ) + producer.start() + assert block_in_put.wait(timeout=2) + + # begin_shutdown must wait for the put's lock, not interleave and strand + shutdown_done = threading.Event() + + def do_shutdown(): + mm.begin_shutdown() + shutdown_done.set() + + t_shutdown = threading.Thread(target=do_shutdown) + t_shutdown.start() + # shutdown should be blocked while producer holds _state_lock inside put + assert not shutdown_done.wait(timeout=0.3) + + release_put.set() + producer.join(timeout=2) + t_shutdown.join(timeout=2) + + # Producer's job was enqueued atomically; simulate flush drain via task_done + assert q.unfinished_tasks == 1 + assert q.qsize() == 1 + # Drain as flush would (join would wait for task_done) + q.get() + q.task_done() + q.join() + assert q.unfinished_tasks == 0 + + # After shutdown, further admission is dropped, not stranded + media2 = LangfuseMedia(content_bytes=b"world", content_type="text/plain") + mm._process_media( + media=media2, trace_id="0" * 32, observation_id="0" * 16, field="input" + ) + assert q.unfinished_tasks == 0 + assert q.qsize() == 0 diff --git a/tests/unit/test_resource_manager.py b/tests/unit/test_resource_manager.py index f66a1e052..88b5434bb 100644 --- a/tests/unit/test_resource_manager.py +++ b/tests/unit/test_resource_manager.py @@ -442,3 +442,347 @@ def signal_shutdown(self, *, count): ("join", 0), ("join", 1), ] + + +# --------------------------------------------------------------------------- +# #1799 regression: durable invariants (added 2026-08) +# --------------------------------------------------------------------------- + + +def test_1799_stale_client_score_trace_dropped_after_shutdown(monkeypatch): + """Original #1799 deadlock: stale same-key client must not enqueue after shutdown. + + Proves flush/shutdown cannot hang via stranded Queue.unfinished_tasks. + Fails on base e3f7e7dd where add_* had no post-shutdown guard. + """ + import atexit + + monkeypatch.setenv("LANGFUSE_MEDIA_UPLOAD_ENABLED", "false") + + with LangfuseResourceManager._lock: + LangfuseResourceManager._instances.clear() + + client_a = Langfuse( + public_key="pk-1799-stale", + secret_key="sk-1799-stale", + base_url="http://localhost:9", + span_exporter=NoOpSpanExporter(), + ) + # same public_key → same manager (cached) + client_b = Langfuse( + public_key="pk-1799-stale", + secret_key="sk-1799-stale", + base_url="http://localhost:9", + span_exporter=NoOpSpanExporter(), + ) + mgr = client_a._resources + assert mgr is not None + assert client_b._resources is mgr + + # Prevent atexit hang on base when queue is stranded + try: + atexit.unregister(mgr.shutdown) + except Exception: + pass + + client_a.shutdown() + + try: + assert mgr._shutdown is True + # stale client attempts work – must be dropped, not enqueued + client_b.create_score(name="stale-score", value=1.0) + client_b._create_trace_tags_via_ingestion(trace_id="0" * 32, tags=["stale-tag"]) + + assert mgr._score_ingestion_queue.unfinished_tasks == 0 + assert mgr._score_ingestion_queue.qsize() == 0 + # flush/shutdown must return boundedly (would hang on base via Queue.join) + import threading + + flush_done = threading.Event() + + def do_flush(): + mgr.flush() + flush_done.set() + + t = threading.Thread(target=do_flush, daemon=True) + t.start() + assert flush_done.wait(timeout=2), "flush hung – stranded Queue task on base" + t.join(timeout=1) + + shutdown_done = threading.Event() + + def do_second_shutdown(): + client_b.shutdown() + shutdown_done.set() + + t2 = threading.Thread(target=do_second_shutdown, daemon=True) + t2.start() + assert shutdown_done.wait(timeout=2), "second shutdown hung" + t2.join(timeout=1) + assert mgr._shutdown_event.is_set() + finally: + # Drain stranded queue on base so atexit/pytest teardown never hangs + try: + while mgr._score_ingestion_queue.unfinished_tasks: + try: + mgr._score_ingestion_queue.get_nowait() + mgr._score_ingestion_queue.task_done() + except Exception: + break + except Exception: + pass + try: + atexit.unregister(mgr.shutdown) + except Exception: + pass + with LangfuseResourceManager._lock: + LangfuseResourceManager._instances.clear() + + +def test_1799_admission_vs_shutdown_atomicity(monkeypatch): + """Score/trace check+put must be atomic vs shutdown – no stranded task. + + Deterministic: producer blocks inside queue.put while holding admission lock, + shutdown must wait for the lock before setting _shutdown and draining. + """ + import atexit + import threading + + monkeypatch.setenv("LANGFUSE_MEDIA_UPLOAD_ENABLED", "false") + + with LangfuseResourceManager._lock: + LangfuseResourceManager._instances.clear() + + client = Langfuse( + public_key="pk-1799-atomic", + secret_key="sk-1799-atomic", + base_url="http://localhost:9", + span_exporter=NoOpSpanExporter(), + ) + mgr = client._resources + assert mgr is not None + try: + atexit.unregister(mgr.shutdown) + except Exception: + pass + + # Make downstream fast so flush/join is bounded + mgr._score_ingestion_client.batch_post = Mock(return_value=None) # type: ignore[attr-defined] + mock_provider = Mock() + mock_provider.force_flush = Mock(return_value=None) + mgr.tracer_provider = mock_provider # type: ignore[assignment] + + orig_put = mgr._score_ingestion_queue.put + block_in_put = threading.Event() + release_put = threading.Event() + + def blocking_put(item, block=False): + block_in_put.set() + assert release_put.wait(timeout=2), "test timed out waiting for release" + return orig_put(item, block=block) + + mgr._score_ingestion_queue.put = blocking_put # type: ignore[assignment] + + # Producer enters admission critical section (holds admission_lock during put) + from unittest.mock import Mock as MockBody + + body = MockBody(trace_id="0" * 32, name="race-score", value=1.0) + event = {"type": "score", "body": body} + + producer = threading.Thread( + target=lambda: mgr.add_score_task(event, force_sample=True) + ) + producer.start() + assert block_in_put.wait(timeout=2), "producer never reached put" + + shutdown_done = threading.Event() + + def do_shutdown(): + mgr.shutdown() + shutdown_done.set() + + t_shutdown = threading.Thread(target=do_shutdown) + t_shutdown.start() + + # shutdown must be blocked on admission_lock while producer holds it + assert not shutdown_done.wait(timeout=0.3), "shutdown must wait for admission lock" + + release_put.set() + producer.join(timeout=2) + t_shutdown.join(timeout=5) + + assert not producer.is_alive() + assert not t_shutdown.is_alive(), "shutdown hung" + assert mgr._score_ingestion_queue.unfinished_tasks == 0 + assert not any(c.is_alive() for c in mgr._ingestion_consumers) + + # post-shutdown admission must be dropped, not stranded + body2 = MockBody(trace_id="0" * 32, name="after", value=2.0) + mgr.add_score_task({"type": "score", "body": body2}, force_sample=True) + assert mgr._score_ingestion_queue.unfinished_tasks == 0 + + with LangfuseResourceManager._lock: + LangfuseResourceManager._instances.clear() + + +def test_1799_concurrent_shutdown_completion(monkeypatch): + """Second shutdown must wait for first to actually complete (no early return).""" + import atexit + import threading + + monkeypatch.setenv("LANGFUSE_MEDIA_UPLOAD_ENABLED", "false") + + with LangfuseResourceManager._lock: + LangfuseResourceManager._instances.clear() + + client = Langfuse( + public_key="pk-1799-concurrent", + secret_key="sk-1799-concurrent", + base_url="http://localhost:9", + span_exporter=NoOpSpanExporter(), + ) + mgr = client._resources + assert mgr is not None + try: + atexit.unregister(mgr.shutdown) + except Exception: + pass + + # Pause inside flush deterministically – only first flush blocks + flush_entered = threading.Event() + release_flush = threading.Event() + orig_flush = mgr.flush + call_count = {"n": 0} + + def pausing_flush(): + call_count["n"] += 1 + if call_count["n"] == 1: + flush_entered.set() + assert release_flush.wait(timeout=5), "flush release timed out" + return orig_flush() + + mgr.flush = pausing_flush # type: ignore[assignment] + + first_done = threading.Event() + second_done = threading.Event() + second_elapsed: list[float] = [] + + def first(): + mgr.shutdown() + first_done.set() + + def second(): + import time + + # ensure first has entered flush (and thus holds _shutdown flag) + assert flush_entered.wait(timeout=2) + start = time.monotonic() + mgr.shutdown() + second_elapsed.append(time.monotonic() - start) + second_done.set() + + t1 = threading.Thread(target=first) + t2 = threading.Thread(target=second) + t1.start() + # wait for first to be inside flush + assert flush_entered.wait(timeout=2) + t2.start() + + # second must be blocked while first is paused inside flush + assert not second_done.wait(timeout=0.5), "second shutdown returned early" + assert not first_done.is_set() + + release_flush.set() + assert first_done.wait(timeout=5) + assert second_done.wait(timeout=5) + + assert second_elapsed[0] >= 0.3, "second should have waited for first" + # On base, _shutdown_event does not exist – this also proves RED + assert hasattr(mgr, "_shutdown_event"), "missing _shutdown_event on base" + assert mgr._shutdown_event.is_set() # type: ignore[attr-defined] + assert not any(c.is_alive() for c in mgr._ingestion_consumers) + assert mgr._score_ingestion_queue.unfinished_tasks == 0 + + t1.join(timeout=2) + t2.join(timeout=2) + + with LangfuseResourceManager._lock: + LangfuseResourceManager._instances.clear() + + +def test_1799_media_during_force_flush_not_dropped_and_after_shutdown_dropped( + monkeypatch, +): + """Media discovered during tracer force_flush must survive; after shutdown it is dropped. + + Guards P1 regression where begin_shutdown before flush silently dropped media. + """ + import atexit + + monkeypatch.setenv("LANGFUSE_MEDIA_UPLOAD_ENABLED", "true") + + with LangfuseResourceManager._lock: + LangfuseResourceManager._instances.clear() + + client = Langfuse( + public_key="pk-1799-media-flush", + secret_key="sk-1799-media-flush", + base_url="http://localhost:9", + span_exporter=NoOpSpanExporter(), + ) + mgr = client._resources + assert mgr is not None + try: + atexit.unregister(mgr.shutdown) + except Exception: + pass + + # Make flush discover media: monkeypatch flush to call + # MediaManager._find_and_process_media before joining. The queue must already + # contain the media job when flush proceeds to queue.join(). + data_uri = "data:text/plain;base64,SGVsbG8=" + + def flush_with_media_discovery(): + # Simulate BatchSpanProcessor export discovering media during force_flush + mgr._media_manager._find_and_process_media( + data=data_uri, + trace_id="0" * 32, + observation_id="0" * 16, + field="input", + ) + # Proceed with normal flush (will join both queues) + # We call the original flush's internals without re-entering our wrapper + if mgr.tracer_provider is not None and not isinstance( + mgr.tracer_provider, __import__("opentelemetry").trace.ProxyTracerProvider + ): + # avoid double media discovery – just join queues as original does after force_flush + pass + mgr._score_ingestion_queue.join() + mgr._media_upload_queue.join() + + # Replace flush with media-discovering variant for this shutdown cycle + mgr.flush = flush_with_media_discovery # type: ignore[assignment] + + # Mock media upload to be fast and not require network: process_next will call + # _process_upload_media_job which we stub to just succeed + mgr._media_manager._process_upload_media_job = Mock(return_value=None) # type: ignore[attr-defined] + + # Shutdown must drain the media discovered during flush (before begin_shutdown) + client.shutdown() + + assert mgr._media_upload_queue.unfinished_tasks == 0 + # queue was drained – if ordering were wrong, media would have been skipped + # and unfinished would still be 0 but qsize would also be 0; we prove it was + # enqueued at all by checking that begin_shutdown happened after flush: + # after shutdown, further media must be dropped + processed = mgr._media_manager._find_and_process_media( + data=data_uri, + trace_id="0" * 32, + observation_id="0" * 16, + field="input", + ) + assert processed == data_uri + assert mgr._media_upload_queue.unfinished_tasks == 0 + + with LangfuseResourceManager._lock: + LangfuseResourceManager._instances.clear()