diff --git a/ddtrace/_trace/context.py b/ddtrace/_trace/context.py index bf1c277f82e..6f1fbf64e7e 100644 --- a/ddtrace/_trace/context.py +++ b/ddtrace/_trace/context.py @@ -13,6 +13,7 @@ from ddtrace.internal.constants import W3C_TRACEPARENT_KEY from ddtrace.internal.constants import W3C_TRACESTATE_KEY from ddtrace.internal.logger import get_logger +from ddtrace.internal.native._native import ContextData from ddtrace.internal.threads import RLock from ddtrace.internal.utils.http import w3c_get_dd_list_member as _w3c_get_dd_list_member @@ -22,7 +23,7 @@ Optional[int], # span_id dict[str, str], # _meta dict[str, NumericType], # _metrics - list[SpanLink], # span_links + list[SpanLink], # span_links dict[str, Any], # baggage bool, # is_remote bool, # _reactivate @@ -34,21 +35,13 @@ log = get_logger(__name__) -class Context(object): +class Context(ContextData): """Represents the state required to propagate a trace across execution boundaries. """ __slots__ = [ - "trace_id", - "span_id", "_lock", - "_meta", - "_metrics", - "_span_links", - "_baggage", - "_is_remote", - "_reactivate", "__weakref__", ] @@ -65,23 +58,12 @@ def __init__( baggage: Optional[dict[str, Any]] = None, is_remote: bool = True, ): - self._meta: dict[str, str] = meta if meta is not None else {} - self._metrics: dict[str, NumericType] = metrics if metrics is not None else {} - self._baggage: dict[str, Any] = baggage if baggage is not None else {} - - self.trace_id: Optional[int] = trace_id - self.span_id: Optional[int] = span_id - self._is_remote: bool = is_remote - self._reactivate: bool = False - + # ContextData.__new__ already populated trace_id/span_id/_meta/_metrics/ + # _baggage/_span_links/_is_remote/_reactivate. if dd_origin is not None and _DD_ORIGIN_INVALID_CHARS_REGEX.search(dd_origin) is None: self._meta[_ORIGIN_KEY] = dd_origin if sampling_priority is not None: self._metrics[_SAMPLING_PRIORITY_KEY] = sampling_priority - if span_links is not None: - self._span_links = span_links - else: - self._span_links = [] if lock is not None: self._lock = lock @@ -223,23 +205,27 @@ def set_baggage_item(self, key: str, value: Any) -> None: """Sets a baggage item in this span context. Note that this operation mutates the baggage of this span context """ - self._baggage[key] = value + with self._lock: + self._baggage[key] = value def copy(self, trace_id: int, span_id: int) -> "Context": """Return a shallow copy of the context with the given correlation IDs.""" - # PERF: copy() is run once per child span, use __new__ + direct assignment to avoid the - # overhead in __init__'s kwargs packing, dd_origin regex, and other processing. This - # optimization holds true if we trust that this data has been validated already. - ctx = Context.__new__(Context) - ctx._meta = self._meta - ctx._metrics = self._metrics - ctx._baggage = self._baggage + # PERF: copy() is run once per child span. Construct via ContextData.__new__ + # directly, passing the shared _meta/_metrics/_baggage references straight into + # native construction, to avoid the overhead of __init__'s dd_origin regex and + # other processing. This optimization holds true if we trust that this data has + # been validated already. + with self._lock: + ctx = Context.__new__( + Context, + trace_id=trace_id, + span_id=span_id, + meta=self._meta, + metrics=self._metrics, + baggage=self._baggage, + is_remote=False, + ) ctx._lock = self._lock - ctx.trace_id = trace_id - ctx.span_id = span_id - ctx._is_remote = False - ctx._reactivate = False - ctx._span_links = [] return ctx def _with_baggage_item(self, key: str, value: Any) -> "Context": @@ -265,14 +251,20 @@ def get_all_baggage_items(self) -> dict[str, Any]: def remove_baggage_item(self, key: str) -> None: """Remove a baggage item from this span context.""" - if key in self._baggage: - del self._baggage[key] + with self._lock: + if key in self._baggage: + del self._baggage[key] def remove_all_baggage_items(self) -> None: """Removes all baggage items from this span context.""" - self._baggage.clear() + with self._lock: + self._baggage.clear() def __eq__(self, other: Any) -> bool: + # NOTE: span_id/_reactivate are deliberately excluded. A Context compares + # equal to any per-span copy() of itself (differing only in span_id) -- + # e.g. Span.context builds one such copy per child span -- so equality + # here means "same trace-level state", not "same span". if isinstance(other, Context): with self._lock: return ( diff --git a/ddtrace/internal/native/_native.pyi b/ddtrace/internal/native/_native.pyi index 2c5855d60d7..2c0e2f436b2 100644 --- a/ddtrace/internal/native/_native.pyi +++ b/ddtrace/internal/native/_native.pyi @@ -19,6 +19,7 @@ from ddtrace._trace.types import _AttributeValueType ActiveTrace = Union[Span, Context] _SpanDataT = TypeVar("_SpanDataT", bound="SpanData") +_ContextDataT = TypeVar("_ContextDataT", bound="ContextData") class DDSketch: def __init__(self): ... @@ -1000,6 +1001,30 @@ class native_flare: def zip_and_send(self, directory: str, send_action: native_flare.FlareAction) -> None: ... def set_current_log_level(self, level: str) -> None: ... +class ContextData: + trace_id: Optional[int] + span_id: Optional[int] + _meta: dict[str, str] + _metrics: dict[str, Any] + _baggage: dict[str, Any] + _span_links: list[Any] + _is_remote: bool + _reactivate: bool + + def __new__( + cls: type[_ContextDataT], + trace_id: Optional[int] = None, + span_id: Optional[int] = None, + dd_origin: Optional[str] = None, # placeholder for Context.__init__ + sampling_priority: Optional[float] = None, # placeholder for Context.__init__ + meta: Optional[dict[str, str]] = None, + metrics: Optional[dict[str, Any]] = None, + lock: Optional[Any] = None, # placeholder for Context.__init__ + span_links: Optional[list[Any]] = None, + baggage: Optional[dict[str, Any]] = None, + is_remote: bool = True, + ) -> _ContextDataT: ... + class SpanData: name: str service: Optional[str] diff --git a/src/native/context/context_data.rs b/src/native/context/context_data.rs new file mode 100644 index 00000000000..33ec305d91d --- /dev/null +++ b/src/native/context/context_data.rs @@ -0,0 +1,231 @@ +use pyo3::{ + types::{PyAny, PyAnyMethods as _, PyDict, PyList, PyTuple}, + Bound, Py, Python, +}; + +/// Native storage layer for `ddtrace._trace.context.Context`. +/// +/// Mirrors the `SpanData`/`Span` split: this struct owns the raw trace-level +/// fields (trace/span id, meta/metrics/baggage dicts, span links, the +/// remote/reactivate flags) with plain get/set properties; all business logic +/// (sampling_priority/dd_origin/dd_user_id, W3C traceparent/tracestate +/// computation, baggage helper methods, copy(), equality, pickling, and the +/// lock) stays in the pure-Python `Context` subclass. +#[pyo3::pyclass(name = "ContextData", module = "ddtrace.internal._native", subclass)] +#[derive(Default)] +pub struct ContextData { + pub trace_id: Option, + pub span_id: Option, + /// dict[str, str]. `None` only in the brief window before `__new__` finishes + /// or after `__clear__` runs during GC teardown -- getters lazily + /// re-materialize an empty dict, mirroring `SpanData::meta_struct`. + meta: Option>, + /// dict[str, NumericType]. + metrics: Option>, + /// dict[str, Any]. + baggage: Option>, + /// list[SpanLink]. + span_links: Option>, + #[pyo3(get, set, name = "_is_remote")] + pub is_remote: bool, + #[pyo3(get, set, name = "_reactivate")] + pub reactivate: bool, +} + +#[pyo3::pymethods] +impl ContextData { + #[new] + #[allow(unused_variables)] + #[allow(clippy::too_many_arguments)] + #[pyo3(signature = ( + trace_id=None, + span_id=None, + dd_origin=None, // placeholder for Context.__init__, handled there + sampling_priority=None, // placeholder for Context.__init__, handled there + meta=None, + metrics=None, + lock=None, // placeholder for Context.__init__, handled there + span_links=None, + baggage=None, + is_remote=true, + *args, + **kwargs + ))] + pub fn __new__<'p>( + py: Python<'p>, + trace_id: Option<&Bound<'p, PyAny>>, + span_id: Option<&Bound<'p, PyAny>>, + dd_origin: Option<&Bound<'p, PyAny>>, // placeholder, not used + sampling_priority: Option<&Bound<'p, PyAny>>, // placeholder, not used + meta: Option<&Bound<'p, PyDict>>, + metrics: Option<&Bound<'p, PyDict>>, + lock: Option<&Bound<'p, PyAny>>, // placeholder, not used + span_links: Option<&Bound<'p, PyList>>, + baggage: Option<&Bound<'p, PyDict>>, + is_remote: bool, + // Accept *args/**kwargs so subclasses don't need to override __new__ + args: &Bound<'p, PyTuple>, + kwargs: Option<&Bound<'p, PyDict>>, + ) -> Self { + Self { + trace_id: trace_id.and_then(|v| v.extract::().ok()), + span_id: span_id.and_then(|v| v.extract::().ok()), + meta: Some( + meta.map(|d| d.clone().unbind()) + .unwrap_or_else(|| PyDict::new(py).unbind()), + ), + metrics: Some( + metrics + .map(|d| d.clone().unbind()) + .unwrap_or_else(|| PyDict::new(py).unbind()), + ), + baggage: Some( + baggage + .map(|d| d.clone().unbind()) + .unwrap_or_else(|| PyDict::new(py).unbind()), + ), + span_links: Some( + span_links + .map(|l| l.clone().unbind()) + .unwrap_or_else(|| PyList::empty(py).unbind()), + ), + is_remote, + reactivate: false, + } + } + + // --- trace_id --- + #[getter] + #[inline(always)] + fn get_trace_id(&self) -> Option { + self.trace_id + } + + #[setter] + #[inline(always)] + fn set_trace_id(&mut self, value: Option<&Bound<'_, PyAny>>) { + match value { + None => self.trace_id = None, + // Silently ignore invalid types (keep existing value), matching SpanData's setters. + Some(v) => { + if let Ok(id) = v.extract::() { + self.trace_id = Some(id); + } + } + } + } + + // --- span_id --- + #[getter] + #[inline(always)] + fn get_span_id(&self) -> Option { + self.span_id + } + + #[setter] + #[inline(always)] + fn set_span_id(&mut self, value: Option<&Bound<'_, PyAny>>) { + match value { + None => self.span_id = None, + Some(v) => { + if let Ok(id) = v.extract::() { + self.span_id = Some(id); + } + } + } + } + + // --- _meta --- + #[getter(_meta)] + #[inline(always)] + fn get_meta<'py>(&mut self, py: Python<'py>) -> Bound<'py, PyDict> { + self.meta + .get_or_insert_with(|| PyDict::new(py).unbind()) + .bind(py) + .clone() + } + + #[setter(_meta)] + #[inline(always)] + fn set_meta(&mut self, value: &Bound<'_, PyDict>) { + self.meta = Some(value.clone().unbind()); + } + + // --- _metrics --- + #[getter(_metrics)] + #[inline(always)] + fn get_metrics<'py>(&mut self, py: Python<'py>) -> Bound<'py, PyDict> { + self.metrics + .get_or_insert_with(|| PyDict::new(py).unbind()) + .bind(py) + .clone() + } + + #[setter(_metrics)] + #[inline(always)] + fn set_metrics(&mut self, value: &Bound<'_, PyDict>) { + self.metrics = Some(value.clone().unbind()); + } + + // --- _baggage --- + #[getter(_baggage)] + #[inline(always)] + fn get_baggage<'py>(&mut self, py: Python<'py>) -> Bound<'py, PyDict> { + self.baggage + .get_or_insert_with(|| PyDict::new(py).unbind()) + .bind(py) + .clone() + } + + #[setter(_baggage)] + #[inline(always)] + fn set_baggage(&mut self, value: &Bound<'_, PyDict>) { + self.baggage = Some(value.clone().unbind()); + } + + // --- _span_links --- + #[getter(_span_links)] + #[inline(always)] + fn get_span_links<'py>(&mut self, py: Python<'py>) -> Bound<'py, PyList> { + self.span_links + .get_or_insert_with(|| PyList::empty(py).unbind()) + .bind(py) + .clone() + } + + #[setter(_span_links)] + #[inline(always)] + fn set_span_links(&mut self, value: &Bound<'_, PyList>) { + self.span_links = Some(value.clone().unbind()); + } + + // --- Cyclic GC support --- + // + // `_meta`/`_metrics`/`_baggage`/`_span_links` are Python containers whose + // contents are arbitrary user data (e.g. a baggage value can reference + // something that references this Context). Without `__traverse__`/`__clear__` + // such a cycle is invisible to CPython's cyclic GC and leaks forever -- see + // the identical rationale on `SpanData`, which hit exactly this class of bug + // for `meta_struct`. + fn __traverse__(&self, visit: pyo3::PyVisit<'_>) -> Result<(), pyo3::PyTraverseError> { + if let Some(d) = &self.meta { + visit.call(d)?; + } + if let Some(d) = &self.metrics { + visit.call(d)?; + } + if let Some(d) = &self.baggage { + visit.call(d)?; + } + if let Some(l) = &self.span_links { + visit.call(l)?; + } + Ok(()) + } + + fn __clear__(&mut self) { + // Reset to Default to drop every owned Python reference so CPython can + // break cycles. See `SpanData::__clear__` for the identical rationale. + *self = Self::default(); + } +} diff --git a/src/native/context/mod.rs b/src/native/context/mod.rs new file mode 100644 index 00000000000..1ec5f70ce85 --- /dev/null +++ b/src/native/context/mod.rs @@ -0,0 +1,10 @@ +use pyo3::types::PyModuleMethods as _; + +mod context_data; + +pub use context_data::ContextData; + +pub fn register_context(m: &pyo3::Bound<'_, pyo3::types::PyModule>) -> pyo3::PyResult<()> { + m.add_class::()?; + Ok(()) +} diff --git a/src/native/context_provider.rs b/src/native/context_provider.rs index 006ff059e20..c852f1fb0f9 100644 --- a/src/native/context_provider.rs +++ b/src/native/context_provider.rs @@ -3,13 +3,9 @@ //! `BaseContextProvider`/`DefaultContextProvider` are the hot path for every //! `tracer.context_provider.active()`/`.activate()` call. `_update_active` in //! particular runs on every `active()` call while a span is active, so it -//! downcasts straight to `SpanData` and reads `duration`, `_parent`, and -//! `_parent_context` as native fields instead of round-tripping through Python -//! attribute lookups. -//! -//! `_reactivate` is read from the parent `Context` via `getattr` because -//! `Context` is still a pure-Python class (porting it to native is a separate -//! effort — it carries substantial logic and many Python imports). +//! downcasts straight to `SpanData` and reads `duration`, `_parent`, +//! `_parent_context`, and (via `ContextData`) `_reactivate` as native fields +//! instead of round-tripping through Python attribute lookups. use std::sync::OnceLock; @@ -272,8 +268,8 @@ impl DefaultContextProvider { }; if parent.is_none() { if let Some(parent_context) = parent_context { - // `_reactivate` lives on the pure-Python Context -- still a getattr. - if parent_context.getattr("_reactivate")?.is_truthy()? { + if parent_context.borrow().reactivate { + let parent_context = parent_context.into_any(); call_activate(slf, py, Some(parent_context.clone()))?; return Ok(Some(parent_context.unbind())); } diff --git a/src/native/lib.rs b/src/native/lib.rs index 9b16a4d07ae..babaa3bc44b 100644 --- a/src/native/lib.rs +++ b/src/native/lib.rs @@ -3,6 +3,7 @@ mod crashtracker; #[cfg(feature = "profiling")] pub use datadog_profiling_ffi::*; mod config; +mod context; mod context_provider; #[cfg(all(Py_3_14, not(any(PyPy, GraalPy))))] mod context_watcher; @@ -85,6 +86,7 @@ fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> { span::register_native_span(m)?; event_hub::register_event_hub(m)?; contextvar::register_contextvar(m)?; + context::register_context(m)?; context_provider::register_context_provider(m)?; rand::register_rand(m)?; m.add_function(wrap_pyfunction!(ddtrace_utils::flatten_key_value, m)?)?; diff --git a/src/native/span/span_data.rs b/src/native/span/span_data.rs index bba1b219481..388f007ce28 100644 --- a/src/native/span/span_data.rs +++ b/src/native/span/span_data.rs @@ -58,9 +58,11 @@ pub struct SpanData { /// Set from Python during span creation; read natively by the context /// provider when walking the ancestor chain in `_update_active`. pub _parent: Option>, - /// The parent `Context` this span was created under, or `None`. - /// Held as `Py` because `Context` is still a pure-Python class. - pub _parent_context: Option>, + /// The parent `Context` this span was created under, or `None`. `Context` + /// (pure-Python) subclasses native `ContextData`, so this is typed against + /// the base class; PyO3 extracts a `Context` instance into `Py` + /// via ordinary covariant pyclass conversion. + pub _parent_context: Option>, } impl SpanData { @@ -461,7 +463,10 @@ impl SpanData { // _parent_context property — the parent Context, or None. #[getter(_parent_context)] #[inline(always)] - fn get_parent_context<'py>(&self, py: Python<'py>) -> Option> { + fn get_parent_context<'py>( + &self, + py: Python<'py>, + ) -> Option> { self._parent_context.as_ref().map(|c| c.bind(py).clone()) } @@ -471,7 +476,8 @@ impl SpanData { self._parent_context = if value.is_none() { None } else { - Some(value.clone().unbind()) + // Silently ignore non-Context values, matching other setters' defensive style. + value.extract::>().ok() }; } diff --git a/tests/tracer/test_context.py b/tests/tracer/test_context.py index c4b6f2f52b9..bbbd264a483 100644 --- a/tests/tracer/test_context.py +++ b/tests/tracer/test_context.py @@ -419,3 +419,53 @@ def test_is_remote(): # is_remote should be set to False on root spans. root = Span("root") assert root.context._is_remote is False + + +# ============================================================================= +# Cyclic GC support +# ============================================================================= +# +# Mirrors the regression tests in tests/tracer/test_span_data.py for SpanData: +# native pyclasses that hold `Py` / `Py` fields without +# implementing `__traverse__` / `__clear__` are invisible to CPython's cyclic +# GC, so a cycle passing through one of those fields (e.g. `_baggage`) leaks +# forever. These tests build the canonical cycle (context -> dict -> list -> +# context) and assert that `gc.collect()` reclaims it. + + +def _count_objects_of_type(typename): + import gc + + return sum(1 for o in gc.get_objects() if type(o).__name__ == typename) + + +def test_context_is_gc_tracked(): + import gc + + assert gc.is_tracked(Context()) + + +def test_context_baggage_cycle_is_collectable(): + """Cycles formed via `_baggage` must be reclaimed by `gc.collect()`.""" + import gc + + initial = _count_objects_of_type("Context") + gc_was_enabled = gc.isenabled() + gc.disable() + try: + N = 200 + for _ in range(N): + ctx = Context() + cycle_list = [] + ctx.set_baggage_item("self_ref", cycle_list) + cycle_list.append(ctx) + del ctx + del cycle_list + # All N cycles still alive (gc disabled, refcount can't break the cycle). + assert _count_objects_of_type("Context") - initial == N + freed = gc.collect() + assert freed > 0, "gc.collect freed nothing — Context is not GC-tracked" + assert _count_objects_of_type("Context") == initial + finally: + if gc_was_enabled: + gc.enable()