context_provider.rs

//! Native port of `ddtrace._trace.provider`.
//!
//! `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).

use std::sync::OnceLock;

use pyo3::{
    exceptions::{PyNotImplementedError, PyTypeError},
    types::{PyAnyMethods as _, PyDict, PyModule, PyModuleMethods as _, PyTuple, PyType},
    Bound, Py, PyAny, PyResult, Python,
};

use crate::contextvar::{contextvar_get, contextvar_new, safe_contextvar_set};
use crate::event_hub;
use crate::span::SpanData;

// The active context/span is tracked in a single process-wide `ContextVar`,
// shared by every `DefaultContextProvider` instance -- this mirrors the
// module-level `_DD_CONTEXTVAR` singleton in the old Python implementation.
static CONTEXTVAR: OnceLock<Py<PyAny>> = OnceLock::new();

/// Borrows the shared contextvar without touching its refcount -- `Py::bind`
/// is a pointer cast, not an INCREF, so the common (post-init) path here is
/// free. Only the one-time init path allocates.
fn contextvar(py: Python<'_>) -> PyResult<&Bound<'_, PyAny>> {
    if let Some(v) = CONTEXTVAR.get() {
        return Ok(v.bind(py));
    }
    let cv = contextvar_new(py, "datadog_contextvar", &py.None().into_bound(py))?;
    // Lost races just drop their extra ContextVar; the getter always reads
    // back through the OnceLock so every caller ends up on the same instance.
    let _ = CONTEXTVAR.set(cv.unbind());
    Ok(CONTEXTVAR.get().unwrap().bind(py))
}

/// Set the active contextvar, routing through the crash-safe setter on
/// CPython < 3.12 (see `contextvar::safe_contextvar_set`).
fn activate_contextvar(py: Python<'_>, ctx: Option<&Bound<'_, PyAny>>) -> PyResult<()> {
    let cv = contextvar(py)?;
    match ctx {
        Some(v) => safe_contextvar_set(py, cv, v),
        None => safe_contextvar_set(py, cv, &py.None().into_bound(py)),
    }
}

const ACTIVATE_EVENT: &str = "ddtrace.context_provider.activate";

/// `core.dispatch("ddtrace.context_provider.activate", (self, ctx))` -- called
/// directly instead of importing `ddtrace.internal.core`, since that module's
/// `dispatch` is itself a re-export of `event_hub::dispatch`.
///
/// `activate()` runs on every span start/finish, but the only listener
/// (profiling's stack sampler) is rarely registered, so check first to skip
/// the tuple allocation on the common no-listener path.
fn dispatch_activate(
    py: Python<'_>,
    slf: &Bound<'_, PyAny>,
    ctx: Option<&Bound<'_, PyAny>>,
) -> PyResult<()> {
    if !event_hub::has_listeners(ACTIVATE_EVENT) {
        return Ok(());
    }
    let ctx_value = match ctx {
        Some(v) => v.clone(),
        None => py.None().into_bound(py),
    };
    let args = PyTuple::new(py, [slf.clone(), ctx_value])?;
    event_hub::dispatch(py, ACTIVATE_EVENT, Some(args.into_any().unbind()), false)
}

enum Resolved<'py> {
    Unchanged,
    /// May hold Python `None`, when the whole ancestor chain is finished.
    Ancestor(Bound<'py, PyAny>),
    ReactivatableContext(Bound<'py, PyAny>),
}

/// Where the active trace *should* be, walking past finished ancestors of `span`.
///
/// Shared by `_update_active`, which applies the result, and `_peek_active`, which
/// only reports it, so the two cannot drift apart. Borrows `span` so that
/// `Resolved::Unchanged` costs the caller nothing -- `_update_active` runs on every
/// span finish, where an extra incref shows up.
#[inline]
fn resolve_active<'py>(py: Python<'py>, span: &Bound<'py, PyAny>) -> PyResult<Resolved<'py>> {
    let mut current = span.clone();
    loop {
        // PERF: read `duration`, `_parent`, and `_parent_context` straight off the
        // native SpanData fields in one borrow -- avoids three Python attribute
        // lookups per ancestor hop.
        let (parent, parent_context) = {
            let Ok(sd) = current.cast::<SpanData>() else {
                break; // not a Span (e.g. None) -- stop walking parents
            };
            let sd = sd.borrow();
            if sd.duration.is_none() {
                break; // unfinished span -- stop walking parents
            }
            (
                sd._parent.as_ref().map(|p| p.bind(py).clone()),
                sd._parent_context.as_ref().map(|c| c.bind(py).clone()),
            )
        };
        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()? {
                    return Ok(Resolved::ReactivatableContext(parent_context));
                }
            }
        }
        // Advance to the parent; `None` ends the walk on the next iteration.
        current = parent.unwrap_or_else(|| py.None().into_bound(py));
    }
    if current.is(span) {
        Ok(Resolved::Unchanged)
    } else {
        Ok(Resolved::Ancestor(current))
    }
}

