refactor: drop duplicated IoTrackingStore for lance-io's

`rust/lancedb/src/io/object_store/io_tracking.rs` was a near-verbatim copy of lance-io's `utils::tracking_store`, and had drifted behind it. It was `#[cfg(test)]`-only with two usages, both in tests, so this PR points those at `lance_io::utils::tracking_store::IOTracker` and deletes the copy.

lance-io's `test-util` feature is enabled as a dev-dependency, which brings the `assert_io_eq!` / `assert_io_lt!` macros. Those print the recorded requests (method, path, byte range) when an assertion fails, instead of just the counter that didn't match, so the IO-counting tests in `table/dataset.rs` now say *which* IO happened when they fail.

`test_stats_does_not_read_data_files` no longer plumbs an `object_store_wrapper` at all: every lance-io `ObjectStore` now carries its own tracker, reachable as `store.io_tracker()`. This also makes the test stronger — the optimized local reader bypasses the `object_store` layer, so reads that would have gone uncounted through a wrapper are recorded by the built-in tracker.

Fixes #4076

## Not included

Two lance-side items from the issue are left for follow-up there, since neither can be fixed from this repo:

- `IOTracker::wrap_paginated` returns `Some(original)` on the premise that a pushed-down listing counts itself, which only holds for the local lister — a paginated listing against S3 is invisible to the tracker. The deleted copy had the identical bug, so this is not a regression.
- The deleted copy had poison-tolerant locking (`unwrap_or_else(|e| e.into_inner())`, from #3196) where lance still `.unwrap()`s. That tolerance, and the two unit tests covering it, go away with the file; it is worth carrying upstream instead of re-forking here.

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.

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<dyn ObjectStore>,

rust/lancedb/src/io/object_store/io_tracking.rs

@@ -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<dyn ObjectStore>,
-    stats: Arc<Mutex<IoStats>>,
-}
-
-impl Display for IoTrackingStore {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{:#?}", self)
-    }
-}
-
-#[derive(Debug, Default, Clone)]
-pub struct IoStatsHolder(Arc<Mutex<IoStats>>);
-
-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<dyn ObjectStore>) -> Arc<dyn ObjectStore> {
-        Arc::new(IoTrackingStore {
-            target,
-            stats: self.0.clone(),
-        })
-    }
-
-    fn wrap_paginated(
-        &self,
-        _store_prefix: &str,
-        original: Arc<dyn PaginatedListStore>,
-    ) -> Option<Arc<dyn PaginatedListStore>> {
-        Some(original)
-    }
-}
-
-impl IoTrackingStore {
-    pub fn new_wrapper() -> (Arc<dyn WrappingObjectStore>, Arc<Mutex<IoStats>>) {
-        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<PutResult> {
-        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<Box<dyn MultipartUpload>> {
-        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<GetResult> {
-        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<u64>],
-    ) -> OSResult<Vec<Bytes>> {
-        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<Path>>,
-    ) -> BoxStream<'static, OSResult<Path>> {
-        self.record_write(0);
-        self.target.delete_stream(locations)
-    }
-
-    fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult<ObjectMeta>> {
-        self.record_read(0);
-        self.target.list(prefix)
-    }
-
-    fn list_with_offset(
-        &self,
-        prefix: Option<&Path>,
-        offset: &Path,
-    ) -> BoxStream<'static, OSResult<ObjectMeta>> {
-        self.record_read(0);
-        self.target.list_with_offset(prefix, offset)
-    }
-
-    async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult<ListResult> {
-        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<dyn MultipartUpload>,
-    stats: Arc<Mutex<IoStats>>,
-}
-
-#[async_trait::async_trait]
-impl MultipartUpload for IoTrackingMultipartUpload {
-    async fn abort(&mut self) -> OSResult<()> {
-        self.target.abort().await
-    }
-
-    async fn complete(&mut self) -> OSResult<PutResult> {
-        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<IoStats> by panicking while holding the lock.
-    fn poison_stats(stats: &Arc<Mutex<IoStats>>) {
-        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);
-    }
-}

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"
         );
     }
 

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.