perf(profiling): datadog-alloc based StringTable

alloc/src/chain.rs

@@ -104,13 +104,15 @@ impl<A: Allocator + Clone> ChainAllocator<A> {
 
     #[cold]
     #[inline(never)]
-    fn grow(&self) -> Result<(), AllocError> {
+    fn grow(&self, min_size: usize) -> Result<(), AllocError> {
         let top = self.top.get();
         let chain_layout = Layout::new::<ChainNode<A>>();
 
+        let node_size = min_size.max(self.node_size);
         let linear = {
-            let layout = Layout::from_size_align(self.node_size, chain_layout.align())
-                .map_err(|_| AllocError)?;
+            let layout = Layout::from_size_align(node_size, chain_layout.align())
+                .map_err(|_| AllocError)?
+                .pad_to_align();
             LinearAllocator::new_in(layout, self.allocator.clone())?
         };
 
@@ -214,16 +216,14 @@ unsafe impl<A: Allocator + Clone> Allocator for ChainAllocator<A> {
     fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
         let layout = layout.pad_to_align();
 
-        // Too large for ChainAllocator to deal with.
-        let header_overhead = size_of::<ChainNode<A>>();
-        let maximum_capacity = self.node_size - header_overhead;
-        if layout.size() > maximum_capacity {
-            return Err(AllocError);
-        }
-
         let remaining_capacity = self.remaining_capacity();
         if layout.size() > remaining_capacity {
-            self.grow()?;
+            let header_overhead = size_of::<ChainNode<A>>();
+            let min_size = layout
+                .size()
+                .checked_add(header_overhead)
+                .ok_or(AllocError)?;
+            self.grow(min_size)?;
         }
 
         // At this point:
@@ -298,6 +298,32 @@ mod tests {
         unsafe { allocator.deallocate(ptr.cast(), layout) };
     }
 
+    #[test]
+    fn test_large_allocations() {
+        let allocator = ChainAllocator::new_in(4096, Global);
+
+        // Force an allocation, so it makes a chunk of the minimum size.
+        {
+            let ptr = allocator.allocate(Layout::new::<u8>()).unwrap();
+            unsafe { allocator.deallocate(ptr.cast(), Layout::new::<u8>()) };
+        }
+        // Should be a bit less than 4096, but use this over hard-coding a
+        // number, to make it more resilient to implementation changes.
+        let remaining_capacity = allocator.remaining_capacity();
+
+        // Now make something bigger than the chunk.
+        let size = 4 * (remaining_capacity + 1);
+        let layout = Layout::from_size_align(size, 1).unwrap();
+        let ptr = allocator.allocate(layout).unwrap();
+        let actual_size = ptr.len();
+        assert!(
+            actual_size >= size,
+            "failed to allocate large allocation, expected at least {size} bytes, saw {actual_size}"
+        );
+        // Doesn't return memory, just ensuring we don't panic.
+        unsafe { allocator.deallocate(ptr.cast(), layout) };
+    }
+
     #[track_caller]
     fn fill_to_capacity<A: Allocator + Clone>(allocator: &ChainAllocator<A>) {
         let remaining_capacity = allocator.remaining_capacity();
@@ -306,9 +332,8 @@ mod tests {
             let ptr = allocator.allocate(layout).unwrap();
             // Doesn't return memory, just ensuring we don't panic.
             unsafe { allocator.deallocate(ptr.cast(), layout) };
+            assert_eq!(0, allocator.remaining_capacity());
         }
-        let remaining_capacity = allocator.remaining_capacity();
-        assert_eq!(0, remaining_capacity);
     }
 
     #[test]

profiling/Cargo.toml

@@ -11,16 +11,22 @@ license.workspace = true
 [lib]
 crate-type = ["lib"]
 
+[[bench]]
+name = "main"
+harness = false
+
 [dependencies]
 anyhow = "1.0"
 bitmaps = "3.2.0"
 bytes = "1.1"
 chrono = {version = "0.4", default-features = false, features = ["std", "clock"]}
+datadog-alloc = {path = "../alloc"}
 ddcommon = {path = "../ddcommon"}
 derivative = "2.2.0"
 futures = { version = "0.3", default-features = false }
 futures-core = {version = "0.3.0", default-features = false}
 futures-util = {version = "0.3.0", default-features = false}
+hashbrown = { version = "0.14", default-features = false, features = ["allocator-api2"] }
 http = "0.2"
 http-body = "0.4"
 hyper = {version = "0.14", features = ["client"], default-features = false}
@@ -38,3 +44,6 @@ serde_json = {version = "1.0"}
 tokio = {version = "1.23", features = ["rt", "macros"]}
 tokio-util = "0.7.1"
 byteorder = { version = "1.5", features = ["std"] }
+
+[dev-dependencies]
+criterion = "0.5.1"

profiling/benches/interning_strings.rs

@@ -0,0 +1,68 @@
+// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+use criterion::*;
+use datadog_profiling::collections::string_table::wordpress_test_data::WORDPRESS_STRINGS;
+
+/// This version is the one we used before having datadog-alloc.
+#[allow(unused)]
+mod old_version {
+    use datadog_profiling::collections::identifiable::{FxIndexSet, Id, StringId};
+    pub struct StringTable {
+        strings: FxIndexSet<Box<str>>,
+    }
+
+    impl StringTable {
+        pub fn new() -> Self {
+            let mut strings = FxIndexSet::<Box<str>>::default();
+            strings.insert("".into());
+            Self { strings }
+        }
+
+        pub fn intern(&mut self, item: &str) -> StringId {
+            // For performance, delay converting the [&str] to a [String] until
+            // after it has been determined to not exist in the set. This avoids
+            // temporary allocations.
+            let index = match self.strings.get_index_of(item) {
+                Some(index) => index,
+                None => {
+                    let (index, _inserted) = self.strings.insert_full(item.into());
+                    debug_assert!(_inserted);
+                    index
+                }
+            };
+            StringId::from_offset(index)
+        }
+
+        #[inline]
+        #[allow(clippy::len_without_is_empty)]
+        pub fn len(&self) -> usize {
+            self.strings.len()
+        }
+    }
+}
+
+// To benchmark a different implementation, import a different one.
+use datadog_profiling::collections::string_table::StringTable;
+// use old_version::StringTable;
+
+pub fn small_wordpress_profile(c: &mut Criterion) {
+    c.bench_function("benching string interning on wordpress profile", |b| {
+        b.iter(|| {
+            let mut table = StringTable::new();
+            let n_strings = WORDPRESS_STRINGS.len();
+            for string in WORDPRESS_STRINGS {
+                black_box(table.intern(string));
+            }
+            assert_eq!(n_strings, table.len());
+
+            // re-insert, should nothing should be inserted.
+            for string in WORDPRESS_STRINGS {
+                black_box(table.intern(string));
+            }
+            assert_eq!(n_strings, table.len())
+        })
+    });
+}
+
+criterion_group!(benches, small_wordpress_profile);

profiling/benches/main.rs

@@ -0,0 +1,8 @@
+// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+use criterion::criterion_main;
+
+mod interning_strings;
+
+criterion_main!(interning_strings::benches);

profiling/src/collections/mod.rs

@@ -2,3 +2,4 @@
 // SPDX-License-Identifier: Apache-2.0
 
 pub mod identifiable;
+pub mod string_table;