/// `Some(v)` unless `v` is Python `None`, matching the `Optional[...]` return
/// convention used throughout this module. Consumes `v` to avoid an incref.
#[inline]
fn none_or_unbind(v: Bound<'_, PyAny>) -> Option<Py<PyAny>> {
    if v.is_none() {
        None
    } else {
        Some(v.unbind())
    }
}

/// Borrowing sibling of [`none_or_unbind`] for when the caller still needs `v`:
/// `Some(clone)` unless `v` is Python `None`.
#[inline]
fn none_or_clone<'py>(v: &Bound<'py, PyAny>) -> Option<Bound<'py, PyAny>> {
    if v.is_none() {
        None
    } else {
        Some(v.clone())
    }
}

/// Invoke `self.activate(ctx)` from within `active`/`_update_active`.
///
/// Fast path: when `slf` is *exactly* a `DefaultContextProvider` (the common
/// case), call the native method directly and skip the Python method-resolution
/// (`_PyType_Lookup` + `GenericGetAttr`) that `call_method1` incurs. A Python
/// subclass (CIContextProvider, LLMObsContextProvider, ...) may override
/// `activate`, so route through Python dispatch for those.
#[inline]
fn call_activate<'py>(
    slf: &Bound<'py, DefaultContextProvider>,
    py: Python<'py>,
    ctx: Option<Bound<'py, PyAny>>,
) -> PyResult<()> {
    if slf.is_exact_instance_of::<DefaultContextProvider>() {
        DefaultContextProvider::activate(slf, py, ctx)
    } else {
        slf.call_method1("activate", (ctx,)).map(|_| ())
    }
}

/// A ``ContextProvider`` is an interface that provides the blueprint
/// for a callable class, capable to retrieve the current active
/// ``Context`` instance. Context providers must inherit this class
/// and implement:
/// * the ``active`` method, that returns the current active ``Context``
/// * the ``activate`` method, that sets the current active ``Context``
// `dict` gives instances a `__dict__`, so external code (mocks, monkeypatches,
// instrumentation) can still set/replace attributes like `activate` on a live
// provider -- native pyclasses have no instance dict by default, which would
// otherwise make method attributes read-only.
#[pyo3::pyclass(subclass, dict, module = "ddtrace.internal._native")]
pub struct BaseContextProvider;

#[pyo3::pymethods]
impl BaseContextProvider {
    /// Takes `cls` so only a *direct* instantiation of `BaseContextProvider`
    /// is rejected -- a concrete subclass (e.g. one passed to
    /// `Tracer.configure(context_provider=...)`) must still be constructible,
    /// mirroring the old `abc.ABCMeta` behavior where only the base type
    /// itself was blocked.
    #[new]
    #[classmethod]
    fn new(cls: &Bound<'_, PyType>) -> PyResult<Self> {
        if cls.is(cls.py().get_type::<Self>()) {
            Err(PyTypeError::new_err(
                "Can't instantiate abstract class BaseContextProvider without an implementation \
                 for abstract methods '_has_active_context', 'active'",
            ))
        } else {
            Ok(Self)
        }
    }

    fn _has_active_context(&self) -> PyResult<bool> {
        Err(PyNotImplementedError::new_err(()))
    }

    #[pyo3(signature = (ctx))]
    fn activate(
        slf: &Bound<'_, Self>,
        py: Python<'_>,
        ctx: Option<Bound<'_, PyAny>>,
    ) -> PyResult<()> {
        dispatch_activate(py, slf.as_any(), ctx.as_ref())
    }

    fn active(&self) -> PyResult<Option<Py<PyAny>>> {
        Err(PyNotImplementedError::new_err(()))
    }

    /// See `DefaultContextProvider::_peek_active`. A provider with its own storage
    /// keeps today's behavior until it opts into a read-only path.
    fn _peek_active(slf: &Bound<'_, Self>) -> PyResult<Option<Py<PyAny>>> {
        slf.call_method0("active").map(none_or_unbind)
    }

    /// Method available for backward-compatibility. It proxies the call to
    /// ``self.active()`` and must not do anything more.
    #[pyo3(signature = (*_args, **_kwargs))]
    fn __call__(
        slf: &Bound<'_, Self>,
        _args: &Bound<'_, PyTuple>,
        _kwargs: Option<&Bound<'_, PyDict>>,
    ) -> PyResult<Py<PyAny>> {
        // Dispatched via `call_method0` (not a direct Rust call) so subclass
        // overrides of `active` -- e.g. CIContextProvider, TracerStackContext --
        // are honored.
        slf.call_method0("active").map(Bound::unbind)
    }
}

