fix(scores): poll gevent's queue instead of using its join(timeout)

When flush() ran on a gevent hub threadpool thread with a score pending,
gevent's Event-based join(timeout) made the consumer's task_done() loop
forever in gevent's cross-thread wakeup while holding the GIL, freezing
the process. v5.0.2's join() raised LoopExit there instead.

The gevent path now polls unfinished_tasks with time.sleep(): sleep(0)
first, so the consumer can take the wake signal and an idle flush()
returns right away, then every 5 ms. The stdlib path is unchanged.

Author: DavidTraina

langfuse/_client/resource_manager.py

@@ -18,9 +18,9 @@
 import os
 import sys
 import threading
+import time
 import urllib.request
 import weakref
-from contextlib import nullcontext
 from queue import Full, Queue
 from typing import Any, Callable, Dict, List, Literal, Optional, cast
 
@@ -610,25 +610,37 @@ def _join_score_queue(self) -> None:
         Queue.join() would block forever when the consumers stop first, e.g.
         when shutdown() runs concurrently and a wake signal is never taken.
         """
-        score_queue = self._score_queue
-        # gevent's monkey-patched Queue has no all_tasks_done, but its join()
-        # takes a timeout.
-        all_tasks_done = getattr(score_queue, "all_tasks_done", None)
-
-        with nullcontext() if all_tasks_done is None else all_tasks_done:
-            while score_queue.unfinished_tasks:
-                if not any(c.is_alive() for c in self._score_consumers):
-                    langfuse_logger.debug(
-                        "Score consumers stopped with %s queue item(s) unprocessed",
-                        score_queue.unfinished_tasks,
-                    )
-
-                    return
-
-                if all_tasks_done is None:
-                    score_queue.join(timeout=0.1)  # type: ignore[call-arg]
-                else:
-                    all_tasks_done.wait(timeout=0.1)
+        all_tasks_done = getattr(self._score_queue, "all_tasks_done", None)
+
+        if all_tasks_done is None:
+            # gevent's patched Queue has no all_tasks_done, and its join()
+            # freezes the process when flush() runs on a hub threadpool thread.
+            # sleep(0) first lets the consumer take the wake signal, so an idle
+            # flush() returns right away.
+            time.sleep(0)
+            while self._should_wait_for_score_queue():
+                time.sleep(0.005)
+
+            return
+
+        with all_tasks_done:
+            while self._should_wait_for_score_queue():
+                all_tasks_done.wait(timeout=0.1)
+
+    def _should_wait_for_score_queue(self) -> bool:
+        unfinished_tasks = self._score_queue.unfinished_tasks
+        if not unfinished_tasks:
+            return False
+
+        if any(c.is_alive() for c in self._score_consumers):
+            return True
+
+        langfuse_logger.debug(
+            "Score consumers stopped with %s queue item(s) unprocessed",
+            unfinished_tasks,
+        )
+
+        return False
 
     def shutdown(self) -> None:
         self._shutdown = True

tests/unit/test_score_consumer.py

@@ -6,7 +6,7 @@
 import time
 from datetime import datetime, timezone
 from queue import Queue
-from typing import Any, Callable, Dict, List, Optional, Sequence
+from typing import Any, Callable, Dict, List, Sequence
 
 import httpx
 import pytest
@@ -419,10 +419,7 @@ def test_create_score_posts_to_scores_endpoint(client_factory):
 
 
 class GeventLikeQueue:
-    """queue.Queue as gevent's monkey patching replaces it.
-
-    It has no all_tasks_done Condition, and join() takes a timeout.
-    """
+    """queue.Queue as gevent's monkey patching replaces it: no all_tasks_done."""
 
     def __init__(self, maxsize: int = 0) -> None:
         self._queue: Queue = Queue(maxsize)
@@ -432,12 +429,6 @@ def __getattr__(self, name: str) -> Any:
             raise AttributeError(name)
         return getattr(self._queue, name)
 
-    def join(self, timeout: Optional[float] = None) -> bool:
-        with self._queue.all_tasks_done:
-            return self._queue.all_tasks_done.wait_for(
-                lambda: not self._queue.unfinished_tasks, timeout
-            )
-
 
 @pytest.fixture(params=[Queue, GeventLikeQueue], ids=["stdlib", "gevent"])
 def queue_class(request, monkeypatch):