simplify

Author: wjones127

Cargo.lock

@@ -16,6 +16,7 @@ arrow-schema = { workspace = true }
 rand = { workspace = true }
 num-traits = { workspace = true }
 lance-arrow = { workspace = true }
+tokio = { workspace = true, features = ["sync"] }
 
 [lints]
 workspace = true

rust/lance-testing/Cargo.toml

@@ -4,12 +4,10 @@
 //! Simple memory allocation tracking for tests
 //!
 //! This module provides a global allocator wrapper that tracks memory usage statistics.
-//! Unlike the previous tracking-allocator implementation, this does not rely on tracing
-//! spans, ensuring all allocations are captured regardless of span propagation.
 
 use std::alloc::{GlobalAlloc, Layout};
+use std::future::Future;
 use std::sync::atomic::{AtomicIsize, AtomicUsize, Ordering};
-use std::sync::Mutex;
 
 /// Statistics about memory allocations
 #[derive(Debug, Clone, Copy, Default)]
@@ -143,6 +141,82 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
     }
 }
 
+/// Record allocation statistics while executing a function
+///
+/// This function automatically:
+/// - Acquires a lock to prevent concurrent tests from interfering
+/// - Resets allocation counters before execution
+/// - Executes the provided closure
+/// - Returns the allocation statistics
+///
+/// # Example
+///
+/// ```rust
+/// use lance_testing::allocator::record_allocations;
+///
+/// let stats = record_allocations(|| {
+///     // Allocate some memory
+///     let v: Vec<u8> = vec![0; 1024 * 1024]; // 1MB
+///     drop(v);
+/// });
+///
+/// assert!(stats.total_bytes_allocated >= 1024 * 1024);
+/// ```
+pub fn record_allocations<F, R>(f: F) -> AllocStats
+where
+    F: FnOnce() -> R,
+{
+    // Acquire lock to prevent concurrent tests from interfering
+    let _guard = ALLOC_TEST_MUTEX.blocking_lock();
+
+    // Reset stats before execution
+    GLOBAL_STATS.reset();
+
+    // Execute the closure
+    let _ = f();
+
+    // Return the stats
+    GLOBAL_STATS.get_stats()
+}
+
+/// Record allocation statistics while executing a future
+///
+/// This function automatically:
+/// - Acquires a lock to prevent concurrent tests from interfering
+/// - Resets allocation counters before execution
+/// - Executes the provided closure
+/// - Returns the allocation statistics
+///
+/// # Example
+///
+/// ```rust
+/// use lance_testing::allocator::record_allocations_async;
+///
+/// let stats = record_allocations_async(async {
+///     // Allocate some memory
+///     let v: Vec<u8> = vec![0; 1024 * 1024]; // 1MB
+///     drop(v);
+/// });
+///
+/// assert!(stats.total_bytes_allocated >= 1024 * 1024);
+/// ```
+pub async fn record_allocations_async<F, R>(f: F) -> AllocStats
+where
+    F: Future<Output = R>,
+{
+    // Acquire lock to prevent concurrent tests from interfering
+    let _guard = ALLOC_TEST_MUTEX.lock().await;
+
+    // Reset stats before execution
+    GLOBAL_STATS.reset();
+
+    // Execute the closure
+    let _ = f.await;
+
+    // Return the stats
+    GLOBAL_STATS.get_stats()
+}
+
 /// Get current allocation statistics
 ///
 /// Returns statistics about memory allocations made by Lance's Rust code.
@@ -153,24 +227,22 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAllocator<A> {
 /// - Allocations made by other native libraries (e.g., PyArrow, NumPy)
 /// - Memory-mapped files or shared memory
 ///
-/// For comprehensive memory tracking in Python tests, use the
-/// `lance.testing.track_memory()` context manager which combines these stats
-/// with `tracemalloc` and PyArrow's memory pool statistics.
+/// **Note**: For tests, prefer using `record_allocations()` which automatically
+/// handles locking and reset.
 pub fn get_alloc_stats() -> AllocStats {
     GLOBAL_STATS.get_stats()
 }
 
 /// Reset allocation statistics to zero
 ///
-/// This resets all counters. Note that for accurate tracking in tests,
-/// you should use the `#[track_alloc]` macro which ensures single-threaded
-/// execution and automatic reset/collection.
+/// **Note**: For tests, prefer using `record_allocations()` which automatically
+/// handles locking and reset.
 pub fn reset_alloc_stats() {
     GLOBAL_STATS.reset();
 }
 
 /// Mutex to ensure single-threaded execution of tests that track allocations
 ///
 /// This prevents stats from different tests running in parallel from mixing.
-/// The `#[track_alloc]` macro uses this automatically.
-pub static ALLOC_TEST_MUTEX: Mutex<()> = Mutex::new(());
+/// The `record_allocations()` function uses this automatically.
+static ALLOC_TEST_MUTEX: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());

