fix: native rust fix for HAMT/ContextVar crash [backport 4.8]

.gitlab/benchmarks/bp-runner.microbenchmarks.fail-on-breach.template.yml

@@ -1070,7 +1070,7 @@ experiments:
               - max_rss_usage < 38.75 MB
           - name: tracer-small
             thresholds:
-              - execution_time < 0.39 ms
+              - execution_time < 0.40 ms
               - max_rss_usage < 38.75 MB
 
           # codeprovenancefork

ddtrace/_trace/provider.py

@@ -1,5 +1,6 @@
 import abc
 import contextvars
+import sys
 from typing import Any
 from typing import Optional
 from typing import Union
@@ -19,6 +20,21 @@
 )
 
 
+# PyContextVar_Set is not atomic before CPython 3.12: an allocation during
+# the HAMT rebuild can trigger a cyclic GC pass that frees a node still in use,
+# crashing with SEGV_MAPERR. On affected versions we route the set through a
+# native helper that secures the context's storage during the call.
+# CPython 3.12+ is unaffected, so use the plain (and faster) set there.
+if sys.version_info < (3, 12):
+    from ddtrace.internal.native._native import safe_contextvar_set
+
+    def _activate_contextvar(ctx: Optional[ActiveTrace]) -> None:
+        safe_contextvar_set(_DD_CONTEXTVAR, ctx)
+
+else:
+    _activate_contextvar = _DD_CONTEXTVAR.set  # type: ignore[assignment]
+
+
 class BaseContextProvider(metaclass=abc.ABCMeta):
     """
     A ``ContextProvider`` is an interface that provides the blueprint
@@ -65,7 +81,7 @@ def _has_active_context(self) -> bool:
 
     def activate(self, ctx: Optional[ActiveTrace]) -> None:
         """Makes the given context active in the current execution."""
-        _DD_CONTEXTVAR.set(ctx)
+        _activate_contextvar(ctx)
         super(DefaultContextProvider, self).activate(ctx)
 
     def active(self) -> Optional[ActiveTrace]:

ddtrace/internal/native/_native.pyi

@@ -1,3 +1,4 @@
+import contextvars
 from enum import Enum
 from typing import Any
 from typing import Iterator
@@ -675,3 +676,5 @@ class config:
     def set_128_bit_trace_id_enabled(val: bool) -> None:
         """Set whether 128-bit trace ID generation is enabled."""
         ...
+
+def safe_contextvar_set(var: contextvars.ContextVar[Any], value: Any) -> None: ...

releasenotes/notes/internal-add-contextvar-helper-d73542f1c6cfceac.yaml

@@ -0,0 +1,4 @@
+fixes:
+  - |
+    tracing: A rare crash happening on versions of CPython prior to 3.12 has been fixed.
+

src/native/contextvar.rs

@@ -0,0 +1,61 @@
+use pyo3::ffi;
+use pyo3::prelude::*;
+
+/// Set a `ContextVar` while protecting against a CPython crash.
+///
+/// `PyContextVar_Set` rebuilds the thread state's context HAMT and is not
+/// atomic: the allocations it performs while building the new trie can trigger
+/// a cyclic GC pass, and it then Py_DECREF's the previous HAMT root, which can
+/// start a dealloc cascade. If a node still referenced by the in-flight rebuild
+/// is collected during that window, the next access to it faults with
+/// `SEGV_MAPERR`.
+///
+/// PyContext_CopyCurrent returns a fresh context object that shares -- and
+/// holds a strong reference to -- the current context's variable storage. Keeping
+/// that snapshot alive across the set pins the whole storage graph as
+/// GC-reachable, so neither the Py_DECREF cascade nor a concurrent GC pass can
+/// free a node that is still in use. The snapshot (and the discarded token) are
+/// released once the set has completed and the context is consistent again.
+#[pyfunction]
+pub fn safe_contextvar_set(
+    py: Python<'_>,
+    var: &Bound<'_, PyAny>,
+    value: &Bound<'_, PyAny>,
+) -> PyResult<()> {
+    // SAFETY: the GIL (`py`) is held for every call below and is never released,
+    // so the thread state and its context cannot change underneath us.
+    unsafe {
+        let snapshot = ffi::PyContext_CopyCurrent();
+        if snapshot.is_null() {
+            // Snapshotting the current context failed (e.g. out of memory); an
+            // exception is already set. Propagate it rather than crashing.
+            return Err(PyErr::take(py).unwrap_or_else(|| {
+                PyErr::new::<pyo3::exceptions::PyRuntimeError, _>("PyContext_CopyCurrent failed")
+            }));
+        }
+
+        let token = ffi::PyContextVar_Set(var.as_ptr(), value.as_ptr());
+
+        // Capture any error before the decrefs below can perturb interpreter state.
+        let err = if token.is_null() {
+            PyErr::take(py)
+        } else {
+            None
+        };
+
+        // The returned token is never used (we don't support reset here), so drop
+        // it. Then release the snapshot now that the context is consistent again.
+        ffi::Py_XDECREF(token);
+        ffi::Py_DECREF(snapshot);
+
+        if let Some(err) = err {
+            return Err(err);
+        }
+    }
+    Ok(())
+}
+
+pub fn register_contextvar(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_function(wrap_pyfunction!(safe_contextvar_set, m)?)?;
+    Ok(())
+}

src/native/lib.rs

@@ -3,6 +3,7 @@ mod crashtracker;
 #[cfg(feature = "profiling")]
 pub use datadog_profiling_ffi::*;
 mod config;
+mod contextvar;
 mod data_pipeline;
 #[cfg(feature = "stats")]
 mod ddsketch;
@@ -49,6 +50,7 @@ fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(library_config::store_metadata))?;
     data_pipeline::register_data_pipeline(m)?;
     span::register_native_span(m)?;
+    contextvar::register_contextvar(m)?;
     rand::register_rand(m)?;
     m.add_function(wrap_pyfunction!(utils::flatten_key_value, m)?)?;
     m.add_function(wrap_pyfunction!(utils::is_sequence, m)?)?;