fix(traces): keep shutdown and fork isolation airtight for the client

Shutdown reads the pipeline under the same lock initialization publishes
it with, so an initialization in flight is either closed by shutdown or
sees the request and stays off. A forked child gets a fresh active-span
variable: an inherited handle exiting in the child resets the old one,
which would have restored the parent process's outer span.

Author: turnipdabeets

posthog/client.py

@@ -2276,9 +2276,11 @@ def _reinit_after_fork(self):
         if self._metrics is not None:
             self._metrics._reinit_after_fork()
         self._traces_lock = threading.Lock()
+        # A fresh variable: an inherited handle resets the old one on exit,
+        # which would restore the parent process's outer span.
+        self._active_span_var = ContextVar("posthog_active_span", default=None)
         if self._traces is not None:
-            self._traces.reinit_after_fork()
-        self._active_span_var.set(None)
+            self._traces.reinit_after_fork(self._active_span_var)
 
         # If using Redis cache, we must reinitialize to get a fresh connection (fork-safe).
         # If using Memory cache, we keep it as-is to benefit from the inherited warm cache.
@@ -2513,7 +2515,9 @@ def _traces_pipeline(self) -> Optional[PostHogTraces]:
             return self._traces
         if self._traces is None:
             with self._traces_lock:
-                if self._traces is None:
+                # Re-checked under the lock, which shutdown takes before it
+                # reads the pipeline, so neither can miss the other.
+                if self._traces is None and not self._shutdown_requested:
                     try:
                         config = resolve_traces_config(
                             self._traces_config, host_resource_attributes()
@@ -2878,15 +2882,16 @@ def _shutdown_once(self, errors: list[Exception]) -> None:
             self._run_lifecycle_cleanup(
                 "Failed to reset metrics on shutdown", self._metrics.reset, errors
             )
-        if self._traces is not None:
+        with self._traces_lock:
             traces = self._traces
+        if traces is not None:
             self._run_lifecycle_cleanup(
                 "Failed to flush spans on shutdown",
                 lambda: traces.flush(_TRACES_SHUTDOWN_FLUSH_SECONDS),
                 errors,
             )
             self._run_lifecycle_cleanup(
-                "Failed to close traces on shutdown", self._traces.close, errors
+                "Failed to close traces on shutdown", traces.close, errors
             )
         self._join_once(errors, flush_queues=False, lanes_prepared=True)
         self._run_lifecycle_cleanup(

posthog/test/tracing/test_client_traces.py

@@ -460,6 +460,24 @@ def test_tracing_is_inert_after_shutdown(self, no_timers):
         assert client.start_span("late") is NOOP_SPAN
         assert client._traces._exporter._flush_timer is None
 
+    def test_shutdown_closes_a_pipeline_still_initializing(self, no_timers):
+        client = make_client(traces={})
+        shutdown = threading.Thread(target=client.shutdown)
+        resolve = posthog.client.resolve_traces_config
+
+        def resolve_while_shutting_down(*args):
+            shutdown.start()
+            time.sleep(0.05)
+            return resolve(*args)
+
+        with mock.patch(
+            "posthog.client.resolve_traces_config", resolve_while_shutting_down
+        ):
+            client.start_span("racing").end()
+        shutdown.join()
+        assert client._traces._closed
+        assert client._traces._exporter._flush_timer is None
+
     def test_tracing_never_starts_after_shutdown(self, no_timers):
         client = make_client(traces={})
         client.shutdown()
@@ -649,6 +667,22 @@ def test_a_forked_child_does_not_inherit_the_active_span(self):
             assert client.get_active_span() is None
         client.shutdown()
 
+    def test_an_inherited_span_exiting_in_the_child_does_not_restore_its_parent(
+        self,
+    ):
+        client = make_client(traces={})
+        with client.start_span("outer") as outer:
+            with client.start_span("inner"):
+                client._reinit_after_fork()
+            assert client.get_active_span() is None
+            child = client.start_span("child")
+            child.end()
+        (record,) = client._traces._exporter._queue
+        assert record.name == "child"
+        assert record.parent_span_id is None
+        assert record.trace_id != outer.traceparent().split("-")[1]
+        client.shutdown()
+
 
 class TestModuleLevelApi:
     def _with_module_client(self, traces, body):

posthog/tracing/_pipeline.py

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