feat(trace-registry): optimize access (#93)

# Motivations

The trace registry is currently a shared data structure protected by a
single lock. This means that 2 threads accessing different traces will
create contention.

The goal of this PR is to optimize access to the trace registry so
creating and finishing spans can be as fast as possible and not block
the user critical path.

# Changes

The optimisations work in two ways:
* For single threaded access: The default std Hashmap has been replaced
by hasbrown. hasbrown uses foldhash, a non cryptographic hash function
which is ~10x faster than the std hash function
* For concurrent access: The trace registry now shards the lock based on
the trace id. This means that we divide the trace registry in 64 locks.
So two different traces now only have 1/64 chances of colliding.

Through testing I have figured out that this should allow us to scale to
up to 32 threads doing frequent concurrent reads on the trace registry

Author: paullegranddc

Cargo.lock

@@ -38,13 +38,19 @@ opentelemetry_sdk = { version = "0.30.0", features = [
 opentelemetry = { version = "0.30.0", features = [
     "trace",
 ], default-features = false }
-opentelemetry-semantic-conventions = { version = "0.30.0", features = ["semconv_experimental"] }
+opentelemetry-semantic-conventions = { version = "0.30.0", features = [
+    "semconv_experimental",
+] }
 tokio = { version = "1.44.1" }
 rand = { version = "0.8", features = ["small_rng"] }
 
 serde = { version = "1.0.194" }
 serde_json = { version = "1.0.140" }
 
+hashbrown = { version = "0.15.5" }
+foldhash = { version = "0.1.5" }
+
+
 [profile.dev]
 debug = 2 # full debug info
 

Cargo.toml

@@ -12,7 +12,9 @@ description.workspace = true
 # Workspace dependencies
 dd-trace = { path = "../dd-trace" }
 dd-trace-sampling = { path = "../dd-trace-sampling" }
-dd-trace-propagation = { path = "../dd-trace-propagation", features = ["opentelemetry"] }
+dd-trace-propagation = { path = "../dd-trace-propagation", features = [
+    "opentelemetry",
+] }
 datadog-opentelemetry-mappings = { path = "../datadog-opentelemetry-mappings" }
 
 # External dependencies
@@ -22,6 +24,8 @@ opentelemetry = { workspace = true }
 opentelemetry-semantic-conventions = { workspace = true }
 rand = { workspace = true }
 serde_json = { workspace = true }
+foldhash = { workspace = true }
+hashbrown = { workspace = true }
 
 # Libdatadog dependencies
 data-pipeline = { workspace = true }
@@ -44,5 +48,8 @@ tracing = { version = "0.1.4" }
 tracing-opentelemetry = { version = "0.31" }
 tracing-subscriber = { version = "0.3" }
 
+# Benchmarking
+criterion = "0.5.1"
+
 [features]
 test-utils = ["dd-trace/test-utils", "data-pipeline/test-utils"]

datadog-opentelemetry/Cargo.toml

@@ -1,8 +1,10 @@
 // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
+use hashbrown::{hash_map, HashMap as BHashMap};
 use std::{
-    collections::{hash_map, HashMap},
+    collections::HashMap,
+    fmt::Debug,
     str::FromStr,
     sync::{Arc, RwLock},
 };
@@ -56,7 +58,7 @@ const EMPTY_PROPAGATION_DATA: TracePropagationData = TracePropagationData {
 
 #[derive(Debug)]
 struct InnerTraceRegistry {
-    registry: HashMap<[u8; 16], Trace>,
+    registry: BHashMap<[u8; 16], Trace>,
 }
 
 pub enum RegisterTracePropagationResult {
@@ -211,6 +213,12 @@ impl InnerTraceRegistry {
     }
 }
 
+const TRACE_REGISTRY_SHARDS: usize = 64;
+
+#[repr(align(128))]
+#[derive(Debug, Clone)]
+struct CachePadded<T>(T);
+
 #[derive(Clone, Debug)]
 /// A registry of traces that are currently running
 ///
@@ -220,23 +228,32 @@ impl InnerTraceRegistry {
 /// - The number of open spans in the trace
 /// - The sampling decision of the trace
 pub(crate) struct TraceRegistry {
-    // TODO: The lock should probably sharded based on the hash of the trace id
-    // so we reduce contention...
     // Example:
     // inner: Arc<[CacheAligned<RwLock<InnerTraceRegistry>>; N]>;
     // to access a trace we do inner[hash(trace_id) % N].read()
-    inner: Arc<RwLock<InnerTraceRegistry>>,
+    inner: Arc<[CachePadded<RwLock<InnerTraceRegistry>>; TRACE_REGISTRY_SHARDS]>,
+    hasher: foldhash::fast::RandomState,
 }
 
 impl TraceRegistry {
     pub fn new() -> Self {
         Self {
-            inner: Arc::new(RwLock::new(InnerTraceRegistry {
-                registry: HashMap::new(),
+            inner: Arc::new(std::array::from_fn(|_| {
+                CachePadded(RwLock::new(InnerTraceRegistry {
+                    registry: BHashMap::new(),
+                }))
             })),
+            hasher: foldhash::fast::RandomState::default(),
         }
     }
 
+    fn get_shard(&self, trace_id: [u8; 16]) -> &RwLock<InnerTraceRegistry> {
+        use std::hash::BuildHasher;
+        let hash = self.hasher.hash_one(u128::from_ne_bytes(trace_id));
+        let shard = hash as usize % TRACE_REGISTRY_SHARDS;
+        &self.inner[shard].0
+    }
+
     /// Register the trace propagation data for a given trace ID
     /// This increases the open span count for the trace by 1, but does not set the root span ID.
     /// You will then need to call `register_local_root_span` to set the root span ID
@@ -249,7 +266,7 @@ impl TraceRegistry {
         propagation_data: TracePropagationData,
     ) -> RegisterTracePropagationResult {
         let mut inner = self
-            .inner
+            .get_shard(trace_id)
             .write()
             .expect("Failed to acquire lock on trace registry");
         inner.register_local_root_trace_propagation_data(trace_id, propagation_data)
@@ -260,7 +277,7 @@ impl TraceRegistry {
     /// If the trace is already registered, it will ignore the new root span ID and log a warning.
     pub fn register_local_root_span(&self, trace_id: [u8; 16], root_span_id: [u8; 8]) {
         let mut inner = self
-            .inner
+            .get_shard(trace_id)
             .write()
             .expect("Failed to acquire lock on trace registry");
         inner.register_local_root_span(trace_id, root_span_id);
@@ -274,7 +291,7 @@ impl TraceRegistry {
         propagation_data: TracePropagationData,
     ) {
         let mut inner = self
-            .inner
+            .get_shard(trace_id)
             .write()
             .expect("Failed to acquire lock on trace registry");
         inner.register_span(trace_id, span_id, propagation_data);
@@ -285,15 +302,15 @@ impl TraceRegistry {
     /// flush
     fn finish_span(&self, trace_id: [u8; 16], span_data: SpanData) -> Option<Trace> {
         let mut inner = self
-            .inner
+            .get_shard(trace_id)
             .write()
             .expect("Failed to acquire lock on trace registry");
         inner.finish_span(trace_id, span_data)
     }
 
     pub fn get_trace_propagation_data(&self, trace_id: [u8; 16]) -> TracePropagationData {
         let inner = self
-            .inner
+            .get_shard(trace_id)
             .read()
             .expect("Failed to acquire lock on trace registry");
 
@@ -532,13 +549,22 @@ impl opentelemetry_sdk::trace::SpanProcessor for DatadogSpanProcessor {
 
 #[cfg(test)]
 mod tests {
-    use std::sync::{Arc, RwLock};
-
-    use dd_trace::Config;
+    use std::{
+        collections::HashMap,
+        hint::black_box,
+        sync::{Arc, RwLock},
+        thread,
+        time::Duration,
+    };
+
+    use dd_trace::{
+        sampling::{mechanism, priority, SamplingDecision},
+        Config,
+    };
     use opentelemetry::{Key, KeyValue, Value};
     use opentelemetry_sdk::{trace::SpanProcessor, Resource};
 
-    use crate::span_processor::{DatadogSpanProcessor, TraceRegistry};
+    use crate::span_processor::{DatadogSpanProcessor, TracePropagationData, TraceRegistry};
 
     #[test]
     fn test_set_resource_from_empty_dd_config() {
@@ -683,4 +709,103 @@ mod tests {
             Some(Value::String("otel-service".into()))
         );
     }
+
+    fn bench_trace_registry(c: &mut criterion::Criterion) {
+        const ITERATIONS: u32 = 10000;
+        const NUM_TRACES: usize = ITERATIONS as usize / 20;
+        let mut group = c.benchmark_group("trace_registry_concurrent_access_threads");
+        group
+            .warm_up_time(Duration::from_millis(100))
+            .measurement_time(Duration::from_millis(1000));
+
+        for concurrency in [1, 2, 4, 8, 16, 32] {
+            group
+                .throughput(criterion::Throughput::Elements(
+                    ITERATIONS as u64 * concurrency,
+                ))
+                .bench_function(
+                    criterion::BenchmarkId::from_parameter(concurrency),
+                    move |g| {
+                        let trace_ids: Vec<_> = (0..concurrency)
+                            .map(|thread| {
+                                std::array::from_fn::<_, NUM_TRACES, _>(|i| {
+                                    ((thread << 16 | i as u64) as u128).to_be_bytes()
+                                })
+                            })
+                            .collect();
+                        g.iter_batched_ref(
+                            {
+                                let trace_ids = trace_ids.clone();
+                                move || {
+                                    let tr: TraceRegistry = TraceRegistry::new();
+                                    for trace_id in trace_ids.iter().flatten() {
+                                        tr.register_local_root_trace_propagation_data(
+                                            *trace_id,
+                                            TracePropagationData {
+                                                sampling_decision: SamplingDecision {
+                                                    priority: Some(priority::AUTO_KEEP),
+                                                    mechanism: Some(mechanism::DEFAULT),
+                                                },
+                                                origin: Some("rum".to_string()),
+                                                tags: Some(HashMap::from_iter([(
+                                                    "dd.p.tid".to_string(),
+                                                    "foobar".to_string(),
+                                                )])),
+                                            },
+                                        );
+                                    }
+                                    tr
+                                }
+                            },
+                            move |tr| {
+                                let tr = &*tr;
+                                let trace_ids = &trace_ids;
+                                thread::scope(move |s| {
+                                    for trace_id in trace_ids {
+                                        s.spawn(move || {
+                                            for _ in 0..(ITERATIONS as usize / NUM_TRACES) {
+                                                for trace_id in trace_id {
+                                                    black_box(tr.get_trace_propagation_data(
+                                                        black_box(*trace_id),
+                                                    ));
+                                                }
+                                            }
+                                        });
+                                    }
+                                })
+                            },
+                            criterion::BatchSize::LargeInput,
+                        );
+                    },
+                );
+        }
+    }
+
+    #[test]
+    fn bench() {
+        // Run with
+        // `cargo test --profile bench -- --nocapture bench -- <benchmark_filter>
+        // Collect cli arguments
+
+        // Interpret sequence of args `[ "...bench", "--", "[filter]" ]` as a trigger and extract
+        // `filter`
+        let filter = std::env::args()
+            .collect::<Vec<_>>()
+            .windows(3)
+            .filter(|p| p.len() >= 2 && p[0].ends_with("bench") && p[1] == "--")
+            .map(|s| s.get(2).unwrap_or(&"".to_string()).clone())
+            .next();
+
+        let filter = match filter {
+            None => return,
+            Some(f) => f,
+        };
+
+        let mut criterion = criterion::Criterion::default()
+            .with_output_color(true)
+            .with_filter(&filter);
+        bench_trace_registry(&mut criterion);
+
+        criterion.final_summary();
+    }
 }