ref: Flush trace buckets when segment spans finish

sentry_sdk/_span_batcher.py

@@ -129,7 +129,8 @@ def add(self, span: "SpanJSON") -> None:
                 self._running_size[span["trace_id"]] += self._estimate_size(span)
 
                 if (
-                    size + 1 >= self.MAX_BEFORE_FLUSH
+                    span["is_segment"] is True
+                    or size + 1 >= self.MAX_BEFORE_FLUSH
                     or self._running_size[span["trace_id"]]
                     >= self.MAX_BYTES_BEFORE_FLUSH
                 ):

tests/conftest.py

@@ -4,6 +4,7 @@
 import json
 import os
 import socket
+import threading
 import warnings
 from collections import namedtuple
 from contextlib import contextmanager
@@ -293,11 +294,71 @@ def inner(identifier):
     return inner
 
 
+def _install_flush_completion_handshake(client: "sentry_sdk.Client") -> None:
+    """Make batcher.flush() wait for the flusher thread to drain.
+
+    Otherwise, test assertions can run before envelopes are captured.
+    The span batcher flushes pending items asynchronously with the main thread.
+    Flushes triggered by segments finishing are asynchronous, and can collect buckets
+    that would have otherwise been flushed synchronously by `sentry_sdk.flush()`.
+    """
+    batcher = client.span_batcher
+    if batcher is None:
+        return
+
+    orig_flush_raw = batcher._flush
+    orig_flush = batcher.flush
+    lock = threading.Lock()
+    drained_count = 0
+    wake = threading.Event()
+
+    def _flush(*args: "Any", **kwargs: "Any") -> "Any":
+        nonlocal drained_count
+        try:
+            return orig_flush_raw(*args, **kwargs)
+        finally:
+            with lock:
+                drained_count += 1
+            wake.set()
+
+    def flush() -> None:
+        nonlocal drained_count
+        # Re-entrancy guard: if `flush()` is invoked from within a drain (e.g. a
+        # custom transport), waiting on the flusher thread would deadlock, because
+        # the flusher is blocked inside our own handler.
+        if getattr(getattr(batcher, "_active", None), "flag", False):
+            orig_flush()
+            return
+
+        # If the background flusher thread was never started (no spans have
+        # been added), there is no thread to drain and the counter will never
+        # advance. Fall back to the original synchronous flush.
+        if batcher._flusher is None or not batcher._flusher.is_alive():
+            orig_flush()
+            return
+
+        with lock:
+            target = drained_count
+
+        batcher._flush_event.set()
+        while True:
+            with lock:
+                if drained_count > target:
+                    break
+            wake.wait()
+            wake.clear()
+        orig_flush()
+
+    object.__setattr__(batcher, "_flush", _flush)
+    object.__setattr__(batcher, "flush", flush)
+
+
 @pytest.fixture
 def sentry_init(request):
     def inner(*a, **kw):
         kw.setdefault("transport", TestTransport())
         client = sentry_sdk.Client(*a, **kw)
+        _install_flush_completion_handshake(client)
         sentry_sdk.get_global_scope().set_client(client)
 
     if request.node.get_closest_marker("forked"):

tests/integrations/django/test_basic.py

@@ -2499,12 +2499,17 @@ def test_transaction_http_method_custom(
 
         client.get("/nomessage")
         client.options("/nomessage")
-        client.head("/nomessage")
 
         sentry_sdk.flush()
         spans = [item.payload for item in items]
 
         assert spans[2]["attributes"][SPANDATA.HTTP_REQUEST_METHOD] == "OPTIONS"
+
+        client.head("/nomessage")
+
+        sentry_sdk.flush()
+        spans = [item.payload for item in items]
+
         assert spans[5]["attributes"][SPANDATA.HTTP_REQUEST_METHOD] == "HEAD"
     else:
         events = capture_events()

tests/integrations/django/test_cache_module.py

@@ -246,7 +246,6 @@ def test_cache_spans_middleware(
     if span_streaming:
         items = capture_items("span")
 
-        client.get(reverse("not_cached_view"))
         client.get(reverse("not_cached_view"))
 
         sentry_sdk.flush()
@@ -269,6 +268,11 @@ def test_cache_spans_middleware(
         )
         assert "cache.hit" not in spans[1]["attributes"]
         assert spans[1]["attributes"]["cache.item_size"] == 2
+
+        client.get(reverse("not_cached_view"))
+
+        sentry_sdk.flush()
+        spans = [item.payload for item in items]
         # second_event - cache.get
         assert spans[4]["attributes"]["sentry.op"] == "cache.get"
         assert spans[4]["name"].startswith("views.decorators.cache.cache_header.")

tests/integrations/flask/test_flask.py

@@ -1211,20 +1211,31 @@ def test_transaction_or_segment_http_method_custom(
     response = client.options("/nomessage")
     assert response.status_code == 200
 
-    response = client.head("/nomessage")
-    assert response.status_code == 200
-
     if span_streaming:
         sentry_sdk.flush()
         spans = [i.payload for i in items]
-        assert len(spans) == 2
-        (options_segment, head_segment) = spans
+        (options_segment,) = spans
         assert options_segment["attributes"]["http.request.method"] == "OPTIONS"
+
+        response = client.head("/nomessage")
+        assert response.status_code == 200
+
+        sentry_sdk.flush()
+        spans = [i.payload for i in items]
+        assert len(spans) == 2
+        (_, head_segment) = spans
+
         assert head_segment["attributes"]["http.request.method"] == "HEAD"
     else:
-        assert len(events) == 2
-        (event1, event2) = events
+        (event1,) = events
         assert event1["request"]["method"] == "OPTIONS"
+
+        response = client.head("/nomessage")
+        assert response.status_code == 200
+
+        assert len(events) == 2
+        (_, event2) = events
+
         assert event2["request"]["method"] == "HEAD"