add tracing instrumentation to the pyo3 layer through ContextVars (#12026)

### What

As part of benchmark testing, we'd like to be able to have a single e2e
trace for each benchmark test. In order to do so, we need to pass trace
context across pyo3 boundary and it seems (per Gemini) that the
recommended way to do this is through ``ContextVar``s which is what
standard tracing libraries in python do.

This PR leverages that and adds a context extraction logic as well as an
API (per @abey79 suggestion) to the rerun sdk so that we can easily
access this trace context var and set it on the python side.

For now I've focused on the query paths and we can see both query and
fetching phase under a single trace thanks to changes in the dataframe
table provider and trace propagation between query execution phases.
Follow up will add instrumentation to the register and index creation
paths.

Example usage:
```
        from rerun.catalog import _rerun_trace_context
        from opentelemetry import context
        from opentelemetry.propagate import get_global_textmap

        trace_ctx = _rerun_trace_context()

        with tracer.start_as_current_span("sample_index_values") as span:
            # Add attributes to help identify this benchmark
            span.set_attribute("benchmark.name", "sample_index_values")
            span.set_attribute("dataset.name", self.dataset_name)
            span.set_attribute("sample_count", 6)

            wrist = self.dataset.dataframe_query_view(
                index="log_tick",
                contents="/camera/wrist/embedding /thumbnail/camera/wrist",
            )

            sampled_times = [0, 100, 200, 500, 1000, 2000]

            current_ctx = context.get_current()
            carrier = {}
            get_global_textmap().inject(carrier, current_ctx)

            if carrier:
                token = trace_ctx.set(carrier)

            result = (
                (wrist.filter_index_values(sampled_times).fill_latest_at()).df().drop("rerun_partition_id").collect()
            )

            if carrier:
                trace_ctx.reset(token)

            span.set_attribute("result_rows", len(result))
``

Author: zehiko

Cargo.lock

@@ -8869,6 +8869,7 @@ dependencies = [
  "re_arrow_util",
  "re_dataframe",
  "re_log_types",
+ "re_perf_telemetry",
  "re_protos",
  "re_redap_client",
  "re_sorbet",
@@ -9209,6 +9210,7 @@ dependencies = [
  "opentelemetry_sdk",
  "parking_lot",
  "prometheus-client",
+ "pyo3",
  "serde",
  "serde_json",
  "tokio",

crates/store/re_datafusion/Cargo.toml

@@ -53,6 +53,9 @@ tokio-stream.workspace = true
 tonic.workspace = true
 tracing.workspace = true
 
+[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
+re_perf_telemetry.workspace = true
+
 [target.'cfg(target_arch = "wasm32")'.dependencies]
 futures.workspace = true
 wasm-bindgen-futures.workspace = true

crates/store/re_datafusion/src/dataframe_query_common.rs

@@ -49,6 +49,11 @@ pub struct DataframeQueryTableProvider {
     sort_index: Option<Index>,
     chunk_info_batches: Arc<Vec<RecordBatch>>,
     client: ConnectionClient,
+
+    /// passing trace headers between phases of execution pipeline helps keep
+    /// the entire operation under a single trace.
+    #[cfg(not(target_arch = "wasm32"))]
+    trace_headers: Option<crate::TraceHeaders>,
 }
 
 impl DataframeQueryTableProvider {
@@ -62,6 +67,7 @@ impl DataframeQueryTableProvider {
         dataset_id: EntryId,
         query_expression: &QueryExpression,
         partition_ids: &[impl AsRef<str> + Sync],
+        #[cfg(not(target_arch = "wasm32"))] trace_headers: Option<crate::TraceHeaders>,
     ) -> Result<Self, DataFusionError> {
         use futures::StreamExt as _;
 
@@ -164,6 +170,8 @@ impl DataframeQueryTableProvider {
             sort_index: query_expression.filtered_index,
             chunk_info_batches,
             client,
+            #[cfg(not(target_arch = "wasm32"))]
+            trace_headers,
         })
     }
 
@@ -259,6 +267,8 @@ impl TableProvider for DataframeQueryTableProvider {
             Arc::clone(&self.chunk_info_batches),
             query_expression,
             self.client.clone(),
+            #[cfg(not(target_arch = "wasm32"))]
+            self.trace_headers.clone(),
         )
         .map(Arc::new)
         .map(|exec| {

crates/store/re_datafusion/src/dataframe_query_provider.rs

@@ -45,6 +45,27 @@ use crate::dataframe_query_common::{
 /// the IO stream.
 const CPU_THREAD_IO_CHANNEL_SIZE: usize = 32;
 
+/// Helper to attach parent trace context if available.
+/// Returns a guard that must be kept alive for the duration of the traced scope.
+/// We can use this to ensure all phases of table provider's execution pipeline are
+/// parented by a single trace.
+#[cfg(not(target_arch = "wasm32"))]
+#[inline]
+fn attach_trace_context(
+    trace_headers: &Option<crate::TraceHeaders>,
+) -> Option<re_perf_telemetry::external::opentelemetry::ContextGuard> {
+    let headers = trace_headers.as_ref()?;
+    if !headers.traceparent.is_empty() {
+        let parent_ctx =
+            re_perf_telemetry::external::opentelemetry::global::get_text_map_propagator(|prop| {
+                prop.extract(headers)
+            });
+        Some(parent_ctx.attach())
+    } else {
+        None
+    }
+}
+
 #[derive(Debug)]
 pub(crate) struct PartitionStreamExec {
     props: PlanProperties,
@@ -61,6 +82,10 @@ pub(crate) struct PartitionStreamExec {
     target_partitions: usize,
     worker_runtime: Arc<CpuRuntime>,
     client: ConnectionClient,
+
+    /// passing trace headers between phases of execution pipeline helps keep
+    /// the entire operation under a single trace.
+    trace_headers: Option<crate::TraceHeaders>,
 }
 
 type ChunksWithPartition = Vec<(Chunk, Option<String>)>;
@@ -79,6 +104,10 @@ pub struct DataframePartitionStreamInner {
     /// so that our worker does not shut down unexpectedly.
     cpu_runtime: Arc<CpuRuntime>,
     cpu_join_handle: Option<JoinHandle<Result<(), DataFusionError>>>,
+
+    /// passing trace headers between phases of execution pipeline helps keep
+    /// the entire operation under a single trace.
+    trace_headers: Option<crate::TraceHeaders>,
 }
 
 /// This is a temporary fix to minimize the impact of leaking memory
@@ -96,13 +125,16 @@ pub struct DataframePartitionStream {
 impl Stream for DataframePartitionStream {
     type Item = Result<RecordBatch, DataFusionError>;
 
-    #[tracing::instrument(level = "info", skip_all)]
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let this_outer = self.get_mut();
         let Some(this) = this_outer.inner.as_mut() else {
             return Poll::Ready(None);
         };
 
+        #[cfg(not(target_arch = "wasm32"))]
+        let _trace_guard = attach_trace_context(&this.trace_headers);
+        let _span = tracing::info_span!("poll_next").entered();
+
         // If we have any errors on the worker thread, we want to ensure we pass them up
         // through the stream.
         if this
@@ -132,11 +164,16 @@ impl Stream for DataframePartitionStream {
                 return Poll::Ready(Some(exec_err!("No tx for chunks from CPU thread")));
             };
 
-            this.io_join_handle = Some(io_handle.spawn(chunk_stream_io_loop(
-                this.client.clone(),
-                this.chunk_infos.clone(),
-                chunk_tx,
-            )));
+            let client = this.client.clone();
+            let chunk_infos = this.chunk_infos.clone();
+            let current_span = tracing::Span::current();
+
+            this.io_join_handle = Some(
+                io_handle.spawn(
+                    async move { chunk_stream_io_loop(client, chunk_infos, chunk_tx).await }
+                        .instrument(current_span.clone()),
+                ),
+            );
         }
 
         let result = this
@@ -163,6 +200,7 @@ impl RecordBatchStream for DataframePartitionStream {
 
 impl PartitionStreamExec {
     #[tracing::instrument(level = "info", skip_all)]
+    #[expect(clippy::too_many_arguments)]
     pub fn try_new(
         table_schema: &SchemaRef,
         sort_index: Option<Index>,
@@ -171,6 +209,7 @@ impl PartitionStreamExec {
         chunk_info_batches: Arc<Vec<RecordBatch>>,
         mut query_expression: QueryExpression,
         client: ConnectionClient,
+        trace_headers: Option<crate::TraceHeaders>,
     ) -> datafusion::common::Result<Self> {
         let projected_schema = match projection {
             Some(p) => Arc::new(table_schema.project(p)?),
@@ -270,6 +309,7 @@ impl PartitionStreamExec {
             target_partitions: num_partitions,
             worker_runtime,
             client,
+            trace_headers,
         })
     }
 }
@@ -476,12 +516,15 @@ impl ExecutionPlan for PartitionStreamExec {
         }
     }
 
-    #[tracing::instrument(level = "info", skip_all)]
     fn execute(
         &self,
         partition: usize,
         _context: Arc<TaskContext>,
     ) -> datafusion::common::Result<SendableRecordBatchStream> {
+        #[cfg(not(target_arch = "wasm32"))]
+        let _trace_guard = attach_trace_context(&self.trace_headers);
+        let _span = tracing::info_span!("execute").entered();
+
         let (chunk_tx, chunk_rx) = tokio::sync::mpsc::channel(CPU_THREAD_IO_CHANNEL_SIZE);
 
         let random_state = ahash::RandomState::with_seeds(0, 0, 0, 0);
@@ -526,6 +569,7 @@ impl ExecutionPlan for PartitionStreamExec {
             io_join_handle: None,
             cpu_join_handle,
             cpu_runtime: Arc::clone(&self.worker_runtime),
+            trace_headers: self.trace_headers.clone(),
         };
         let stream = DataframePartitionStream {
             inner: Some(stream),
@@ -552,6 +596,7 @@ impl ExecutionPlan for PartitionStreamExec {
             target_partitions,
             worker_runtime: Arc::new(CpuRuntime::try_new(target_partitions)?),
             client: self.client.clone(),
+            trace_headers: self.trace_headers.clone(),
         };
 
         plan.props.partitioning = match plan.props.partitioning {

crates/store/re_datafusion/src/lib.rs

@@ -23,3 +23,6 @@ pub use dataset_manifest::DatasetManifestProvider;
 pub use partition_table::SegmentTableProvider;
 pub use search_provider::SearchResultsTableProvider;
 pub use table_entry_provider::TableEntryTableProvider;
+
+#[cfg(not(target_arch = "wasm32"))]
+pub(crate) type TraceHeaders = re_perf_telemetry::TraceHeaders;

crates/utils/re_perf_telemetry/Cargo.toml

@@ -42,6 +42,9 @@ tracy_enabled = []
 ## E.g. an async function that yields 50 times will be counted as 51 (the first call + 50 yields).
 tracy = ["dep:tracing-tracy"]
 
+## PyO3 integration for cross-boundary tracing
+pyo3 = ["dep:pyo3"]
+
 
 [dependencies]
 
@@ -78,6 +81,7 @@ tokio.workspace = true
 
 # External (optional)
 tracing-tracy = { workspace = true, optional = true }
+pyo3 = { workspace = true, optional = true }
 
 [lints]
 workspace = true

crates/utils/re_perf_telemetry/src/lib.rs

@@ -192,6 +192,82 @@ impl From<&TraceHeaders> for opentelemetry::Context {
 
 // ---
 
+/// The name of the `ContextVar` used for trace context propagation
+pub const TRACE_CONTEXT_VAR_NAME: &str = "TRACE_CONTEXT";
+
+#[cfg(feature = "pyo3")]
+/// Get the trace context `ContextVar` object.
+///
+/// This returns the same Python `ContextVar` instance every time, ensuring that
+/// values set on it can be read back later. It is up to the caller to ensure trace context
+/// is reset and cleared as needed.
+pub fn get_trace_context_var(py: pyo3::Python<'_>) -> pyo3::PyResult<pyo3::Bound<'_, pyo3::PyAny>> {
+    use pyo3::prelude::*;
+
+    static CONTEXT_VAR: parking_lot::Mutex<Option<pyo3::Py<pyo3::PyAny>>> =
+        parking_lot::Mutex::new(None);
+
+    let mut guard = CONTEXT_VAR.lock();
+
+    if let Some(var) = guard.as_ref() {
+        return Ok(var.bind(py).clone());
+    }
+
+    // Create the trace context ContextVar
+    let module = py.import("contextvars")?;
+    let contextvar_class = module.getattr("ContextVar")?;
+    let trace_ctx_var = contextvar_class.call1((TRACE_CONTEXT_VAR_NAME,))?;
+    let trace_ctx_unbound = trace_ctx_var.clone().unbind();
+
+    *guard = Some(trace_ctx_unbound);
+
+    Ok(trace_ctx_var)
+}
+
+#[cfg(feature = "pyo3")]
+/// Extract trace context from Python `ContextVar` for cross-boundary propagation.
+pub fn extract_trace_context_from_contextvar(py: pyo3::Python<'_>) -> TraceHeaders {
+    use pyo3::prelude::*;
+    use pyo3::types::PyDict;
+
+    fn try_extract(py: pyo3::Python<'_>) -> PyResult<TraceHeaders> {
+        let context_var = get_trace_context_var(py)?;
+
+        match context_var.call_method0("get") {
+            Ok(trace_data) => {
+                if let Ok(dict) = trace_data.downcast::<PyDict>() {
+                    let traceparent = dict
+                        .get_item(TraceHeaders::TRACEPARENT_KEY)?
+                        .and_then(|v| v.extract::<String>().ok())
+                        .unwrap_or_default();
+
+                    let tracestate = dict
+                        .get_item(TraceHeaders::TRACESTATE_KEY)?
+                        .and_then(|v| v.extract::<String>().ok());
+
+                    let headers = TraceHeaders {
+                        traceparent,
+                        tracestate,
+                    };
+
+                    tracing::debug!("Trace headers: {:?}", headers);
+                    Ok(headers)
+                } else {
+                    Ok(TraceHeaders::empty())
+                }
+            }
+            Err(_) => Ok(TraceHeaders::empty()),
+        }
+    }
+
+    try_extract(py).unwrap_or_else(|err| {
+        tracing::debug!("Failed to extract trace context: {}", err);
+        TraceHeaders::empty()
+    })
+}
+
+// ---
+
 // Extension to [`tracing_subscriber:EnvFilter`] that allows to
 // add a directive only if not already present in the base filter
 pub trait EnvFilterExt

rerun_py/Cargo.toml

@@ -15,7 +15,7 @@ name = "rerun_bindings" # name of the .so library that the Python module will im
 
 
 [features]
-default = ["extension-module"]
+default = ["extension-module", "perf_telemetry"]
 
 ## Extra features that aren't ready to be included in release builds yet.
 extra = []
@@ -38,7 +38,11 @@ nasm = ["re_video/nasm"]
 ## Enables integration with `re_perf_telemetry` (Tracy, Jaeger).
 ##
 ## This only works on native.
-perf_telemetry = ["dep:re_perf_telemetry", "re_redap_client/perf_telemetry"]
+perf_telemetry = [
+  "dep:re_perf_telemetry",
+  "re_redap_client/perf_telemetry",
+  "re_perf_telemetry/pyo3",
+]
 
 server = ["re_sdk/server", "dep:re_grpc_server"]
 
@@ -109,7 +113,7 @@ thiserror.workspace = true
 
 # Native dependencies:
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
-re_perf_telemetry = { workspace = true, features = ["tracy"], optional = true }
+re_perf_telemetry = { workspace = true, features = ["tracy", "pyo3"], optional = true }
 
 
 [build-dependencies]

rerun_py/rerun_bindings/rerun_bindings.pyi

@@ -1851,3 +1851,6 @@ class Credentials:
 
 def get_credentials() -> Credentials | None:
     """Returns the credentials for the current user."""
+
+def rerun_trace_context() -> Any:
+    """Get the trace context ContextVar for distributed tracing propagation."""

rerun_py/rerun_sdk/rerun/catalog/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+# Conditionally compiled function - always exists at runtime but mypy can't verify
 from rerun_bindings import (
     AlreadyExistsError as AlreadyExistsError,
     DataframeQueryView as DataframeQueryView,
@@ -14,6 +15,7 @@
     Task as Task,
     Tasks as Tasks,
     VectorDistanceMetric as VectorDistanceMetric,
+    rerun_trace_context as _rerun_trace_context,
 )
 from rerun_bindings.types import (
     IndexValuesLike as IndexValuesLike,