ref: Flush trace buckets when segment spans finish (#7170)

Mark the corresponding bucket as pending when a segment span is added in the span buffer.
This aims to keep memory pressure low. Segment spans typically finish after their children.
Adapt span batcher tests by adding an outer segment span. As a result, the various flush conditions are still exercised since the assertions run before the segment span has finished (finishing the segment span otherwise flushes the buffer as well).

Author: alexander-alderman-webb

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

@@ -2496,12 +2496,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"