fix(traces): detach before ending and record only Exception on scoped exit

A generator closed early or a cancelled task raised GeneratorExit or
CancelledError through __exit__ and shipped as an error span. The scoped
form now records an Exception only, still ending on every exit. The span
is detached before end() so nothing started from the on_end path nests
under a span that is over. end() flips its flag under the handle's lock,
and a deactivate that finds no token for the context says so at debug.

Author: turnipdabeets

posthog/test/tracing/test_span.py

@@ -1,4 +1,5 @@
 import asyncio
+import contextvars
 import gc
 import sys
 import threading
@@ -360,6 +361,32 @@ def test_records_a_raised_exception_and_reraises_it_unchanged(self):
             "exception.message": "boom",
         }
 
+    @pytest.mark.parametrize(
+        "error", [GeneratorExit(), asyncio.CancelledError(), KeyboardInterrupt()]
+    )
+    def test_ends_without_recording_a_base_exception_that_is_control_flow(self, error):
+        records: list = []
+        with pytest.raises(type(error)):
+            with make_span(records):
+                raise error
+        assert len(records) == 1
+        assert records[0].status is None
+        assert records[0].events == []
+
+    def test_a_closed_generator_holding_a_span_is_not_an_error(self):
+        records: list = []
+
+        def stream():
+            with make_span(records):
+                yield 1
+                yield 2
+
+        consumer = stream()
+        next(consumer)
+        consumer.close()
+        assert len(records) == 1
+        assert records[0].status is None
+
     def test_treats_an_explicit_ok_status_as_final_when_the_block_raises(self):
         records: list = []
         with pytest.raises(ValueError):
@@ -377,6 +404,25 @@ def test_deactivates_even_when_the_block_raises(self):
                 raise RuntimeError("x")
         assert active.get() is None
 
+    def test_is_no_longer_active_when_on_end_runs(self):
+        active: ContextVar = ContextVar("active", default=None)
+        seen: list = []
+        span = make_span(active_var=active, on_end=lambda _: seen.append(active.get()))
+        with span:
+            pass
+        assert seen == [None]
+
+    def test_exiting_in_a_context_that_never_entered_leaves_it_active(self, caplog):
+        active: ContextVar = ContextVar("active", default=None)
+        span = make_span(active_var=active)
+        span.__enter__()
+        with caplog.at_level("DEBUG", logger="posthog"):
+            contextvars.copy_context().run(span.__exit__, None, None, None)
+        assert active.get() is span
+        assert "never entered it" in caplog.text
+        span.__exit__(None, None, None)
+        assert active.get() is None
+
     def test_ending_inside_the_block_does_not_double_record(self):
         records: list = []
         with make_span(records) as span:

posthog/tracing/_span.py

@@ -86,6 +86,10 @@ def _deactivate(self) -> None:
                     continue
                 del self._tokens[index]
                 return
+        log.debug(
+            "Span exited in a context that never entered it; it stays active "
+            "where it was entered"
+        )
 
 
 class PassThroughSpan(_Activatable, NoopSpan):
@@ -299,10 +303,11 @@ def _child_context(self) -> ParentContext:
         )
 
     def end(self, end_time: Optional[SpanTimeInput] = None) -> None:
-        if self._ended:
-            log.debug("Ignoring end() on a span that has already ended")
-            return
-        self._ended = True
+        with self._tokens_lock:
+            if self._ended:
+                log.debug("Ignoring end() on a span that has already ended")
+                return
+            self._ended = True
 
         derived = self._now_ns()
         resolved = resolve_supplied_ns(end_time, derived, "end time")
@@ -331,9 +336,11 @@ def __enter__(self) -> "Span":
         return self
 
     def __exit__(self, exc_type, exc, tb) -> None:
-        try:
-            if exc is not None and not self._ended:
-                self._record_exception(exc, keep_ok=True)
-            self.end()
-        finally:
-            self._deactivate()
+        # Detached before it ends, so a span started from the on_end path
+        # does not become a child of one that is already over.
+        self._deactivate()
+        # GeneratorExit, CancelledError and the like are control flow, not
+        # failures of the span's work.
+        if isinstance(exc, Exception) and not self._ended:
+            self._record_exception(exc, keep_ok=True)
+        self.end()

posthog/tracing/span.py

@@ -12,7 +12,10 @@ class Span:
     Every method is safe to call on any handle, including after ``end()`` and
     on the inert handles returned when tracing is off, so calling code never
     branches on whether tracing is running. Entering a handle (``with span:``)
-    makes it the active span for the block and ends it on exit.
+    makes it the active span for the block and ends it on exit, recording an
+    ``Exception`` raised inside it. A ``BaseException`` that is not an
+    ``Exception`` (``GeneratorExit``, ``CancelledError``, ``KeyboardInterrupt``)
+    still ends the span but is not recorded as a failure.
     """
 
     def set_attribute(self, key: str, value: Any) -> "Span":
@@ -65,7 +68,7 @@ def set_status(self, code: str, message: Optional[str] = None) -> "Span":
         """Set the span status to ``"ok"`` or ``"error"``. Ignored after ``end()``.
 
         Unset by default. Any other ``code`` is ignored. ``ok`` is final for the
-        scoped form: an exception raised inside ``with span:`` does not override
+        scoped form: an ``Exception`` raised inside ``with span:`` does not override
         it. Returns the span, so calls chain.
 
         Examples:
@@ -80,8 +83,8 @@ def record_exception(self, exception: BaseException) -> "Span":
 
         Ignored after ``end()``. The event carries ``exception.type`` and
         ``exception.message``. Returns the span, so calls chain. Inside
-        ``with span:`` a raised exception is recorded automatically, so this
-        is for exceptions that are caught and handled.
+        ``with span:`` a raised ``Exception`` is recorded automatically, so
+        this is for exceptions that are caught and handled.
 
         Examples:
             ```python