rust/lance-testing/src/allocator.rs

@@ -445,8 +445,11 @@ struct WriteContext<'a> {
 
 #[cfg(test)]
 mod test {
+    use all_asserts::assert_gt;
     use arrow_array::StructArray;
     use arrow_schema::{DataType, Field, Schema};
+    use lance_datafusion::datagen::DatafusionDatagenExt;
+    use lance_datagen::{array, gen_batch, BatchCount, ByteCount, RoundingBehavior};
 
     use crate::session::Session;
 
@@ -502,57 +505,43 @@ mod test {
     async fn test_insert_memory_tracking() {
         // Test that insert doesn't load all data into memory at once
         // when writing 100MB of data in 10MB batches
-        use arrow_array::Int32Array;
-        use arrow_schema::ArrowError;
-
-        // Create 100MB of data in 10 batches of 10MB each
-        let batch_size_mb = 10;
-        let num_batches = 10;
-        let total_rows_per_batch = (batch_size_mb * 1024 * 1024) / (4 * 3); // Int32 = 4 bytes, 3 columns
-
-        let schema = Arc::new(Schema::new(vec![
-            Field::new("a", DataType::Int32, false),
-            Field::new("b", DataType::Int32, false),
-            Field::new("c", DataType::Int32, false),
-        ]));
-
-        // Create batches on demand
-        let batches: Vec<std::result::Result<RecordBatch, ArrowError>> = (0..num_batches)
-            .map(|i| {
-                let start = i * total_rows_per_batch;
-                let end = (i + 1) * total_rows_per_batch;
-                let a = Int32Array::from_iter_values(start..end);
-                let b = Int32Array::from_iter_values(start..end);
-                let c = Int32Array::from_iter_values(start..end);
-                RecordBatch::try_new(schema.clone(), vec![Arc::new(a), Arc::new(b), Arc::new(c)])
-            })
-            .collect();
-
-        lance_testing::allocator::reset_alloc_stats();
+        let batch_size = 10 * 1024 * 1024; // 10MB
+
+        let stats = lance_testing::allocator::record_allocations_async(async {
+            let tmp_dir = tempfile::tempdir().unwrap();
+            let tmp_path = tmp_dir.path().to_str().unwrap().to_string();
+            // Create batches on demand
+            // Create a stream of 100MB of data, in batches
+            let num_batches = BatchCount::from(10);
+            let data = gen_batch()
+                .col("a", array::rand_type(&DataType::Int32))
+                .into_df_stream_bytes(
+                    ByteCount::from(batch_size),
+                    num_batches,
+                    RoundingBehavior::RoundDown,
+                )
+                .unwrap();
 
-        let _dataset = InsertBuilder::new("memory://")
-            .execute_stream(RecordBatchIterator::new(
-                batches.into_iter(),
-                schema.clone(),
-            ))
-            .await
-            .unwrap();
+            let _dataset = InsertBuilder::new(&tmp_path)
+                .execute_stream(data)
+                .await
+                .unwrap();
+        })
+        .await;
 
-        let stats = lance_testing::allocator::get_alloc_stats();
+        assert_gt!(stats.total_bytes_allocated, 100 * 1024 * 1024);
 
         // The key test: we shouldn't load all 100MB at once
-        // Allow 5x the batch size to account for overhead and buffering
-        // (this would still catch if we buffered all 100MB)
-        let max_allowed_mb = batch_size_mb * 5;
-        let max_allowed_bytes = (max_allowed_mb * 1024 * 1024) as isize;
+        // Allow 2x the batch size to account for overhead and buffering
+        let max_allowed_bytes = (batch_size * 2) as isize;
 
         let peak_mb = stats.max_bytes_allocated as f64 / (1024.0 * 1024.0);
 
         assert!(
             stats.max_bytes_allocated <= max_allowed_bytes,
-            "Peak memory {} MB exceeded limit {} MB",
+            "Peak memory {:.2} MB exceeded limit {:.2} MB",
             peak_mb,
-            max_allowed_mb
+            max_allowed_bytes as f64 / (1024.0 * 1024.0)
         );
     }
 }

rust/lance/src/dataset/write/insert.rs

@@ -108,3 +108,13 @@ pub mod deps {
     pub use arrow_schema;
     pub use datafusion;
 }
+
+#[cfg(test)]
+mod tests {
+    use lance_testing::allocator::TrackingAllocator;
+    use std::alloc::System;
+
+    // Enable allocation statistics for memory tracking tests.
+    #[global_allocator]
+    static GLOBAL: TrackingAllocator<System> = TrackingAllocator::new(System);
+}