[profiling] Add fuzz test for StringTable (#437)

.github/workflows/fuzz.yml

@@ -0,0 +1,25 @@
+name: Fuzz test
+on:
+  push:
+
+jobs:
+  run-fuzz:
+    runs-on: ubuntu-latest
+    env:
+      CARGO_TERM_COLOR: always
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Rust
+      run: |
+        set -e
+        rustup set profile minimal
+        rustup toolchain install nightly 
+        rustup default nightly
+    - uses: taiki-e/install-action@v2
+      with: 
+        tool: cargo-bolero
+    - run: |
+        set -e
+        (cd profiling && cargo bolero test collections::string_table::tests::fuzz_string_table -T 1min)
+        (cd profiling && cargo bolero test collections::string_table::tests::fuzz_arena_allocator -T 1min)
+

Cargo.toml

@@ -53,3 +53,10 @@ opt-level = "s" # optimize for size
 
 [profile.release.package.datadog-serverless-trace-mini-agent]
 strip = true
+
+# https://camshaft.github.io/bolero/library-installation.html
+[profile.fuzz]
+inherits = "dev"
+opt-level = 3
+incremental = false
+codegen-units = 1

alloc/src/chain.rs

@@ -214,6 +214,9 @@ impl<A: Allocator + Clone> ChainAllocator<A> {
 unsafe impl<A: Allocator + Clone> Allocator for ChainAllocator<A> {
     #[cfg_attr(debug_assertions, track_caller)]
     fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        if layout.size() == 0 {
+            return Err(AllocError);
+        }
         let layout = layout.pad_to_align();
 
         let remaining_capacity = self.remaining_capacity();

alloc/src/linear.rs

@@ -78,6 +78,10 @@ impl<A: Allocator> Drop for LinearAllocator<A> {
 
 unsafe impl<A: Allocator> Allocator for LinearAllocator<A> {
     fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        if layout.size() == 0 {
+            return Err(AllocError);
+        }
+
         // Find the needed allocation size including the necessary alignment.
         let size = self.used_bytes();
         // SAFETY: base_ptr + size will always be in the allocated range, or

alloc/src/virtual.rs

@@ -94,6 +94,10 @@ pub mod os {
         }
 
         fn allocate_zeroed(&self, layout: Layout) -> Result<ptr::NonNull<[u8]>, AllocError> {
+            if layout.size() == 0 {
+                return Err(AllocError);
+            }
+
             let size = super::layout_to_page_size(layout)?;
 
             let null = ptr::null_mut();

profiling/Cargo.toml

@@ -46,4 +46,6 @@ tokio-util = "0.7.1"
 byteorder = { version = "1.5", features = ["std"] }
 
 [dev-dependencies]
+bolero = "0.10.1"
 criterion = "0.5.1"
+

profiling/src/collections/string_table/mod.rs

@@ -15,6 +15,11 @@ use std::alloc::Layout;
 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> {
+        // TODO: We might want each allocator to return its own empty string
+        // so we can debug where the value came from.
+        if str.is_empty() {
+            return Ok("");
+        }
         let layout = Layout::for_value(str);
         let uninit_ptr = Allocator::allocate(self, layout)?;
 
@@ -202,6 +207,78 @@ impl IntoLendingIterator for StringTable {
 mod tests {
     use super::*;
 
+    #[test]
+    fn fuzz_arena_allocator() {
+        bolero::check!()
+            .with_type::<(usize, Vec<String>)>()
+            .for_each(|(size_hint, strings)| {
+                // If the size_hint is insanely large, get allowed allocation
+                // failures.  These are not interesting, so avoid them.
+                if *size_hint > 4 * 1024 * 1024 * 1024 {
+                    return;
+                }
+                let bytes = ChainAllocator::new_in(*size_hint, VirtualAllocator {});
+                let mut allocated_strings = vec![];
+                for string in strings {
+                    let s =
+                        ArenaAllocator::allocate(&bytes, string).expect("allocation to succeed");
+                    assert_eq!(s, string);
+                    allocated_strings.push(s);
+                }
+                assert_eq!(strings.len(), allocated_strings.len());
+                strings
+                    .iter()
+                    .zip(allocated_strings.iter())
+                    .for_each(|(s, t)| assert_eq!(s, t));
+            });
+    }
+
+    /// This is a fuzz test for the allocation optimized `StringTable`.
+    /// It checks both safety (lack of crashes / sanitizer failures),
+    /// as well as functional correctness (the table should behave like an
+    /// ordered set).  
+    /// Limitations:
+    ///   - The crate used here to generate Strings internally has a default range for the length of
+    ///     a string, (0..=64) We should experiment with longer strings to see what happens. https://github.com/camshaft/bolero/blob/f401669697ffcbe7f34cbfd09fd57b93d5df734c/lib/bolero-generator/src/alloc/mod.rs#L17
+    ///   - Since iterating is destructive, can only check the string values once.
+    /// `cargo +nightly bolero test collections::string_table::tests::fuzz_string_table -T 1min`
+    #[test]
+    fn fuzz_string_table() {
+        bolero::check!()
+            .with_type::<Vec<String>>()
+            .for_each(|strings| {
+                // Compare our optimized implementation against a "golden" version
+                // from the standard library.
+                let mut golden_list = vec![""];
+                let mut golden_set = std::collections::HashSet::from([""]);
+                let mut st = StringTable::new();
+
+                for string in strings {
+                    assert_eq!(st.len(), golden_set.len());
+                    if golden_set.insert(string) {
+                        golden_list.push(string);
+                    }
+
+                    let str_id = st.intern(string);
+                    // The str_id should refer to the id_th string interned
+                    // on the list.  We can't look inside the `StringTable`
+                    // in a non-desctrive way, but fortunately we have the
+                    // `golden_list` to compare against.
+                    assert_eq!(string, golden_list[str_id.to_offset()]);
+                }
+                assert_eq!(st.len(), golden_list.len());
+                assert_eq!(st.len(), golden_set.len());
+
+                // Check that the strings remain in order
+                let mut it = st.into_lending_iter();
+                let mut idx = 0;
+                while let Some(s) = it.next() {
+                    assert_eq!(s, golden_list[idx]);
+                    idx += 1;
+                }
+            })
+    }
+
     #[test]
     fn test_basics() {
         let mut table = StringTable::new();