ref: Flush trace bucket when segment span finishes

Author: alexander-alderman-webb

sentry_sdk/_span_batcher.py

@@ -100,7 +100,7 @@ def _flush_loop(self) -> None:
                 self._flush()
                 self._last_full_flush = time.monotonic()
 
-    def add(self, span: "SpanJSON") -> None:
+    def add(self, span: "SpanJSON", flush_trace_bucket: "bool" = False) -> None:
         # Bail out if the current thread is already executing batcher code.
         # This prevents deadlocks when code running inside the batcher (e.g.
         # _add_to_envelope during flush, or _flush_event.wait/set) triggers
@@ -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
+                    flush_trace_bucket
+                    or size + 1 >= self.MAX_BEFORE_FLUSH
                     or self._running_size[span["trace_id"]]
                     >= self.MAX_BYTES_BEFORE_FLUSH
                 ):

sentry_sdk/client.py

@@ -1270,6 +1270,12 @@ def _capture_telemetry(
                 if serialized is None:
                     return
 
+                if ty == "log":
+                    self.log_batcher.add(serialized)  # type: ignore
+
+                elif ty == "metric":
+                    self.metrics_batcher.add(serialized)  # type: ignore
+
             elif ty == "span" and isinstance(telemetry, StreamedSpan):
                 # Reset the span to its original value before we attempted
                 # to call the `before_send_span` callback
@@ -1292,21 +1298,12 @@ def _capture_telemetry(
 
                 serialized = telemetry._to_json()
 
-        batcher = None
-        if ty == "log":
-            batcher = self.log_batcher
-
-        elif ty == "metric":
-            batcher = self.metrics_batcher
-
-        elif ty == "span":
-            # We need a reference to the segment span in the batcher to populate
-            # the dynamic sampling context (DSC)
-            serialized["_segment_span"] = telemetry._segment  # type: ignore
-            batcher = self.span_batcher
-
-        if batcher is not None:
-            batcher.add(serialized)  # type: ignore
+                # We need a reference to the segment span in the batcher to populate
+                # the dynamic sampling context (DSC)
+                serialized["_segment_span"] = telemetry._segment  # type: ignore
+                self.span_batcher.add(
+                    serialized, flush_trace_bucket=telemetry._is_segment()
+                )  # type: ignore
 
     def _capture_log(self, log: "Optional[Log]", scope: "Scope") -> None:
         self._capture_telemetry(log, "log", scope)

tests/tracing/test_span_batcher.py

@@ -432,6 +432,25 @@ def test_transport_format(sentry_init, capture_envelopes):
         assert value["type"] in ("string", "boolean", "integer", "double", "array")
 
 
+def test_trace_bucket_flushes_when_segment_ends(
+    sentry_init, capture_items, monkeypatch
+):
+    """All currently completed spans in a trace are flushed when the segment is finished."""
+    monkeypatch.setattr(SpanBatcher, "FLUSH_WAIT_TIME", 100000)
+
+    sentry_init(traces_sample_rate=1.0, trace_lifecycle="stream")
+    items = capture_items("span")
+
+    with sentry_sdk.traces.start_span(name="segment span"):
+        with sentry_sdk.traces.start_span(name="child"):
+            pass
+
+    time.sleep(0.1)
+
+    assert len(items) == 3
+    assert items[0].payload["name"] == "span"
+
+
 @pytest.mark.skipif(
     sys.platform == "win32"
     or not hasattr(os, "fork")