Add allocation measurement to stress2 proper

Author: spacejam

benchmarks/stress2/Cargo.toml

@@ -20,11 +20,11 @@ event_log = ["sled/event_log"]
 compression = ["sled/compression"]
 no_logs = ["sled/no_logs"]
 metrics = ["sled/metrics"]
-measure_allocs = ["sled/measure_allocs"]
 jemalloc = ["jemallocator"]
 logging = ["env_logger", "log", "color-backtrace"]
 dh = ["dhat"]
 memshred = []
+measure_allocs = []
 
 [dependencies]
 rand = "0.7.3"

benchmarks/stress2/src/main.rs

@@ -14,16 +14,16 @@ use rand::{thread_rng, Rng};
 
 #[cfg(feature = "jemalloc")]
 mod alloc {
-    use std::alloc::Layout;
     use jemallocator::Jemalloc;
+    use std::alloc::Layout;
 
     #[global_allocator]
     static ALLOCATOR: Jemalloc = Jemalloc;
 }
 
 #[cfg(feature = "memshred")]
 mod alloc {
-    use std::alloc::{System, Layout};
+    use std::alloc::{Layout, System};
 
     #[global_allocator]
     static ALLOCATOR: Alloc = Alloc;
@@ -42,11 +42,35 @@ mod alloc {
         unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
             std::ptr::write_bytes(ptr, 0xde, layout.size());
             System.dealloc(ptr, layout)
-
         }
     }
 }
 
+#[cfg(feature = "measure_allocs")]
+mod alloc {
+    use std::alloc::{Layout, System};
+    use std::sync::atomic::{AtomicUsize, Ordering::Release};
+
+    pub static ALLOCATIONS: AtomicUsize = AtomicUsize::new(0);
+    pub static ALLOCATED_BYTES: AtomicUsize = AtomicUsize::new(0);
+
+    #[global_allocator]
+    static ALLOCATOR: Alloc = Alloc;
+
+    #[derive(Default, Debug, Clone, Copy)]
+    struct Alloc;
+
+    unsafe impl std::alloc::GlobalAlloc for Alloc {
+        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+            ALLOCATIONS.fetch_add(1, Release);
+            ALLOCATED_BYTES.fetch_add(layout.size(), Release);
+            System.alloc(layout)
+        }
+        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+            System.dealloc(ptr, layout)
+        }
+    }
+}
 
 #[global_allocator]
 #[cfg(feature = "dh")]
@@ -385,6 +409,17 @@ fn main() {
         ((ops * 1_000) / (time * 1_000)).to_formatted_string(&Locale::en)
     );
 
+    #[cfg(feature = "measure_allocs")]
+    println!(
+        "allocated {} bytes in {} allocations",
+        alloc::ALLOCATED_BYTES
+            .load(Ordering::Acquire)
+            .to_formatted_string(&Locale::en),
+        alloc::ALLOCATIONS
+            .load(Ordering::Acquire)
+            .to_formatted_string(&Locale::en),
+    );
+
     #[cfg(feature = "metrics")]
     sled::print_profile();
 }