fix(traces): flush the lanes before an inline span flush, and skip the span thread when nothing is queued

When no thread could start at exit, the span flush ran before the event
lanes and could spend the whole exit budget. _start_span_flush now hands
back a waiter the caller runs after the lanes, and starts nothing when
the span queue is empty. The flush docstring describes the retry-within-
budget contract. Tracing test fixtures move to a conftest, and the client
tests reuse the shared ids and one slow sender.

Author: turnipdabeets

posthog/client.py

@@ -2639,10 +2639,13 @@ def flush(self, timeout_seconds: Optional[float] = 10) -> None:
                 Defaults to 10 seconds. Pass ``None`` to wait indefinitely.
                 Queued spans are sent at the same time, within the same
                 budget: at least one span request is attempted even when the
-                budget is already spent, no further one starts once it is, and
-                each request is bounded by ``timeout``. The wait for that first
-                request is not cut short, so a flush can take up to
-                ``timeout_seconds`` plus ``timeout`` in the worst case.
+                budget is already spent, a retriable failure is retried after
+                its backoff while budget remains, no other request starts once
+                it is spent, and each request is bounded by ``timeout``. The
+                wait for the last request is not cut short, so a flush can
+                take up to ``timeout_seconds`` plus ``timeout`` in the worst
+                case. A span flush already in flight for the whole wait is
+                left to finish instead.
 
         Examples:
             ```python
@@ -2664,19 +2667,25 @@ def flush(self, timeout_seconds: Optional[float] = 10) -> None:
                 for lane in self._lanes:
                     lane.flush(max(0.0, deadline - time.monotonic()))
             if span_flush is not None:
-                # The first span request is exempt from the budget and bounded
-                # only by the request timeout, so the join is not.
-                span_flush.join()
+                # The last span request is bounded only by the request
+                # timeout, so the wait is not.
+                span_flush(None)
         except Exception as e:
             self.log.exception("error flushing queue: %s", e)
             return
 
     def _start_span_flush(
         self, timeout_seconds: Optional[float]
-    ) -> Optional[threading.Thread]:
-        """Flush spans alongside the events, so a handler waits one round trip, not two."""
+    ) -> Optional[Callable[[Optional[float]], None]]:
+        """Flush spans alongside the events, so a handler waits one round trip, not two.
+
+        Returns a waiter taking the seconds to wait, or ``None`` when nothing
+        is queued. When no thread can start (interpreter shutdown), the waiter
+        runs the flush on the calling thread, after the caller has flushed the
+        event lanes.
+        """
         traces = self._traces
-        if traces is None:
+        if traces is None or not traces.has_queued_spans():
             return None
 
         def flush_spans() -> None:
@@ -2691,10 +2700,8 @@ def flush_spans() -> None:
         try:
             flusher.start()
         except RuntimeError:
-            # No new threads at interpreter shutdown; flush on this one.
-            flush_spans()
-            return None
-        return flusher
+            return lambda _seconds: flush_spans()
+        return flusher.join
 
     def _is_consumer_thread(self) -> bool:
         current = threading.current_thread()
@@ -3009,10 +3016,10 @@ def _atexit_spans(self) -> None:
         self._join_span_flush(span_flush, deadline)
 
     def _join_span_flush(
-        self, flusher: Optional[threading.Thread], deadline: float
+        self, waiter: Optional[Callable[[Optional[float]], None]], deadline: float
     ) -> None:
-        if flusher is not None:
-            flusher.join(max(0.0, deadline - time.monotonic()))
+        if waiter is not None:
+            waiter(max(0.0, deadline - time.monotonic()))
         if self._traces is not None:
             # Not close(): an app's own shutdown() hook may still run and send them.
             self._traces.warn_if_queued()

posthog/test/tracing/conftest.py

@@ -0,0 +1,32 @@
+"""Fixtures shared by the tracing tests."""
+
+import threading
+import time
+from unittest import mock
+
+import pytest
+
+from posthog.test.tracing.helpers import FakeTimer
+from posthog.tracing import _export as export_module
+
+
+@pytest.fixture
+def fake_timers():
+    """Timers that fire only when a test says so."""
+    FakeTimer.instances = []
+    with mock.patch.object(threading, "Timer", FakeTimer):
+        yield FakeTimer
+
+
+@pytest.fixture
+def no_jitter():
+    # Backoff delays are asserted exactly; TestJitter covers the spread.
+    with mock.patch.object(export_module, "_draw_jitter", return_value=1.0):
+        yield
+
+
+@pytest.fixture
+def clock():
+    state = {"now": 1000.0}
+    with mock.patch.object(time, "monotonic", lambda: state["now"]):
+        yield state

posthog/test/tracing/helpers.py

@@ -1,14 +1,9 @@
 """Shared fakes for the tracing pipeline and export tests."""
 
 import threading
-import time
 from contextvars import ContextVar
 from types import SimpleNamespace
-from unittest import mock
 
-import pytest
-
-from posthog.tracing import _export as export_module
 from posthog.tracing._config import resolve_traces_config
 from posthog.tracing._drops import DropLog
 from posthog.tracing._export import SpanExporter
@@ -80,31 +75,13 @@ def close(self):
     def warn_if_queued(self):
         pass
 
+    def has_queued(self):
+        return bool(self.records)
+
     def reinit_after_fork(self):
         self.reinitialized = True
 
 
-@pytest.fixture(autouse=True)
-def fake_timers():
-    FakeTimer.instances = []
-    with mock.patch.object(threading, "Timer", FakeTimer):
-        yield FakeTimer
-
-
-@pytest.fixture(autouse=True)
-def no_jitter():
-    # Backoff delays are asserted exactly; TestJitter covers the spread.
-    with mock.patch.object(export_module, "_draw_jitter", return_value=1.0):
-        yield
-
-
-@pytest.fixture
-def clock():
-    state = {"now": 1000.0}
-    with mock.patch.object(time, "monotonic", lambda: state["now"]):
-        yield state
-
-
 def make(client=None, context=None, **config):
     """A pipeline whose ended spans collect on a ``RecordingExporter``."""
     client = client or SimpleNamespace(disabled=False, send=True)

posthog/test/tracing/test_client_traces.py

@@ -11,14 +11,12 @@
 from posthog import Posthog
 from posthog.client import Client
 from posthog.contexts import identify_context, new_context, set_context_session
-from posthog.test.tracing.helpers import FakeTimer
+from posthog.test.tracing.helpers import SPAN_ID, TRACE_ID
 from posthog.tracing._transport import OK
 from posthog.tracing._span import NOOP_SPAN, RecordingSpan, Span
 from posthog.version import VERSION
 
 FAKE_API_KEY = "phc_test_key"
-TRACE_ID = "4bf92f3577b34da6a3ce929d0e0e4736"
-SPAN_ID = "00f067aa0ba902b7"
 
 
 def make_client(**kwargs):
@@ -33,11 +31,15 @@ def mock_session(status_code=200):
     return session
 
 
-@pytest.fixture
-def no_timers():
-    # No background drain racing the test.
-    with mock.patch.object(threading, "Timer", FakeTimer):
-        yield
+def slow_send(requests, delay=0.2):
+    """A sender that records each payload and takes ``delay`` seconds to answer."""
+
+    def send(pipeline_client, payload):
+        requests.append(payload)
+        time.sleep(delay)
+        return OK
+
+    return send
 
 
 @pytest.fixture(autouse=True)
@@ -135,6 +137,14 @@ def test_a_failed_init_is_not_retried_on_the_next_call(self):
         assert resolve.call_count == 1
         client.shutdown()
 
+    def test_a_non_callable_hook_turns_tracing_off(self, caplog):
+        caplog.set_level("ERROR", logger="posthog")
+        client = make_client(traces={"before_span_send": "scrub"})
+        assert client.start_span("x") is NOOP_SPAN
+        assert "Error initializing traces" in caplog.text
+        assert "not callable" in caplog.text
+        client.shutdown()
+
     def test_never_starts_a_pipeline_on_a_client_without_traces(self):
         client = make_client()
         client.flush()
@@ -370,19 +380,13 @@ def test_flush_resolves_when_the_span_export_fails(self):
         client.shutdown()
 
     def test_flush_stops_starting_span_requests_once_its_budget_is_spent(
-        self, no_timers
+        self, fake_timers
     ):
         client = make_client(traces={"max_export_batch_size": 1})
         client.start_span("x").end()
         client.start_span("y").end()
         requests = []
-
-        def slow_send(pipeline_client, payload):
-            requests.append(payload)
-            time.sleep(0.2)
-            return OK
-
-        client._traces._exporter._send = slow_send
+        client._traces._exporter._send = slow_send(requests)
         client.flush(timeout_seconds=0.05)
         assert len(requests) == 1
         assert len(client._traces._exporter._queue) == 1
@@ -404,7 +408,7 @@ def test_flush_without_a_timeout_drains_every_queued_span(self):
         assert client._traces._exporter._queue == []
         client.shutdown()
 
-    def test_flush_sends_spans_while_events_are_still_draining(self, no_timers):
+    def test_flush_sends_spans_while_events_are_still_draining(self, fake_timers):
         client = make_client(traces={}, sync_mode=False)
         client.start_span("x").end()
         span_sent = threading.Event()
@@ -421,7 +425,7 @@ def send(pipeline_client, payload):
         assert overlapped and all(overlapped)
         client.shutdown()
 
-    def test_flush_sends_spans_inline_when_no_thread_can_start(self, no_timers):
+    def test_flush_sends_spans_inline_when_no_thread_can_start(self, fake_timers):
         client = make_client(traces={})
         client.start_span("x").end()
         with mock.patch("posthog.client.threading.Thread") as thread:
@@ -432,6 +436,32 @@ def test_flush_sends_spans_inline_when_no_thread_can_start(self, no_timers):
         assert len(spans_from(payload)) == 1
         client.shutdown()
 
+    def test_flush_starts_no_span_thread_when_nothing_is_queued(self, fake_timers):
+        client = make_client(traces={})
+        client.start_span("x").end()
+        client.flush()
+        with mock.patch("posthog.client.threading.Thread") as thread:
+            client.flush()
+        thread.assert_not_called()
+        client.shutdown()
+
+    def test_exit_flushes_the_lanes_before_an_inline_span_flush(self, fake_timers):
+        client = make_client(traces={}, sync_mode=False)
+        client.start_span("x").end()
+        order = []
+        client._traces.flush = lambda timeout: order.append("spans")
+        for lane in client._lanes:
+            lane.flush = lambda timeout, _lane=lane: order.append("lanes")
+        with (
+            mock.patch("posthog.client.threading.Thread") as thread,
+            mock.patch("posthog.client._atexit_deadline", None),
+        ):
+            thread.return_value.start.side_effect = RuntimeError("no threads")
+            client._atexit()
+        assert order[0] == "lanes"
+        assert order[-1] == "spans"
+        client.shutdown()
+
     def test_shutdown_flushes_pending_spans(self):
         client = make_client(traces={})
         client.start_span("x").end()
@@ -444,33 +474,27 @@ def test_shutdown_flushes_pending_spans(self):
         assert client._traces._exporter._queue == []
 
     def test_shutdown_bounds_the_final_span_flush_and_warns_about_the_rest(
-        self, no_timers, caplog
+        self, fake_timers, caplog
     ):
         caplog.set_level("WARNING", logger="posthog")
         client = make_client(traces={"max_export_batch_size": 1})
         for name in ("a", "b", "c"):
             client.start_span(name).end()
         requests = []
-
-        def slow_send(pipeline_client, payload):
-            requests.append(payload)
-            time.sleep(0.2)
-            return OK
-
-        client._traces._exporter._send = slow_send
+        client._traces._exporter._send = slow_send(requests)
         with mock.patch("posthog.client._TRACES_SHUTDOWN_FLUSH_SECONDS", 0.05):
             client.shutdown()
         assert len(requests) == 1
         assert any("Discarding 2 span(s)" in r.getMessage() for r in caplog.records)
 
-    def test_tracing_is_inert_after_shutdown(self, no_timers):
+    def test_tracing_is_inert_after_shutdown(self, fake_timers):
         client = make_client(traces={})
         client.start_span("before").end()
         client.shutdown()
         assert client.start_span("late") is NOOP_SPAN
         assert client._traces._exporter._flush_timer is None
 
-    def test_shutdown_closes_a_pipeline_still_initializing(self, no_timers):
+    def test_shutdown_closes_a_pipeline_still_initializing(self, fake_timers):
         client = make_client(traces={})
         shutdown = threading.Thread(target=client.shutdown)
         resolve = posthog.client.resolve_traces_config
@@ -488,13 +512,13 @@ def resolve_while_shutting_down(*args):
         assert client._traces._closed
         assert client._traces._exporter._flush_timer is None
 
-    def test_tracing_never_starts_after_shutdown(self, no_timers):
+    def test_tracing_never_starts_after_shutdown(self, fake_timers):
         client = make_client(traces={})
         client.shutdown()
         assert client.start_span("late") is NOOP_SPAN
         assert client._traces is None
 
-    def test_exit_drains_spans_the_timer_would_have_sent(self, no_timers):
+    def test_exit_drains_spans_the_timer_would_have_sent(self, fake_timers):
         client = make_client(traces={}, sync_mode=False)
         client.start_span("x").end()
         session = mock_session()
@@ -516,7 +540,7 @@ def test_exit_drains_spans_the_timer_would_have_sent(self, no_timers):
         assert session.post.called
 
     def test_exit_flushes_spans_alongside_events_that_use_up_the_budget(
-        self, no_timers
+        self, fake_timers
     ):
         client = make_client(traces={}, sync_mode=False)
         client.start_span("x").end()
@@ -535,7 +559,7 @@ def slow_lane_flush(timeout_seconds):
         assert session.post.called
         client.shutdown()
 
-    def test_exit_does_not_wait_on_a_hung_span_request(self, no_timers, caplog):
+    def test_exit_does_not_wait_on_a_hung_span_request(self, fake_timers, caplog):
         caplog.set_level("WARNING", logger="posthog")
         client = make_client(traces={}, sync_mode=False)
         client.start_span("x").end()
@@ -570,7 +594,7 @@ def hung_post(*args, **kwargs):
         "sync_mode, hook", [(False, "_atexit"), (True, "_atexit_spans")]
     )
     def test_exit_warns_about_spans_it_could_not_send(
-        self, no_timers, caplog, sync_mode, hook
+        self, fake_timers, caplog, sync_mode, hook
     ):
         caplog.set_level("WARNING", logger="posthog")
         client = make_client(traces={}, sync_mode=sync_mode)
@@ -591,18 +615,13 @@ def test_exit_warns_about_spans_it_could_not_send(
 
     @pytest.mark.parametrize("sync_mode", [True, False])
     def test_an_app_exit_hook_registered_earlier_still_gets_to_flush_spans(
-        self, no_timers, caplog, sync_mode
+        self, fake_timers, caplog, sync_mode
     ):
         caplog.set_level("WARNING", logger="posthog")
         hooks = []
         holder = {}
         requests = []
 
-        def slow_send(pipeline_client, payload):
-            requests.append(payload)
-            time.sleep(0.2)
-            return OK
-
         with (
             mock.patch("posthog.client.atexit.register", side_effect=hooks.append),
             mock.patch("posthog.client._ATEXIT_FLUSH_TIMEOUT_SECONDS", 0.1),
@@ -614,7 +633,7 @@ def slow_send(pipeline_client, payload):
                 traces={"max_export_batch_size": 1}, sync_mode=sync_mode
             )
             client.start_span("a").end()
-            client._traces._exporter._send = slow_send
+            client._traces._exporter._send = slow_send(requests)
             client.start_span("b").end()
             client.start_span("c").end()
             assert len(hooks) == 2