chore: root_span_id -> local_root_span_id: [u8; 8]

Author: yannham

libdd-profiling/src/otel_thread_ctx.rs

@@ -107,8 +107,8 @@ pub mod linux {
         }
     }
 
-    // We maintain the convention in libdatadog that the `root_span_id` attribute key is always the
-    // very first in the string table, so its key index is guaranteed to be zero.
+    // We maintain the convention in libdatadog that the `local_root_span_id` attribute key is
+    // always the very first in the string table, so its key index is guaranteed to be zero.
     const ROOT_SPAN_KEY_INDEX: u8 = 0;
 
     /// Maximum size in bytes of the `attrs_data` field.
@@ -169,13 +169,13 @@ pub mod linux {
     }
 
     impl ThreadContextRecord {
-        /// Build a record with the given trace id, span id and attributes. The `root_span_id` is a
-        /// distinguished attribute with special handling for convenience, but it ends up as other
-        /// attributes in `attrs_data`.
+        /// Build a record with the given trace id, span id and attributes. The
+        /// `local_root_span_id` is a distinguished attribute with special handling for
+        /// convenience, but it ends up as other attributes in `attrs_data`.
         fn new(
             trace_id: [u8; 16],
             span_id: [u8; 8],
-            root_span_id: &str,
+            local_root_span_id: [u8; 8],
             attrs: &[(u8, &str)],
         ) -> Self {
             const { assert!(size_of::<ThreadContextRecord>() == 640) }
@@ -185,7 +185,7 @@ pub mod linux {
                 span_id,
                 ..Default::default()
             };
-            record.set_attrs(root_span_id, attrs);
+            record.set_attrs(local_root_span_id, attrs);
             record
         }
 
@@ -210,13 +210,15 @@ pub mod linux {
         /// recovery would require us to be able to rollback to the previous attributes which would
         /// hurt the happy path, or leave the record in a inconsistent state. Another possibility
         /// would be to error out and reset the record in that situation.
-        fn set_attrs(&mut self, root_span_id: &str, attributes: &[(u8, &str)]) -> bool {
-            let mut offset = 0;
+        fn set_attrs(&mut self, local_root_span_id: [u8; 8], attributes: &[(u8, &str)]) -> bool {
             let mut fully_encoded = true;
-            let root_span_attr = (ROOT_SPAN_KEY_INDEX, root_span_id);
-            let full_attrs = std::iter::once(&root_span_attr).chain(attributes);
 
-            for &(key_index, val) in full_attrs {
+            self.attrs_data[0] = ROOT_SPAN_KEY_INDEX;
+            self.attrs_data[1] = 8;
+            self.attrs_data[2..10].copy_from_slice(local_root_span_id.as_slice());
+            let mut offset = 10;
+
+            for &(key_index, val) in attributes {
                 let val_bytes = val.as_bytes();
                 let val_len = val_bytes.len();
                 let val_len = if val_len > 255 {
@@ -277,13 +279,13 @@ pub mod linux {
         pub fn new(
             trace_id: [u8; 16],
             span_id: [u8; 8],
-            root_span_id: &str,
+            local_root_span_id: [u8; 8],
             attrs: &[(u8, &str)],
         ) -> Self {
             Self::from(ThreadContextRecord::new(
                 trace_id,
                 span_id,
-                root_span_id,
+                local_root_span_id,
                 attrs,
             ))
         }
@@ -361,7 +363,7 @@ pub mod linux {
         pub fn update(
             trace_id: [u8; 16],
             span_id: [u8; 8],
-            root_span_id: &str,
+            local_root_span_id: [u8; 8],
             attrs: &[(u8, &str)],
         ) {
             let slot = get_tls_slot();
@@ -372,7 +374,7 @@ pub mod linux {
 
                 current.trace_id = trace_id;
                 current.span_id = span_id;
-                current.set_attrs(root_span_id, attrs);
+                current.set_attrs(local_root_span_id, attrs);
 
                 compiler_fence(Ordering::SeqCst);
                 current.valid.store(1, Ordering::Relaxed);
@@ -382,7 +384,7 @@ pub mod linux {
                 // `ThreadContext::new` already initialises `valid = 1`.
                 let _ = Self::swap(
                     slot,
-                    ThreadContext::new(trace_id, span_id, root_span_id, attrs).into_raw(),
+                    ThreadContext::new(trace_id, span_id, local_root_span_id, attrs).into_raw(),
                 );
             }
         }
@@ -422,12 +424,13 @@ pub mod linux {
         fn tls_lifecycle_basic() {
             let trace_id = [1u8; 16];
             let span_id = [2u8; 8];
+            let root_span_id = [3u8; 8];
 
             assert!(
                 read_tls_context_ptr().is_null(),
                 "TLS must be null initially"
             );
-            ThreadContext::new(trace_id, span_id, "", &[]).attach();
+            ThreadContext::new(trace_id, span_id, root_span_id, &[]).attach();
             assert!(
                 !read_tls_context_ptr().is_null(),
                 "TLS must not be null after attach"
@@ -457,8 +460,9 @@ pub mod linux {
         fn raw_tls_pointer_read() {
             let trace_id = [1u8; 16];
             let span_id = [2u8; 8];
+            let root_span_id = [3u8; 8];
 
-            ThreadContext::new(trace_id, span_id, "", &[]).attach();
+            ThreadContext::new(trace_id, span_id, root_span_id, &[]).attach();
 
             let ptr = read_tls_context_ptr();
             assert!(!ptr.is_null(), "TLS must be non-null after attach");
@@ -468,32 +472,34 @@ pub mod linux {
             assert_eq!(record.trace_id, trace_id);
             assert_eq!(record.span_id, span_id);
             assert_eq!(record.valid.load(Ordering::Relaxed), 1);
-            assert_eq!(record.attrs_data_size, 2);
+            // 1 (key) + 1 (len) + 8 (root_span_id bytes) = 10
+            assert_eq!(record.attrs_data_size, 10);
 
             let _ = ThreadContext::detach();
         }
 
         #[test]
         #[cfg_attr(miri, ignore)]
         fn attribute_encoding_basic() {
-            let root_span_id = "aaaa-bbbb";
+            let root_span_id = [0u8; 8];
             let attrs: &[(u8, &str)] = &[(1, "GET"), (2, "/api/v1")];
             ThreadContext::new([0u8; 16], [0u8; 8], root_span_id, attrs).attach();
 
             let ptr = read_tls_context_ptr();
             assert!(!ptr.is_null());
             let record = unsafe { &*ptr };
-            let expected_size: u16 = (2 + 9 + 2 + 3 + 2 + 7) as u16;
+            // 1+1+8 (root_span_id) + 1+1+3 (GET) + 1+1+7 (/api/v1)
+            let expected_size: u16 = (2 + 8 + 2 + 3 + 2 + 7) as u16;
             assert_eq!(record.attrs_data_size, expected_size);
             assert_eq!(record.attrs_data[0], 0);
-            assert_eq!(record.attrs_data[1], 9);
-            assert_eq!(&record.attrs_data[2..11], b"aaaa-bbbb");
-            assert_eq!(record.attrs_data[11], 1);
-            assert_eq!(record.attrs_data[12], 3);
-            assert_eq!(&record.attrs_data[13..16], b"GET");
-            assert_eq!(record.attrs_data[16], 2);
-            assert_eq!(record.attrs_data[17], 7);
-            assert_eq!(&record.attrs_data[18..25], b"/api/v1");
+            assert_eq!(record.attrs_data[1], 8);
+            assert_eq!(&record.attrs_data[2..10], &root_span_id);
+            assert_eq!(record.attrs_data[10], 1);
+            assert_eq!(record.attrs_data[11], 3);
+            assert_eq!(&record.attrs_data[12..15], b"GET");
+            assert_eq!(record.attrs_data[15], 2);
+            assert_eq!(record.attrs_data[16], 7);
+            assert_eq!(&record.attrs_data[17..24], b"/api/v1");
 
             let _ = ThreadContext::detach();
         }
@@ -503,30 +509,31 @@ pub mod linux {
         fn attribute_truncation_on_overflow() {
             // Build attributes whose combined encoded size exceeds MAX_ATTRS_DATA_SIZE.
             // Each max entry: 1 (key) + 1 (len) + 255 (val) = 257 bytes.
-            // Two such entries: 514 bytes, plus an empty root_span_id: 516.
-            // A third entry of 100 chars would need 102 bytes, bringing the total to 618 > 612, so
+            // root_span_id: 1 (key) + 1 (len) + 8 (val) = 10 bytes.
+            // Two such entries: 514 bytes, plus root_span_id: 524.
+            // A third entry of 100 chars would need 102 bytes, bringing the total to 626 > 612, so
             // the third entry must be dropped.
             let val_a = "a".repeat(255); // 257 bytes encoded
             let val_b = "b".repeat(255); // 257 bytes encoded → 514 total
-            let val_c = "c".repeat(100); // 102 bytes encoded → 616 total: must be dropped
+            let val_c = "c".repeat(100); // 102 bytes encoded → 626 total: must be dropped
 
             let attrs: &[(u8, &str)] = &[
                 (1, val_a.as_str()),
                 (2, val_b.as_str()),
                 (3, val_c.as_str()),
             ];
 
-            ThreadContext::new([0u8; 16], [0u8; 8], "", attrs).attach();
+            ThreadContext::new([0u8; 16], [0u8; 8], [0u8; 8], attrs).attach();
 
             let ptr = read_tls_context_ptr();
             assert!(!ptr.is_null());
             let record = unsafe { &*ptr };
-            // Only the first two entries fit (514 bytes + 2 bytes for empty root span id).
-            assert_eq!(record.attrs_data_size, 516);
-            assert_eq!(record.attrs_data[2], 1);
-            assert_eq!(record.attrs_data[3], 255);
-            assert_eq!(record.attrs_data[259], 2);
-            assert_eq!(record.attrs_data[260], 255);
+            // Only the first two entries fit (514 bytes + 10 bytes for root_span_id).
+            assert_eq!(record.attrs_data_size, 524);
+            assert_eq!(record.attrs_data[10], 1);
+            assert_eq!(record.attrs_data[11], 255);
+            assert_eq!(record.attrs_data[267], 2);
+            assert_eq!(record.attrs_data[268], 255);
 
             let _ = ThreadContext::detach();
         }
@@ -536,10 +543,10 @@ pub mod linux {
         fn update_record_in_place() {
             let trace_id1 = [1u8; 16];
             let span_id1 = [1u8; 8];
-            let root_span_id1 = "xxxx";
+            let root_span_id1 = [0x78u8; 8];
             let trace_id2 = [2u8; 16];
             let span_id2 = [2u8; 8];
-            let root_span_id2 = "yyyy";
+            let root_span_id2 = [0x79u8; 8];
 
             // Updating before any context is attached should be equivalent to `attach()`
             ThreadContext::update(trace_id1, span_id1, root_span_id1, &[(0, "v1")]);
@@ -551,11 +558,11 @@ pub mod linux {
             assert_eq!(record.span_id, span_id1);
             assert_eq!(record.valid.load(Ordering::Relaxed), 1);
             assert_eq!(record.attrs_data[0], 0);
-            assert_eq!(record.attrs_data[1], 4);
-            assert_eq!(&record.attrs_data[2..6], root_span_id1.as_bytes());
-            assert_eq!(record.attrs_data[6], 0);
-            assert_eq!(record.attrs_data[7], 2);
-            assert_eq!(&record.attrs_data[8..10], b"v1");
+            assert_eq!(record.attrs_data[1], 8);
+            assert_eq!(&record.attrs_data[2..10], &root_span_id1);
+            assert_eq!(record.attrs_data[10], 0);
+            assert_eq!(record.attrs_data[11], 2);
+            assert_eq!(&record.attrs_data[12..14], b"v1");
 
             ThreadContext::update(trace_id2, span_id2, root_span_id2, &[(0, "v2")]);
 
@@ -570,11 +577,11 @@ pub mod linux {
             assert_eq!(record.span_id, span_id2);
             assert_eq!(record.valid.load(Ordering::Relaxed), 1);
             assert_eq!(record.attrs_data[0], 0);
-            assert_eq!(record.attrs_data[1], 4);
-            assert_eq!(&record.attrs_data[2..6], root_span_id2.as_bytes());
-            assert_eq!(record.attrs_data[6], 0);
-            assert_eq!(record.attrs_data[7], 2);
-            assert_eq!(&record.attrs_data[8..10], b"v2");
+            assert_eq!(record.attrs_data[1], 8);
+            assert_eq!(&record.attrs_data[2..10], &root_span_id2);
+            assert_eq!(record.attrs_data[10], 0);
+            assert_eq!(record.attrs_data[11], 2);
+            assert_eq!(&record.attrs_data[12..14], b"v2");
 
             let _ = ThreadContext::detach();
             assert!(read_tls_context_ptr().is_null());
@@ -583,7 +590,7 @@ pub mod linux {
         #[test]
         #[cfg_attr(miri, ignore)]
         fn explicit_detach_nulls_tls() {
-            ThreadContext::new([3u8; 16], [3u8; 8], "aaaa", &[]).attach();
+            ThreadContext::new([0u8; 16], [0u8; 8], [0u8; 8], &[]).attach();
             assert!(!read_tls_context_ptr().is_null());
 
             let _ = ThreadContext::detach();
@@ -598,14 +605,16 @@ pub mod linux {
         #[cfg_attr(miri, ignore)]
         fn long_value_capped_at_255_bytes() {
             let long_val = "a".repeat(300);
-            ThreadContext::new([0u8; 16], [0u8; 8], "bbbb", &[(0, long_val.as_str())]).attach();
+            ThreadContext::new([0u8; 16], [0u8; 8], [0u8; 8], &[(0, long_val.as_str())]).attach();
 
             let ptr = read_tls_context_ptr();
             assert!(!ptr.is_null());
             let record = unsafe { &*ptr };
-            let val_len = record.attrs_data[2 + 4 + 1];
+            // root_span_id occupies offset 0..10, then the attr entry starts at 10: key at [10],
+            // len at [11]
+            let val_len = record.attrs_data[2 + 8 + 1];
             assert_eq!(val_len, 255, "value must be capped at 255 bytes");
-            assert_eq!(record.attrs_data_size, 2 + 4 + 2 + 255);
+            assert_eq!(record.attrs_data_size, 2 + 8 + 2 + 255);
 
             let _ = ThreadContext::detach();
         }
@@ -621,10 +630,10 @@ pub mod linux {
 
             let spawned_trace_id = [0xABu8; 16];
             let spawned_span_id = [0xCDu8; 8];
-            let spawned_root_span_id = "xxxx";
+            let spawned_root_span_id = [0xEFu8; 8];
             let main_trace_id = [0x11u8; 16];
             let main_span_id = [0x22u8; 8];
-            let main_root_span_id = "yyyy";
+            let main_root_span_id = [0x33u8; 8];
 
             let handle = std::thread::spawn(move || {
                 ThreadContext::new(spawned_trace_id, spawned_span_id, spawned_root_span_id, &[])
@@ -641,10 +650,7 @@ pub mod linux {
                 let record = unsafe { &*ptr };
                 assert_eq!(record.trace_id, spawned_trace_id);
                 assert_eq!(record.span_id, spawned_span_id);
-                assert_eq!(
-                    &record.attrs_data[2..(record.attrs_data_size as usize)],
-                    spawned_root_span_id.as_bytes()
-                );
+                assert_eq!(&record.attrs_data[2..10], &spawned_root_span_id);
 
                 let _ = ThreadContext::detach();
                 assert!(read_tls_context_ptr().is_null());
@@ -665,10 +671,7 @@ pub mod linux {
             let record = unsafe { &*ptr };
             assert_eq!(record.trace_id, main_trace_id);
             assert_eq!(record.span_id, main_span_id);
-            assert_eq!(
-                &record.attrs_data[2..(record.attrs_data_size as usize)],
-                main_root_span_id.as_bytes()
-            );
+            assert_eq!(&record.attrs_data[2..10], &main_root_span_id);
 
             barrier.wait();