fix(traces): keep the module-level setup from turning an explicit traces config off

setup() runs on every module-level call and re-applied posthog.traces
whenever no pipeline existed yet, so a default client built with its own
traces config lost it on the first capture(), and a failed init was
retried and re-logged on every call. The module option now applies only
where the client has none, and a failed init latches as False. Shutdown
unregisters the sync-mode exit drain so the client is collectable, the
flush docstring states its worst case, Span is re-exported for pyright
strict and checked in CI, the fork ContextVar is required, and the
start_span docstring says what a forked child inherits.

Author: turnipdabeets

.github/scripts/check_strict_types.sh

@@ -26,8 +26,11 @@ all_flags: dict[str, FlagValue] | None = posthog.get_all_flags("user", groups=gr
 enabled: bool | None = posthog.feature_enabled("flag", "user", groups=groups)
 payload: object | None = client.get_feature_flag_payload("flag", "user", groups=groups)
 evaluations: FeatureFlagEvaluations = posthog.evaluate_flags(123, groups=groups)
+span: posthog.Span = client.start_span("job")
+active: posthog.Span | None = posthog.get_active_span()
+span.end()
 
-_ = (flag_value, all_flags, enabled, payload, evaluations)
+_ = (flag_value, all_flags, enabled, payload, evaluations, active)
 PY
 
 "$tmp/.venv/bin/python" - <<'PY' > "$tmp/public_api_access.py"

posthog/__init__.py

@@ -12,7 +12,7 @@
 from posthog.capture_compression import CaptureCompression as CaptureCompression
 from posthog.capture_mode import CaptureMode as CaptureMode
 from posthog.client import Client
-from posthog.tracing.span import Span
+from posthog.tracing.span import Span as Span
 from posthog.async_client import AsyncClient as AsyncClient
 from posthog.async_client import AsyncPosthog as AsyncPosthog
 from posthog.exception_capture import ExceptionCapture
@@ -1393,7 +1393,9 @@ def setup() -> Client:
     # already forced setup()) still applies until the metrics API is first used.
     if default_client._metrics is None:
         default_client._metrics_config = metrics
-    if default_client._traces is None:
+    # traces=None means off, so the module option applies only where the client
+    # has none; False is the latch of an init that failed and stays off.
+    if traces is not None and default_client._traces_config is None:
         default_client._traces_config = traces
 
     return default_client

posthog/client.py

@@ -2539,7 +2539,7 @@ def _traces_pipeline(self) -> Optional[PostHogTraces]:
                         # Off rather than defaults: defaults would drop a
                         # before_span_send hook and export unscrubbed spans.
                         self.log.exception("Error initializing traces; tracing is off")
-                        self._traces_config = None
+                        self._traces_config = False
         return self._traces
 
     def _tracing_context(self) -> Dict[str, Optional[str]]:
@@ -2578,6 +2578,8 @@ def start_span(
             attributes: Initial attributes.
             parent: A span handle, or an inbound W3C ``traceparent`` header
                 value to continue a remote trace. Defaults to the active span.
+                A forked child starts with no active span; pass the parent
+                span to continue a trace across a fork.
             tracestate: The inbound ``tracestate`` header accompanying a
                 ``traceparent`` string ``parent``; preserved and propagated.
             start_time: A ``datetime`` or epoch seconds, to backdate the span.
@@ -2638,7 +2640,9 @@ def flush(self, timeout_seconds: Optional[float] = 10) -> None:
                 Queued spans are sent at the same time, within the same
                 budget: at least one span request is attempted even when the
                 budget is already spent, no further one starts once it is, and
-                each request is bounded by ``timeout``.
+                each request is bounded by ``timeout``. The wait for that first
+                request is not cut short, so a flush can take up to
+                ``timeout_seconds`` plus ``timeout`` in the worst case.
 
         Examples:
             ```python
@@ -2893,6 +2897,8 @@ def _shutdown_once(self, errors: list[Exception]) -> None:
             self._run_lifecycle_cleanup(
                 "Failed to close traces on shutdown", traces.close, errors
             )
+            # The sync-mode exit drain, so a shut-down client is collectable.
+            atexit.unregister(self._atexit_spans)
         self._join_once(errors, flush_queues=False, lanes_prepared=True)
         self._run_lifecycle_cleanup(
             "Failed to clear feature flag deduplication state on shutdown",

posthog/test/tracing/test_client_traces.py

@@ -125,6 +125,16 @@ def test_a_bad_traces_config_degrades_to_defaults(self):
         assert isinstance(client.start_span("x"), RecordingSpan)
         client.shutdown()
 
+    def test_a_failed_init_is_not_retried_on_the_next_call(self):
+        client = make_client(traces={})
+        with mock.patch(
+            "posthog.client.resolve_traces_config", side_effect=RuntimeError("no")
+        ) as resolve:
+            assert client.start_span("x") is NOOP_SPAN
+            assert client.start_span("y") is NOOP_SPAN
+        assert resolve.call_count == 1
+        client.shutdown()
+
     def test_never_starts_a_pipeline_on_a_client_without_traces(self):
         client = make_client()
         client.flush()
@@ -623,6 +633,13 @@ def test_sync_mode_registers_the_span_exit_drain_when_tracing_starts(self):
         register.assert_called_once_with(client._atexit_spans)
         client.shutdown()
 
+    def test_shutdown_unregisters_the_sync_mode_exit_drain(self):
+        client = make_client(traces={}, sync_mode=True)
+        client.start_span("x").end()
+        with mock.patch("posthog.client.atexit.unregister") as unregister:
+            client.shutdown()
+        unregister.assert_called_once_with(client._atexit_spans)
+
     @pytest.mark.parametrize("traces", [False, None])
     def test_sync_mode_without_tracing_registers_no_exit_hook(self, traces):
         with mock.patch("posthog.client.atexit.register") as register:
@@ -735,6 +752,25 @@ def body():
 
         self._with_module_client(None, body)
 
+    def test_setup_leaves_an_explicitly_configured_default_client_alone(self):
+        def body():
+            posthog.default_client = make_client(traces={"service_name": "explicit"})
+            posthog.setup()
+            assert isinstance(posthog.start_span("x"), RecordingSpan)
+
+        self._with_module_client(None, body)
+
+    def test_setup_does_not_retry_a_traces_init_that_failed(self):
+        def body():
+            with mock.patch(
+                "posthog.client.resolve_traces_config", side_effect=RuntimeError("no")
+            ) as resolve:
+                assert posthog.start_span("x") is NOOP_SPAN
+                assert posthog.start_span("y") is NOOP_SPAN
+            assert resolve.call_count == 1
+
+        self._with_module_client({"service_name": "broken"}, body)
+
     def test_module_start_span_is_inert_without_config(self):
         def body():
             assert posthog.start_span("x") is NOOP_SPAN

posthog/test/tracing/test_export.py

@@ -1,4 +1,5 @@
 import threading
+from contextvars import ContextVar
 from types import SimpleNamespace
 from unittest import mock
 
@@ -903,7 +904,7 @@ def test_a_forked_child_drops_the_inherited_queue_and_timer(self):
         pipeline.start_span("parent-span").end()
         assert queued(pipeline) and pipeline._exporter._flush_timer is not None
         pipeline._exporter._max_export_batch_size = 1
-        pipeline.reinit_after_fork()
+        pipeline.reinit_after_fork(ContextVar("child-active", default=None))
         assert queued(pipeline) == []
         assert pipeline._exporter._flush_timer is None
         assert (

posthog/test/tracing/test_pipeline.py

@@ -1,6 +1,7 @@
 import gc
 import logging
 import threading
+from contextvars import ContextVar
 import time
 import warnings
 import weakref
@@ -582,14 +583,14 @@ def test_close_makes_later_spans_inert_and_closes_the_exporter(self):
     def test_a_forked_child_drops_the_parents_live_spans(self):
         pipeline, exporter, _ = make()
         pipeline.start_span("live")
-        pipeline.reinit_after_fork()
+        pipeline.reinit_after_fork(ContextVar("child-active", default=None))
         assert pipeline._live_spans == {}
         assert exporter.reinitialized
 
     def test_reinit_after_fork_replaces_locks_without_acquiring_them(self):
         pipeline, _, _ = make()
         pipeline._lock.acquire()
-        pipeline.reinit_after_fork()
+        pipeline.reinit_after_fork(ContextVar("child-active", default=None))
         assert not pipeline._lock.locked()
         pipeline.start_span("a").end()
 

posthog/tracing/_pipeline.py

@@ -114,12 +114,11 @@ def warn_if_queued(self) -> None:
         self._exporter.warn_if_queued()
         self._drops.warn_if_due(force=True)
 
-    def reinit_after_fork(self, active_var: Optional[ContextVar] = None) -> None:
+    def reinit_after_fork(self, active_var: ContextVar) -> None:
         # Runs in the forked child before user code; the parent's spans stay
-        # with the parent.
+        # with the parent, and the child's active span is the fresh var.
         self._lock = threading.Lock()
-        if active_var is not None:
-            self._active_var = active_var
+        self._active_var = active_var
         self._live_spans.clear()
         self._drops.reinit_after_fork()
         self._exporter.reinit_after_fork()