chore(cubestore): Add simple allocator tracking stats for benchmarks

Author: ovr

rust/cubestore/cubestore/benches/metastore.rs

@@ -113,8 +113,11 @@ fn do_get_tables_with_path_bench(
 ) {
     let total_tables = num_schemas * tables_per_schema;
     let metastore = runtime.block_on(async {
-        let metastore = prepare_metastore(&format!("get_tables_with_path_{}", total_tables)).unwrap();
-        populate_metastore(&metastore, num_schemas, tables_per_schema).await.unwrap();
+        let metastore =
+            prepare_metastore(&format!("get_tables_with_path_{}", total_tables)).unwrap();
+        populate_metastore(&metastore, num_schemas, tables_per_schema)
+            .await
+            .unwrap();
         metastore
     });
 
@@ -139,7 +142,9 @@ fn do_get_tables_with_path_bench(
 
 async fn do_cold_cache_test(num_schemas: usize, tables_per_schema: usize) {
     let fresh_metastore = prepare_metastore("cold_cache_fresh").unwrap();
-    populate_metastore(&fresh_metastore, num_schemas, tables_per_schema).await.unwrap();
+    populate_metastore(&fresh_metastore, num_schemas, tables_per_schema)
+        .await
+        .unwrap();
     let result = fresh_metastore.get_tables_with_path(false).await;
     assert!(result.is_ok());
 }
@@ -157,12 +162,15 @@ fn do_cold_vs_warm_cache_bench(
 ) {
     let metastore = runtime.block_on(async {
         let metastore = prepare_metastore("cold_warm_cache").unwrap();
-        populate_metastore(&metastore, num_schemas, tables_per_schema).await.unwrap();
+        populate_metastore(&metastore, num_schemas, tables_per_schema)
+            .await
+            .unwrap();
         metastore
     });
 
     c.bench_function("get_tables_with_path_cold_cache", |b| {
-        b.to_async(runtime).iter(|| do_cold_cache_test(num_schemas, tables_per_schema));
+        b.to_async(runtime)
+            .iter(|| do_cold_cache_test(num_schemas, tables_per_schema));
     });
 
     c.bench_function("get_tables_with_path_warm_cache", |b| {
@@ -177,7 +185,7 @@ fn do_benches(c: &mut Criterion) {
     do_get_tables_with_path_bench(c, &runtime, 10, 10, 100);
     do_get_tables_with_path_bench(c, &runtime, 50, 20, 50);
     do_get_tables_with_path_bench(c, &runtime, 25, 1000, 10);
-    
+
     do_cold_vs_warm_cache_bench(c, &runtime, 20, 50);
 }
 

rust/cubestore/cubestore/benches/tracking_allocator.rs

@@ -33,7 +33,6 @@ impl TrackingAllocator {
         self.total_allocated.store(0, Ordering::Relaxed);
     }
 
-
     pub fn print_stats(&self) {
         let allocations = self.allocations.load(Ordering::Relaxed);
         let deallocations = self.deallocations.load(Ordering::Relaxed);
@@ -46,9 +45,21 @@ impl TrackingAllocator {
         println!("Total allocations: {}", allocations);
         println!("Total deallocations: {}", deallocations);
         println!("Total reallocations: {}", reallocations);
-        println!("Current allocated: {} bytes ({:.2} MB)", current_allocated, current_allocated as f64 / 1024.0 / 1024.0);
-        println!("Peak allocated: {} bytes ({:.2} MB)", peak_allocated, peak_allocated as f64 / 1024.0 / 1024.0);
-        println!("Total allocated: {} bytes ({:.2} MB)", total_allocated, total_allocated as f64 / 1024.0 / 1024.0);
+        println!(
+            "Current allocated: {} bytes ({:.2} MB)",
+            current_allocated,
+            current_allocated as f64 / 1024.0 / 1024.0
+        );
+        println!(
+            "Peak allocated: {} bytes ({:.2} MB)",
+            peak_allocated,
+            peak_allocated as f64 / 1024.0 / 1024.0
+        );
+        println!(
+            "Total allocated: {} bytes ({:.2} MB)",
+            total_allocated,
+            total_allocated as f64 / 1024.0 / 1024.0
+        );
         println!("===============================");
     }
 
@@ -57,7 +68,7 @@ impl TrackingAllocator {
             self.allocations.fetch_add(1, Ordering::Relaxed);
             self.total_allocated.fetch_add(size, Ordering::Relaxed);
             let current = self.current_allocated.fetch_add(size, Ordering::Relaxed) + size;
-            
+
             // Update peak if current exceeds it
             let mut peak = self.peak_allocated.load(Ordering::Relaxed);
             while current > peak {
@@ -96,12 +107,16 @@ unsafe impl GlobalAlloc for TrackingAllocator {
         let new_ptr = self.inner.realloc(ptr, layout, new_size);
         if !new_ptr.is_null() {
             self.reallocations.fetch_add(1, Ordering::Relaxed);
-            
+
             // Update counters: subtract old size, add new size
-            self.current_allocated.fetch_sub(layout.size(), Ordering::Relaxed);
+            self.current_allocated
+                .fetch_sub(layout.size(), Ordering::Relaxed);
             self.total_allocated.fetch_add(new_size, Ordering::Relaxed);
-            let current = self.current_allocated.fetch_add(new_size, Ordering::Relaxed) + new_size;
-            
+            let current = self
+                .current_allocated
+                .fetch_add(new_size, Ordering::Relaxed)
+                + new_size;
+
             // Update peak if current exceeds it
             let mut peak = self.peak_allocated.load(Ordering::Relaxed);
             while current > peak {
@@ -118,4 +133,4 @@ unsafe impl GlobalAlloc for TrackingAllocator {
         }
         new_ptr
     }
-}
+}