diff --git a/Cargo.lock b/Cargo.lock index 50064f363..60f7da9ea 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -830,11 +830,14 @@ dependencies = [ "byteorder", "bytes", "chrono", + "criterion", + "datadog-alloc", "ddcommon", "derivative", "futures", "futures-core", "futures-util", + "hashbrown 0.14.3", "http", "http-body", "hyper", diff --git a/alloc/src/chain.rs b/alloc/src/chain.rs index 7df9c9588..7defb6e17 100644 --- a/alloc/src/chain.rs +++ b/alloc/src/chain.rs @@ -104,13 +104,15 @@ impl ChainAllocator { #[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::>(); + 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 Allocator for ChainAllocator { fn allocate(&self, layout: Layout) -> Result, AllocError> { let layout = layout.pad_to_align(); - // Too large for ChainAllocator to deal with. - let header_overhead = size_of::>(); - 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::>(); + 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::()).unwrap(); + unsafe { allocator.deallocate(ptr.cast(), Layout::new::()) }; + } + // 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(allocator: &ChainAllocator) { 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] diff --git a/profiling/Cargo.toml b/profiling/Cargo.toml index 1dea03960..570c9a775 100644 --- a/profiling/Cargo.toml +++ b/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" diff --git a/profiling/benches/interning_strings.rs b/profiling/benches/interning_strings.rs new file mode 100644 index 000000000..2bec040a4 --- /dev/null +++ b/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>, + } + + impl StringTable { + pub fn new() -> Self { + let mut strings = FxIndexSet::>::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); diff --git a/profiling/benches/main.rs b/profiling/benches/main.rs new file mode 100644 index 000000000..e2d220f76 --- /dev/null +++ b/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); diff --git a/profiling/src/collections/mod.rs b/profiling/src/collections/mod.rs index 435ca2bd6..66281e671 100644 --- a/profiling/src/collections/mod.rs +++ b/profiling/src/collections/mod.rs @@ -2,3 +2,4 @@ // SPDX-License-Identifier: Apache-2.0 pub mod identifiable; +pub mod string_table; diff --git a/profiling/src/collections/string_table/mod.rs b/profiling/src/collections/string_table/mod.rs new file mode 100644 index 000000000..e47f33c20 --- /dev/null +++ b/profiling/src/collections/string_table/mod.rs @@ -0,0 +1,277 @@ +// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +#[allow(unused)] +pub mod wordpress_test_data; + +use crate::collections::identifiable::{Id, StringId}; +use crate::iter::{IntoLendingIterator, LendingIterator}; +use datadog_alloc::{AllocError, Allocator, ChainAllocator, VirtualAllocator}; +use std::alloc::Layout; + +/// A trait that indicates an allocator is arena allocator, meaning it doesn't +/// deallocate individual items, but deallocates their memory as a group when +/// the arena is dropped. +pub trait ArenaAllocator: Allocator { + /// Copies the str into the arena, and returns a slice to the new str. + fn allocate(&self, str: &str) -> Result<&str, AllocError> { + let layout = Layout::for_value(str); + let uninit_ptr = Allocator::allocate(self, layout)?; + + // Copy the bytes of the string into the allocated memory. + // SAFETY: this is guaranteed to not be overlapping because an + // allocator must not return aliasing bytes in its allocations. + unsafe { + let src = str.as_ptr(); + let dst = uninit_ptr.as_ptr() as *mut u8; + let count = str.len(); + core::ptr::copy_nonoverlapping(src, dst, count); + } + + // SAFETY: The bytes were properly initialized, and they cannot be + // misaligned because they have an alignment of 1, so it is safe to + // create a slice of the given data and length. The lifetime matches + // the arena allocator's lifetime. + let slice: &[u8] = + unsafe { core::slice::from_raw_parts(uninit_ptr.as_ptr() as *const u8, str.len()) }; + + // SAFETY: Since the bytes were copied from a valid str without + // slicing, the bytes must also be utf-8. + Ok(unsafe { core::str::from_utf8_unchecked(slice) }) + } +} + +impl ArenaAllocator for ChainAllocator {} + +type Hasher = core::hash::BuildHasherDefault; +type HashSet = indexmap::IndexSet; + +/// Holds unique strings and provides [StringId]s that correspond to the order +/// that the strings were inserted. +pub struct StringTable { + /// The bytes of each string stored in `strings` are allocated here. + bytes: ChainAllocator, + + /// The ordered hash set of unique strings. The order becomes the StringId. + /// The static lifetime is a lie, it is tied to the `bytes`, which is only + /// moved if the string table is moved e.g. + /// [StringTable::into_lending_iterator]. + /// References to the underlying strings should generally not be handed, + /// but if they are, they should be bound to the string table's lifetime + /// or the lending iterator's lifetime. + strings: HashSet<&'static str>, +} + +impl Default for StringTable { + fn default() -> Self { + Self::new() + } +} + +impl StringTable { + /// Creates a new string table, which initially holds the empty string and + /// no others. + pub fn new() -> Self { + // Keep in mind 32-bit .NET. There is only 2 GiB of virtual memory + // total available to an application, and we're not the application, + // we're just a piece inside it. Additionally, there may be 2 or more + // string tables in memory at a given time. Talk to .NET profiling + // engineers before making this any bigger. + const SIZE_HINT: usize = 4 * 1024 * 1024; + let bytes = ChainAllocator::new_in(SIZE_HINT, VirtualAllocator {}); + + let mut strings = HashSet::with_hasher(Hasher::default()); + // It varies by implementation, but frequently I've noticed that the + // capacity after the first insertion is quite small, as in 3. This is + // a bit too small and there are frequent reallocations. For one sample + // with endpoint + code hotspots, we'd have at least these strings: + // - "" + // - At least one sample type + // - At least one sample unit--already at 3 without any samples. + // - "local root span id" + // - "span id" + // - "trace endpoint" + // - A file and/or function name per frame. + // So with a capacity like 3, we end up reallocating a bunch on or + // before the very first sample. The number here is not fine-tuned, + // just skipping some obviously bad, tiny sizes. + strings.reserve(32); + + // Always hold the empty string as item 0. Do not insert it via intern + // because that will try to allocate zero-bytes from the storage, + // which is sketchy. + strings.insert(""); + + Self { bytes, strings } + } + + /// Returns the number of strings currently held in the string table. + #[inline] + #[allow(clippy::len_without_is_empty)] + pub fn len(&self) -> usize { + self.strings.len() + } + + /// Adds the string to the string table if it isn't present already, and + /// returns a [StringId] that corresponds to the order that this string + /// was originally inserted. + /// + /// # Panics + /// This panics if the allocator fails to allocate a new chunk/node. + pub fn intern(&mut self, str: &str) -> StringId { + let set = &mut self.strings; + match set.get_index_of(str) { + Some(offset) => StringId::from_offset(offset), + None => { + // No match. Get the current size of the table, which + // corresponds to the StringId it will have when inserted. + let string_id = StringId::from_offset(set.len()); + + // Make a new string in the arena, and fudge its lifetime + // to appease the borrow checker. + let new_str = { + // PANIC: the intern API doesn't allow for failure, so if + // this allocation fails, panic. The current + // implementation of `ChainAllocator` will fail if the + // underlying allocator fails when asking for a new chunk. + // This is expected to be rare. + let s = ArenaAllocator::allocate(&self.bytes, str) + .expect("allocator for StringTable::intern to succeed"); + + // SAFETY: all references to this value get re-narrowed to + // the lifetime of the string table or iterator when + // exposed to the user. The string table and iterator will + // keep the arena alive, making the access safe. + unsafe { core::mem::transmute::<&str, &'static str>(s) } + }; + + // Add it to the set. + self.strings.insert(new_str); + + string_id + } + } + } +} + +/// A [LendingIterator] for a [StringTable]. Make one by calling +/// [StringTable::into_lending_iter]. +pub struct StringTableIter { + /// This is actually used, the compiler doesn't know that the static + /// references in `iter` actually point in here. + #[allow(unused)] + bytes: ChainAllocator, + + /// The strings of the string table, in order of insertion. + /// The static lifetimes are a lie, they are tied to the `bytes`. When + /// handing out references, bind the lifetime to the iterator's lifetime, + /// which is a [LendingIterator] is needed. + iter: as IntoIterator>::IntoIter, +} + +impl StringTableIter { + fn new(string_table: StringTable) -> StringTableIter { + StringTableIter { + bytes: string_table.bytes, + iter: string_table.strings.into_iter(), + } + } +} + +impl LendingIterator for StringTableIter { + type Item<'a> = &'a str where Self: 'a; + + fn next(&mut self) -> Option> { + self.iter.next() + } + + fn count(self) -> usize { + self.iter.count() + } +} + +impl IntoLendingIterator for StringTable { + type Iter = StringTableIter; + + fn into_lending_iter(self) -> Self::Iter { + StringTableIter::new(self) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_basics() { + let mut table = StringTable::new(); + // The empty string should already be present. + assert_eq!(1, table.len()); + assert_eq!(StringId::ZERO, table.intern("")); + + // Intern a string literal to ensure ?Sized works. + let string = table.intern("datadog"); + assert_eq!(StringId::from_offset(1), string); + assert_eq!(2, table.len()); + } + + #[track_caller] + fn test_from_src(src: &[&str]) { + // Insert all the strings. + let mut table = StringTable::new(); + let n_strings = src.len(); + for string in src { + table.intern(string); + } + assert_eq!(n_strings, table.len()); + + // Re-inserting doesn't change the size. + for string in src { + table.intern(string); + } + assert_eq!(n_strings, table.len()); + + // Check that they are ordered correctly when iterating. + let mut actual_iter = table.into_lending_iter(); + let mut expected_iter = src.iter(); + while let (Some(expected), Some(actual)) = (expected_iter.next(), actual_iter.next()) { + assert_eq!(*expected, actual); + } + + // The iterators should be exhausted at this point. + assert_eq!(None, expected_iter.next()); + assert_eq!(0, actual_iter.count()); + } + + #[test] + fn test_small_set_of_strings() { + let cases: &[_] = &[ + "", + "local root span id", + "span id", + "trace endpoint", + "samples", + "count", + "wall-time", + "nanoseconds", + "cpu-time", + ", stack_traces: FxIndexSet, start_time: SystemTime, - strings: FxIndexSet>, + strings: StringTable, timestamp_key: StringId, upscaling_rules: UpscalingRules, } @@ -257,7 +259,8 @@ impl Profile { encoder.encode(ProfileFunctionsEntry::from(item))?; } - for item in self.strings.into_iter() { + let mut lender = self.strings.into_lending_iter(); + while let Some(item) = lender.next() { encoder.encode_string_table_entry(item)?; } @@ -406,21 +409,9 @@ impl Profile { /// Interns the `str` as a string, returning the id in the string table. /// The empty string is guaranteed to have an id of [StringId::ZERO]. + #[inline] 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()); - // This wouldn't make any sense; the item couldn't be found so - // we try to insert it, but suddenly it exists now? - debug_assert!(_inserted); - index - } - }; - StringId::from_offset(index) + self.strings.intern(item) } /// Creates a profile from the period, sample types, and start time using @@ -847,7 +838,7 @@ mod api_tests { assert_eq!(profile.sample_types, prev.sample_types); // The string table should have at least the empty string. - assert!(!profile.strings.is_empty()); + assert!(profile.strings.len() > 0); } #[test] @@ -855,12 +846,12 @@ mod api_tests { /* The previous test (reset) checked quite a few properties already, so * this one will focus only on the period. */ - let sample_types: &[api::ValueType] = &[api::ValueType::new("wall-time", "nanoseconds")]; + let sample_types = [api::ValueType::new("wall-time", "nanoseconds")]; let period = api::Period { r#type: sample_types[0], value: 10_000_000, }; - let mut profile = Profile::new(SystemTime::now(), sample_types, Some(period)); + let mut profile = Profile::new(SystemTime::now(), &sample_types, Some(period)); let prev = profile .reset_and_return_previous(None) @@ -868,11 +859,16 @@ mod api_tests { // Resolve the string values to check that they match (their string // table offsets may not match). + let mut strings = Vec::with_capacity(profile.strings.len()); + let mut strings_iter = profile.strings.into_lending_iter(); + while let Some(item) = strings_iter.next() { + strings.push(Box::from(String::from(item))); + } + for (value, period_type) in [profile.period.unwrap(), prev.period.unwrap()] { assert_eq!(value, period.value); - let strings = profile.strings.iter().collect::>(); - let r#type: &str = strings[period_type.r#type.to_offset()]; - let unit: &str = strings[period_type.unit.to_offset()]; + let r#type: &str = &strings[period_type.r#type.to_offset()]; + let unit: &str = &strings[period_type.unit.to_offset()]; assert_eq!(r#type, period.r#type.r#type); assert_eq!(unit, period.r#type.unit); } diff --git a/profiling/src/iter.rs b/profiling/src/iter.rs new file mode 100644 index 000000000..b1afbead3 --- /dev/null +++ b/profiling/src/iter.rs @@ -0,0 +1,23 @@ +// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +/// The [LendingIterator] is a version of an [Iterator] that can yield items +/// with references into the lender. It is a well-known name and there are +/// multiple crates which offer it, with differences. The needs here are +/// small, and so rather than bring in a pre-1.0 crate, just make our own. +pub trait LendingIterator { + type Item<'a> + where + Self: 'a; + + fn next(&mut self) -> Option>; + + #[allow(unused)] + fn count(self) -> usize; +} + +/// Turn a collection of some sort into a [LendingIterator]. +pub trait IntoLendingIterator { + type Iter: LendingIterator; + fn into_lending_iter(self) -> Self::Iter; +} diff --git a/profiling/src/lib.rs b/profiling/src/lib.rs index 1da8d7e5e..d450ab68e 100644 --- a/profiling/src/lib.rs +++ b/profiling/src/lib.rs @@ -5,5 +5,6 @@ pub mod api; pub mod collections; pub mod exporter; pub mod internal; +pub mod iter; pub mod pprof; pub mod serializer; diff --git a/tools/docker/Dockerfile.build b/tools/docker/Dockerfile.build index 5f1227660..58124a47c 100644 --- a/tools/docker/Dockerfile.build +++ b/tools/docker/Dockerfile.build @@ -106,7 +106,7 @@ COPY "data-pipeline/Cargo.toml" "data-pipeline/" COPY "data-pipeline-ffi/Cargo.toml" "data-pipeline-ffi/" COPY "bin_tests/Cargo.toml" "bin_tests/" RUN find -name "Cargo.toml" | sed -e s#Cargo.toml#src/lib.rs#g | xargs -n 1 sh -c 'mkdir -p $(dirname $1); touch $1; echo $1' create_stubs -RUN echo trace-obfuscation/benches/trace_obfuscation.rs tools/src/bin/dedup_headers.rs tools/sidecar_mockgen/src/bin/sidecar_mockgen.rs ddtelemetry/examples/tm-worker-test.rs ipc/tarpc/tarpc/examples/compression.rs ipc/tarpc/tarpc/examples/custom_transport.rs ipc/tarpc/tarpc/examples/pubsub.rs ipc/tarpc/tarpc/examples/readme.rs ipc/tarpc/tarpc/examples/tracing.rs ipc/tarpc/tarpc/tests/compile_fail.rs ipc/tarpc/tarpc/tests/dataservice.rs ipc/tarpc/tarpc/tests/service_functional.rs bin_tests/src/bin/crashtracker_bin_test.rs bin_tests/src/bin/test_the_tests.rs | xargs -n 1 sh -c 'mkdir -p $(dirname $1); touch $1; echo $1' create_stubs +RUN echo profiling/benches/main.rs profiling/benches/interning_strings.rs trace-obfuscation/benches/trace_obfuscation.rs tools/src/bin/dedup_headers.rs tools/sidecar_mockgen/src/bin/sidecar_mockgen.rs ddtelemetry/examples/tm-worker-test.rs ipc/tarpc/tarpc/examples/compression.rs ipc/tarpc/tarpc/examples/custom_transport.rs ipc/tarpc/tarpc/examples/pubsub.rs ipc/tarpc/tarpc/examples/readme.rs ipc/tarpc/tarpc/examples/tracing.rs ipc/tarpc/tarpc/tests/compile_fail.rs ipc/tarpc/tarpc/tests/dataservice.rs ipc/tarpc/tarpc/tests/service_functional.rs bin_tests/src/bin/crashtracker_bin_test.rs bin_tests/src/bin/test_the_tests.rs | xargs -n 1 sh -c 'mkdir -p $(dirname $1); touch $1; echo $1' create_stubs # cache dependencies RUN cargo fetch --locked @@ -130,4 +130,5 @@ COPY ./ ./ RUN ./build-profiling-ffi.sh /build/output FROM scratch as ffi_build_output + COPY --from=ffi_build /build/output/ ./