fix(traces): close the shutdown and failure gaps in span creation

Re-check closed under the lock so a close() that lands mid-start cannot
reserve a live span after the registry was cleared. Release the reserved
slot when building the span fails, so it does not wait for age eviction.
Contain a raising logging handler inside the drop warning, and reset the
warning throttle in a forked child with the rest of its state.

Author: turnipdabeets

posthog/test/tracing/test_pipeline.py

@@ -1,4 +1,5 @@
 import gc
+import logging
 import threading
 import time
 import weakref
@@ -15,6 +16,7 @@
     make,
     queued,
 )
+from posthog.tracing import _pipeline as pipeline_module
 from posthog.tracing import _span as span_module
 from posthog.tracing._drops import DropLog
 from posthog.tracing._span import NOOP_SPAN, PassThroughSpan, RecordingSpan
@@ -392,8 +394,30 @@ def test_drops_a_span_whose_client_was_disabled_mid_trace_without_raising(self):
         assert queued(pipeline) == []
         assert pipeline._live_spans == {}
 
+    def test_a_close_that_lands_mid_start_makes_the_span_inert(self):
+        pipeline, _, _ = make()
+        resolve = pipeline._resolve_parent
+
+        def close_then_resolve(parent, tracestate):
+            pipeline.close()
+            return resolve(parent, tracestate)
+
+        with mock.patch.object(pipeline, "_resolve_parent", close_then_resolve):
+            span = pipeline.start_span("late")
+        assert span is NOOP_SPAN
+        assert pipeline._live_spans == {}
+
 
 class TestLiveSpanBounds:
+    def test_returns_the_slot_when_building_the_span_fails(self):
+        pipeline, _, _ = make(max_live_spans=1)
+        with mock.patch.object(
+            pipeline_module, "copy_user_attributes", side_effect=RuntimeError("no")
+        ):
+            assert pipeline.start_span("a") is NOOP_SPAN
+        assert pipeline._live_spans == {}
+        assert isinstance(pipeline.start_span("b"), RecordingSpan)
+
     def test_returns_an_inert_handle_once_max_live_spans_are_live(self):
         pipeline, _, _ = make(max_live_spans=2)
         a, b = pipeline.start_span("a"), pipeline.start_span("b")
@@ -456,6 +480,32 @@ def test_names_reasons_in_the_order_they_happened(self, caplog):
             "Dropping 4 span(s): the queue is full; before_span_send dropped it"
         )
 
+    def test_a_raising_log_handler_does_not_escape(self):
+        class Raising(logging.Handler):
+            def emit(self, record):
+                raise RuntimeError("handler broke")
+
+        handler = Raising()
+        logging.getLogger("posthog").addHandler(handler)
+        try:
+            drops = DropLog(5)
+            drops.record(1, "the queue is full")
+            drops.warn_if_due(force=True)
+        finally:
+            logging.getLogger("posthog").removeHandler(handler)
+
+    def test_a_forked_child_does_not_wait_out_the_parents_warning_interval(
+        self, caplog
+    ):
+        caplog.set_level("WARNING", logger="posthog")
+        drops = DropLog(5)
+        drops.record(1, "the queue is full")
+        drops.warn_if_due()
+        drops.reinit_after_fork()
+        drops.record(1, "the queue is full")
+        drops.warn_if_due()
+        assert len(caplog.records) == 2
+
 
 class TestCloseAndFork:
     def test_close_counts_spans_still_open_and_drops_them_when_they_end(self, caplog):

posthog/tracing/_drops.py

@@ -41,9 +41,14 @@ def warn_if_due(self, force: bool = False) -> None:
             self._count = 0
             self._reasons.clear()
             self._last_warning_at = now
-        log.warning(message)
+        try:
+            log.warning(message)
+        except Exception:
+            # A raising logging handler must not surface through span creation.
+            pass
 
     def reinit_after_fork(self) -> None:
         self._lock = threading.Lock()
         self._count = 0
         self._reasons.clear()
+        self._last_warning_at = 0.0

posthog/tracing/_pipeline.py

@@ -146,6 +146,9 @@ def _start_span(
         parent_context = self._resolve_parent(parent, tracestate)
 
         with self._lock:
+            if self._is_closed():
+                # close() may have run since the check above.
+                return inert_span(parent, tracestate, self._active_var)
             # Swept first, so a process that leaked its way to the bound
             # recovers once the leaks age out.
             aged = self._evict_aged_spans_locked()
@@ -167,6 +170,25 @@ def _start_span(
             )
             return inert_span(parent, tracestate, self._active_var)
 
+        try:
+            return self._build_span(
+                span_id, name, kind, attributes, parent_context, start_time
+            )
+        except Exception:
+            # The reservation would otherwise hold its slot until age eviction.
+            with self._lock:
+                self._live_spans.pop(span_id, None)
+            raise
+
+    def _build_span(
+        self,
+        span_id: str,
+        name: str,
+        kind: Optional[str],
+        attributes: Optional[Mapping[str, Any]],
+        parent_context: Optional[ParentContext],
+        start_time: Any,
+    ) -> RecordingSpan:
         now_ns = time.time_ns()
         start_ns = resolve_start_ns(start_time, now_ns)
         auto_attributes = self._auto_context_attributes()
@@ -192,6 +214,10 @@ def _start_span(
             else None,
         )
 
+    def _is_closed(self) -> bool:
+        # A method, so the re-check under the lock is not narrowed away.
+        return self._closed
+
     def _resolve_parent(self, parent: Any, tracestate: Any) -> Optional[ParentContext]:
         """An explicit parent, else the active span, else a fresh root."""
         if isinstance(parent, str):