diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index 64252f772..2fcffb297 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -111,6 +111,8 @@ sha2.workspace = true [dev-dependencies] anyhow.workspace = true +# `test-util` adds the per-request records that the `assert_io_*` macros print. +lance-io = { workspace = true, features = ["test-util"] } lance-testing = { workspace = true } # RestAdapter is only used by cfg(test) code. Keep rest-adapter off the # production remote feature so Python wheels do not resolve its Axum 0.7 stack. diff --git a/rust/lancedb/src/io/object_store.rs b/rust/lancedb/src/io/object_store.rs index c4a9a4f7e..a40694463 100644 --- a/rust/lancedb/src/io/object_store.rs +++ b/rust/lancedb/src/io/object_store.rs @@ -15,9 +15,6 @@ use object_store::{ use async_trait::async_trait; -#[cfg(test)] -pub mod io_tracking; - #[derive(Debug)] struct MirroringObjectStore { primary: Arc, diff --git a/rust/lancedb/src/io/object_store/io_tracking.rs b/rust/lancedb/src/io/object_store/io_tracking.rs deleted file mode 100644 index 7f9750216..000000000 --- a/rust/lancedb/src/io/object_store/io_tracking.rs +++ /dev/null @@ -1,250 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The LanceDB Authors - -use std::{ - fmt::{Display, Formatter}, - sync::{Arc, Mutex}, -}; - -use bytes::Bytes; -use futures::stream::BoxStream; -use lance::io::WrappingObjectStore; -use object_store::{ - CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore, - PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, Result as OSResult, - UploadPart, list::PaginatedListStore, path::Path, -}; - -#[derive(Debug, Default)] -pub struct IoStats { - pub read_iops: u64, - pub read_bytes: u64, - pub write_iops: u64, - pub write_bytes: u64, -} - -impl Display for IoStats { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "{:#?}", self) - } -} - -#[derive(Debug, Clone)] -pub struct IoTrackingStore { - target: Arc, - stats: Arc>, -} - -impl Display for IoTrackingStore { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "{:#?}", self) - } -} - -#[derive(Debug, Default, Clone)] -pub struct IoStatsHolder(Arc>); - -impl IoStatsHolder { - pub fn incremental_stats(&self) -> IoStats { - std::mem::take(&mut self.0.lock().expect("failed to lock IoStats")) - } -} - -impl WrappingObjectStore for IoStatsHolder { - fn wrap(&self, _store_prefix: &str, target: Arc) -> Arc { - Arc::new(IoTrackingStore { - target, - stats: self.0.clone(), - }) - } - - fn wrap_paginated( - &self, - _store_prefix: &str, - original: Arc, - ) -> Option> { - Some(original) - } -} - -impl IoTrackingStore { - pub fn new_wrapper() -> (Arc, Arc>) { - let stats = Arc::new(Mutex::new(IoStats::default())); - (Arc::new(IoStatsHolder(stats.clone())), stats) - } - - fn record_read(&self, num_bytes: u64) { - let mut stats = self.stats.lock().unwrap_or_else(|e| e.into_inner()); - stats.read_iops += 1; - stats.read_bytes += num_bytes; - } - - fn record_write(&self, num_bytes: u64) { - let mut stats = self.stats.lock().unwrap_or_else(|e| e.into_inner()); - stats.write_iops += 1; - stats.write_bytes += num_bytes; - } -} - -#[async_trait::async_trait] -#[deny(clippy::missing_trait_methods)] -impl ObjectStore for IoTrackingStore { - async fn put_opts( - &self, - location: &Path, - bytes: PutPayload, - opts: PutOptions, - ) -> OSResult { - self.record_write(bytes.content_length() as u64); - self.target.put_opts(location, bytes, opts).await - } - - async fn put_multipart_opts( - &self, - location: &Path, - opts: PutMultipartOptions, - ) -> OSResult> { - let target = self.target.put_multipart_opts(location, opts).await?; - Ok(Box::new(IoTrackingMultipartUpload { - target, - stats: self.stats.clone(), - })) - } - - async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult { - let result = self.target.get_opts(location, options).await; - if let Ok(result) = &result { - let num_bytes = result.range.end - result.range.start; - self.record_read(num_bytes); - } - result - } - - async fn get_ranges( - &self, - location: &Path, - ranges: &[std::ops::Range], - ) -> OSResult> { - let result = self.target.get_ranges(location, ranges).await; - if let Ok(result) = &result { - self.record_read(result.iter().map(|b| b.len() as u64).sum()); - } - result - } - - fn delete_stream( - &self, - locations: BoxStream<'static, OSResult>, - ) -> BoxStream<'static, OSResult> { - self.record_write(0); - self.target.delete_stream(locations) - } - - fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult> { - self.record_read(0); - self.target.list(prefix) - } - - fn list_with_offset( - &self, - prefix: Option<&Path>, - offset: &Path, - ) -> BoxStream<'static, OSResult> { - self.record_read(0); - self.target.list_with_offset(prefix, offset) - } - - async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult { - self.record_read(0); - self.target.list_with_delimiter(prefix).await - } - - async fn copy_opts(&self, from: &Path, to: &Path, options: CopyOptions) -> OSResult<()> { - self.record_write(0); - self.target.copy_opts(from, to, options).await - } - - async fn rename_opts(&self, from: &Path, to: &Path, options: RenameOptions) -> OSResult<()> { - self.record_write(0); - self.target.rename_opts(from, to, options).await - } -} - -#[derive(Debug)] -struct IoTrackingMultipartUpload { - target: Box, - stats: Arc>, -} - -#[async_trait::async_trait] -impl MultipartUpload for IoTrackingMultipartUpload { - async fn abort(&mut self) -> OSResult<()> { - self.target.abort().await - } - - async fn complete(&mut self) -> OSResult { - self.target.complete().await - } - - fn put_part(&mut self, payload: PutPayload) -> UploadPart { - { - let mut stats = self.stats.lock().unwrap_or_else(|e| e.into_inner()); - stats.write_iops += 1; - stats.write_bytes += payload.content_length() as u64; - } - self.target.put_part(payload) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - /// Helper: poison a Mutex by panicking while holding the lock. - fn poison_stats(stats: &Arc>) { - let stats_clone = stats.clone(); - let handle = std::thread::spawn(move || { - let _guard = stats_clone.lock().unwrap(); - panic!("intentional panic to poison stats mutex"); - }); - let _ = handle.join(); - assert!(stats.lock().is_err(), "mutex should be poisoned"); - } - - #[test] - fn test_record_read_recovers_from_poisoned_lock() { - let stats = Arc::new(Mutex::new(IoStats::default())); - let store = IoTrackingStore { - target: Arc::new(object_store::memory::InMemory::new()), - stats: stats.clone(), - }; - - poison_stats(&stats); - - // record_read should not panic - store.record_read(1024); - - // Verify the stats were updated despite poisoning - let s = stats.lock().unwrap_or_else(|e| e.into_inner()); - assert_eq!(s.read_iops, 1); - assert_eq!(s.read_bytes, 1024); - } - - #[test] - fn test_record_write_recovers_from_poisoned_lock() { - let stats = Arc::new(Mutex::new(IoStats::default())); - let store = IoTrackingStore { - target: Arc::new(object_store::memory::InMemory::new()), - stats: stats.clone(), - }; - - poison_stats(&stats); - - // record_write should not panic - store.record_write(2048); - - let s = stats.lock().unwrap_or_else(|e| e.into_inner()); - assert_eq!(s.write_iops, 1); - assert_eq!(s.write_bytes, 2048); - } -} diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 173ddc815..6f1c92147 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -4267,9 +4267,10 @@ mod tests { use tempfile::tempdir; use super::*; + use lance_io::assert_io_lt; + use crate::connect; use crate::connection::ConnectBuilder; - use crate::io::object_store::io_tracking::IoTrackingStore; use crate::query::Select; use crate::query::{ExecutableQuery, QueryBase}; use crate::test_utils::connection::new_test_connection; @@ -6301,34 +6302,31 @@ mod tests { table.add(batch.clone()).execute().await.unwrap(); } - // Reopen through a tracking store so the counters cover `stats()` alone and - // not the writes above. - let (wrapper, io_stats) = IoTrackingStore::new_wrapper(); - let table = conn - .open_table("test_stats_io") - .lance_read_params(ReadParams { - store_options: Some(ObjectStoreParams { - object_store_wrapper: Some(wrapper), - ..Default::default() - }), - ..Default::default() - }) - .execute() - .await - .unwrap(); - io_stats.lock().unwrap().read_iops = 0; + // Reopen, then count against the store's own tracker so the counters cover + // `stats()` alone and not the writes above. + let table = conn.open_table("test_stats_io").execute().await.unwrap(); + let io_tracker = { + let dataset = table.as_native().unwrap().dataset.get().await.unwrap(); + dataset + .object_store(None) + .await + .unwrap() + .io_tracker() + .clone() + }; + io_tracker.incremental_stats(); let stats = table.stats().await.unwrap(); - let read_iops = io_stats.lock().unwrap().read_iops; + let io_stats = io_tracker.incremental_stats(); assert_eq!(stats.fragment_stats.num_fragments, NUM_APPENDS + 1); assert!(stats.total_bytes > 0); // Reading the fragments' data files would take at least one IOP each. - assert!( - read_iops < stats.fragment_stats.num_fragments as u64, - "stats() issued {} read IOPs across {} fragments", + assert_io_lt!( + io_stats, read_iops, - stats.fragment_stats.num_fragments + stats.fragment_stats.num_fragments as u64, + "stats() should not read the fragments' data files" ); } diff --git a/rust/lancedb/src/table/dataset.rs b/rust/lancedb/src/table/dataset.rs index 69d02776f..2739522bc 100644 --- a/rust/lancedb/src/table/dataset.rs +++ b/rust/lancedb/src/table/dataset.rs @@ -373,10 +373,10 @@ mod tests { use super::*; - use crate::{ - connect, io::object_store::io_tracking::IoStatsHolder, table::WriteOptions, - utils::background_cache::clock, - }; + use lance_io::assert_io_eq; + use lance_io::utils::tracking_store::IOTracker; + + use crate::{connect, table::WriteOptions, utils::background_cache::clock}; async fn create_test_dataset(uri: &str) -> Dataset { let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); @@ -626,7 +626,7 @@ mod tests { .execute() .await .expect("Failed to connect to database"); - let io_stats = IoStatsHolder::default(); + let io_stats = IOTracker::default(); let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); @@ -651,7 +651,7 @@ mod tests { // latest version. table.schema().await.unwrap(); let stats = io_stats.incremental_stats(); - assert_eq!(stats.read_iops, 1); + assert_io_eq!(stats, read_iops, 1); } /// Regression test: a write that races with as_time_travel() must not panic. @@ -689,7 +689,7 @@ mod tests { .execute() .await .unwrap(); - let io_stats = IoStatsHolder::default(); + let io_stats = IOTracker::default(); let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); let table = db .create_empty_table("test", schema) @@ -712,12 +712,12 @@ mod tests { // Step 1: within interval — no list table.schema().await.unwrap(); let s = io_stats.incremental_stats(); - assert_eq!(s.read_iops, 0, "step 1, elapsed={:?}", start.elapsed()); + assert_io_eq!(s, read_iops, 0, "step 1, elapsed={:?}", start.elapsed()); // Step 2: still within interval — no list table.schema().await.unwrap(); let s = io_stats.incremental_stats(); - assert_eq!(s.read_iops, 0, "step 2, elapsed={:?}", start.elapsed()); + assert_io_eq!(s, read_iops, 0, "step 2, elapsed={:?}", start.elapsed()); // Step 3: sleep past the 1s boundary tokio::time::sleep(Duration::from_secs(1)).await; @@ -725,14 +725,14 @@ mod tests { // Step 4: interval expired — exactly 1 list, timer resets table.schema().await.unwrap(); let s = io_stats.incremental_stats(); - assert_eq!(s.read_iops, 1, "step 4, elapsed={:?}", start.elapsed()); + assert_io_eq!(s, read_iops, 1, "step 4, elapsed={:?}", start.elapsed()); // Step 5: 10 more calls — timer just reset, no lists (THIS is the regression test). for _ in 0..10 { table.schema().await.unwrap(); } let s = io_stats.incremental_stats(); - assert_eq!(s.read_iops, 0, "step 5, elapsed={:?}", start.elapsed()); + assert_io_eq!(s, read_iops, 0, "step 5, elapsed={:?}", start.elapsed()); } /// Helper: poison the mutex inside a DatasetConsistencyWrapper.