/// Context provider that retrieves contexts from a context variable.
///
/// It is suitable for synchronous programming and for asynchronous executors
/// that support contextvars.
#[pyo3::pyclass(extends = BaseContextProvider, subclass, module = "ddtrace.internal._native")]
pub struct DefaultContextProvider;

#[pyo3::pymethods]
impl DefaultContextProvider {
    #[new]
    fn new() -> (Self, BaseContextProvider) {
        (DefaultContextProvider, BaseContextProvider)
    }

    /// Returns whether there is an active context in the current execution.
    fn _has_active_context(&self, py: Python<'_>) -> PyResult<bool> {
        let item = contextvar_get(py, contextvar(py)?)?;
        Ok(!item.is_none())
    }

    /// Makes the given context active in the current execution.
    #[pyo3(signature = (ctx))]
    fn activate<'py>(
        slf: &Bound<'py, Self>,
        py: Python<'py>,
        ctx: Option<Bound<'py, PyAny>>,
    ) -> PyResult<()> {
        activate_contextvar(py, ctx.as_ref())?;
        // `super(DefaultContextProvider, self).activate(ctx)` -- calls
        // BaseContextProvider's dispatch directly (not via `call_method`,
        // which would just re-resolve back to this same override). `as_super()`
        // upcasts by reference, so `base` still points at the same `self`.
        let base = slf.as_super();
        BaseContextProvider::activate(base, py, ctx)
    }

    /// Returns the active span or context for the current execution.
    fn active<'py>(slf: &Bound<'py, Self>, py: Python<'py>) -> PyResult<Option<Py<PyAny>>> {
        let item = contextvar_get(py, contextvar(py)?)?;
        if item.is_none() {
            return Ok(None);
        }
        match item.cast::<SpanData>() {
            Ok(span) => {
                // Fast path: an exact `DefaultContextProvider` calls the native
                // `_update_active` directly. A subclass (e.g. LLMObsContextProvider)
                // may override it, so route those through Python dispatch.
                if slf.is_exact_instance_of::<DefaultContextProvider>() {
                    Self::_update_active(slf, py, span.clone().into_any())
                } else {
                    slf.call_method1("_update_active", (span.clone(),))
                        .map(none_or_unbind)
                }
            }
            Err(_) => Ok(Some(item.unbind())),
        }
    }

    /// Updates the active trace in an executor.
    ///
    /// When a span finishes, the active span becomes its parent.
    /// If no parent exists and the context is reactivatable, that context is restored.
    fn _update_active<'py>(
        slf: &Bound<'py, Self>,
        py: Python<'py>,
        span: Bound<'py, PyAny>,
    ) -> PyResult<Option<Py<PyAny>>> {
        match resolve_active(py, &span)? {
            Resolved::Unchanged => Ok(none_or_unbind(span)),
            Resolved::ReactivatableContext(parent_context) => {
                call_activate(slf, py, Some(parent_context.clone()))?;
                Ok(Some(parent_context.unbind()))
            }
            Resolved::Ancestor(current) => {
                call_activate(slf, py, none_or_clone(&current))?;
                Ok(none_or_unbind(current))
            }
        }
    }

    /// `active` without its repair: no `activate`, so no contextvar write and no
    /// `ddtrace.context_provider.activate` dispatch.
    ///
    /// For observers that must not perturb the active trace. A CPython
    /// context-switch watcher is one: `activate` there would write to whichever
    /// context the switch just made current, and re-enter the event hub mid-switch.
    ///
    /// Reads this contextvar and resolves the way `_update_active` does, so **a subclass
    /// that overrides `active` or `_update_active` must override this too** -- otherwise
    /// it reports from storage it does not use. Subclasses that only add behaviour
    /// elsewhere inherit it safely.
    fn _peek_active<'py>(_slf: &Bound<'py, Self>, py: Python<'py>) -> PyResult<Option<Py<PyAny>>> {
        let item = contextvar_get(py, contextvar(py)?)?;
        if item.is_none() {
            return Ok(None);
        }
        if item.cast::<SpanData>().is_err() {
            return Ok(Some(item.unbind()));
        }
        match resolve_active(py, &item)? {
            Resolved::Unchanged => Ok(none_or_unbind(item)),
            Resolved::ReactivatableContext(parent_context) => Ok(Some(parent_context.unbind())),
            Resolved::Ancestor(current) => Ok(none_or_unbind(current)),
        }
    }
}

pub fn register_context_provider(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<BaseContextProvider>()?;
    m.add_class::<DefaultContextProvider>()?;
    // Exposed so `ddtrace._trace.provider` (and downstream code like the
    // gevent monkeypatch) can keep importing `_DD_CONTEXTVAR` directly.
    let cv = contextvar(m.py())?;
    m.add("DD_CONTEXTVAR", cv)?;
    Ok(())
}