fix(io): publish object store metrics for local reads and writes - #7994
Conversation
📝 WalkthroughWalkthroughLocal filesystem readers, writers, io_uring paths, and copy/delete shortcuts now record object-store metrics through ChangesLocal I/O metrics
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant ObjectStore
participant LocalObjectReader
participant LocalFilesystem
participant IOTracker
ObjectStore->>LocalObjectReader: open with local_io_tracker()
LocalObjectReader->>IOTracker: begin_io("get")
LocalObjectReader->>LocalFilesystem: perform blocking read
LocalFilesystem-->>LocalObjectReader: return Result and bytes
LocalObjectReader->>IOTracker: record outcome and byte count
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
The optimized local paths — `LocalWriter`, `LocalObjectReader`, the io_uring readers, and the local `copy` / recursive-delete shortcuts — go straight to the filesystem, so they never reach `MeteredObjectStore` and published no metrics at all. Writing a dataset to a local path produced zero object store metrics even though the IO tracker recorded it. `IOTracker` now carries the metrics `base` label of the store it belongs to and hands out an `IoMetricsGuard` for IO that bypasses the `object_store` layer. Every such path records requests, bytes, latency, errors and the in-flight gauge under the same labels the metered store uses, so local IO aggregates with cloud IO. A caller-supplied `ObjectStore` (the deprecated `ObjectStoreParams::object_store`) is now metered too. `IoStats` is unchanged: the local `copy`, recursive delete and size lookup were never counted there and still aren't, so existing IO assertions are unaffected. Fixes lance-format#7993 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
abda9ee to
4b64b21
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Note
Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.
🟡 Other comments (1)
rust/lance-io/src/object_store/metrics.rs-1517-1536 (1)
1517-1536: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAssert on the error variant/message, not just
is_err().
test_local_read_error_is_countedonly checksreader.get_range(0..100).await.is_err()). As per coding guidelines: "Assert on both the error variant and the message content in tests; do not check onlyis_err()."✅ Proposed fix
let reader = store.open(&path).await.unwrap(); // Reading past the end of the file fails. - assert!(reader.get_range(0..100).await.is_err()); + let err = reader.get_range(0..100).await.unwrap_err(); + assert!( + matches!(err, object_store::Error::Generic { .. }), + "unexpected error variant: {err:?}" + );🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance-io/src/object_store/metrics.rs` around lines 1517 - 1536, Update test_local_read_error_is_counted to capture the get_range error and assert both its expected error variant and message content instead of checking only is_err(). Preserve the existing metrics assertions and the out-of-bounds read scenario.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rust/lance-io/src/utils/tracking_store.rs`:
- Around line 104-144: The local-bypass metrics label must honor BaseLabelMode
scoping and avoid per-operation allocation. In
rust/lance-io/src/utils/tracking_store.rs lines 104-144, update
with_metrics_base to materialize the same scoped label used by
ObjectStoreMetricsExt::metered(), including the configured off/full/scheme
behavior. In rust/lance-io/src/object_store.rs lines 653-661, precompute and
cache the scoped local IOTracker alongside the existing io_tracker/store_prefix
state, then reuse it across open, open_with_size, create, copy_impl, and
remove_dir_all instead of calling local_io_tracker() for each operation.
---
Other comments:
In `@rust/lance-io/src/object_store/metrics.rs`:
- Around line 1517-1536: Update test_local_read_error_is_counted to capture the
get_range error and assert both its expected error variant and message content
instead of checking only is_err(). Preserve the existing metrics assertions and
the out-of-bounds read scenario.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: 7e76efbb-517b-4558-9642-4ee3f9999754
📒 Files selected for processing (8)
rust/lance-io/src/local.rsrust/lance-io/src/object_reader.rsrust/lance-io/src/object_store.rsrust/lance-io/src/object_store/metrics.rsrust/lance-io/src/object_writer.rsrust/lance-io/src/uring/current_thread.rsrust/lance-io/src/uring/reader.rsrust/lance-io/src/utils/tracking_store.rs
|
|
||
| /// Label the metrics published through [`Self::begin_io`] with the prefix of | ||
| /// the store this tracker belongs to, so IO that bypasses the `object_store` | ||
| /// layer carries the same `base` label as the store's metered operations. | ||
| #[cfg(feature = "metrics")] | ||
| pub fn with_metrics_base(mut self, base: &str) -> Self { | ||
| self.metrics_base = Some(base.into()); | ||
| self | ||
| } | ||
|
|
||
| /// Without the `metrics` feature there is nothing to label. | ||
| #[cfg(not(feature = "metrics"))] | ||
| pub fn with_metrics_base(self, _base: &str) -> Self { | ||
| self | ||
| } | ||
|
|
||
| /// Begin an operation that talks to storage without going through the | ||
| /// `object_store` layer, and so is invisible to the `MeteredObjectStore` | ||
| /// wrapper: the optimized local reads and writes go straight to the | ||
| /// filesystem. `operation` must be one of the labels that wrapper uses | ||
| /// (`get`, `put`, `head`, ...) so this IO aggregates with the rest. | ||
| /// | ||
| /// The returned guard keeps the in-flight gauge raised until it is dropped. | ||
| #[cfg(feature = "metrics")] | ||
| pub fn begin_io(&self, operation: &'static str) -> IoMetricsGuard { | ||
| IoMetricsGuard { | ||
| state: self.metrics_base.as_ref().map(|base| IoMetricsState { | ||
| _in_flight: InFlightGuard::new(base, operation), | ||
| base: base.clone(), | ||
| operation, | ||
| start: Instant::now(), | ||
| }), | ||
| } | ||
| } | ||
|
|
||
| /// Without the `metrics` feature there is nothing to publish. | ||
| #[cfg(not(feature = "metrics"))] | ||
| pub fn begin_io(&self, _operation: &'static str) -> IoMetricsGuard { | ||
| IoMetricsGuard {} | ||
| } | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Local-bypass metrics label ignores BaseLabelMode scoping and is re-allocated on every local I/O call.
with_metrics_base stores the raw, unscoped store_prefix as-is, unlike ObjectStoreMetricsExt::metered() which applies scoped_base(parse_base_label_mode(...), prefix) (see metrics.rs test_scoped_base/test_base_label_defaults_to_scheme). Because local_io_tracker() calls with_metrics_base afresh on every local read/write/copy/delete, this also means: (1) the runtime LANCE_METRICS_BASE_LABEL_MODE (off/full/scheme) opt-out isn't honored for local bypass paths the way it is for cloud/MeteredObjectStore paths, and (2) a new Arc<str> is allocated on every single local I/O call instead of once per ObjectStore.
rust/lance-io/src/utils/tracking_store.rs#L104-L144: apply the sameBaseLabelMode/scoped_baselogic used by.metered()when materializingmetrics_base(or accept an already-scoped label from the caller) instead of storing the raw prefix.rust/lance-io/src/object_store.rs#L653-L661: precompute and cache the scoped localIOTrackeronce (e.g., alongside theio_tracker/store_prefixfields at construction) instead of callinglocal_io_tracker()— and re-allocating — on everyopen/open_with_size/create/copy_impl/remove_dir_allcall.
📍 Affects 2 files
rust/lance-io/src/utils/tracking_store.rs#L104-L144(this comment)rust/lance-io/src/object_store.rs#L653-L661
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/lance-io/src/utils/tracking_store.rs` around lines 104 - 144, The
local-bypass metrics label must honor BaseLabelMode scoping and avoid
per-operation allocation. In rust/lance-io/src/utils/tracking_store.rs lines
104-144, update with_metrics_base to materialize the same scoped label used by
ObjectStoreMetricsExt::metered(), including the configured off/full/scheme
behavior. In rust/lance-io/src/object_store.rs lines 653-661, precompute and
cache the scoped local IOTracker alongside the existing io_tracker/store_prefix
state, then reuse it across open, open_with_size, create, copy_impl, and
remove_dir_all instead of calling local_io_tracker() for each operation.
There was a problem hiding this comment.
The BaseLabelMode half of this is a misread. Scoping isn't applied when the base is stored, it's applied when the labels are built: operation_labels() calls scoped_base(base_label_mode(), base), and both record_outcome and InFlightGuard::new go through it. MeteredObjectStore also holds the raw, unscoped prefix in its base field, so passing the raw store_prefix here was already consistent with the cloud path and off/full/scheme were honored identically on both. (The env var is LANCE_OBJECT_STORE_METRICS_LABEL, not LANCE_METRICS_BASE_LABEL_MODE.)
The per-call allocation was real, and it's fixed in aa0642c as a side effect of a larger correctness fix — see the thread on local_io_tracker(). The label is now materialized once per store at construction, and local_io_tracker() is gone.
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
hamersaw
left a comment
There was a problem hiding this comment.
Overall this makes sense, just a few clarifications to make sure it's doing exactly what we think it's doing.
| source: err.into(), | ||
| }); | ||
|
|
||
| metrics.record(&result, num_bytes); |
There was a problem hiding this comment.
JoinError on the await? mean that this metrics.record is never reached. A very uncommon corner case, but we could replace await.? with a match or .map_err(Into::into).and_then(...)... so that the metrics are still reported on JoinError.
There was a problem hiding this comment.
Good catch, fixed in aa0642c. There was a third site with the same problem: stream_local_range in object_reader.rs also did .await? before metrics.record.
Rather than three spellings of the same await, all of them now go through a shared join_local_io helper that flattens the JoinError and the IO error into one object_store::Error — which is what size() was already doing inline, so that one is folded in too.
No test for the JoinError branch: triggering it means making the blocking read panic, which I couldn't do deterministically without adding an injection seam just for the test. Same reason the io_uring readers are uncovered here.
| fn local_io_tracker(&self) -> IOTracker { | ||
| self.io_tracker | ||
| .clone() | ||
| .with_metrics_base(&self.store_prefix) | ||
| } |
There was a problem hiding this comment.
Note: Below is a comment from AI review. Honestly, it's a bit deeper than I have the time to dissect, but feels worthy of a look for somebody in the know. It sounds like we can error in counting metrics for some build patterns.
Local IO metrics published for stores whose inner was never metered: local_io_tracker() labels metrics from self.store_prefix unconditionally, so every store publishes local read/write metrics regardless of whether its inner was wrapped with MeteredObjectStore. ObjectStoreRegistry::get_store and the caller-supplied branch of from_uri_and_params (newly metered at line 495) both wrap inner, but ObjectStore::new (object_store.rs:1136) and the direct provider constructors ObjectStore::local() / ObjectStore::memory() do not. DatasetBuilder::build_object_store (rust/lance/src/dataset/builder.rs:556-572) routes a caller-supplied options.object_store through ObjectStore::new, not through from_uri_and_params, so the branch the diff just metered is bypassed by the main consumer of that option. Before this change such a store emitted no metrics at all; now it emits a partial set.
Fix: Apply the same #[cfg(feature = "metrics")] inner = inner.metered(store_prefix.clone()) wrapping in ObjectStore::new (object_store.rs:1162, alongside the existing wrapper.wrap), so every constructor that sets store_prefix also meters inner. Alternatively, only set metrics_base on the tracker for stores whose inner was metered, so coverage is all-or-nothing per store.
There was a problem hiding this comment.
This one was right, and worse than stated — the fix is in aa0642c.
The part the AI review got right: local_io_tracker() labelled unconditionally while inner was only metered in the registry path and the caller-supplied branch of from_uri_and_params. The part worth underlining is the consequence it spotted at the end — DatasetBuilder::build_object_store routes a caller-supplied object_store through ObjectStore::new, so the branch this PR metered was bypassed by the main consumer of that option. That second bullet in the description was largely ineffective as written.
Labelling the tracker and wrapping inner are now inseparable, in meter_store, called from all three constructors including ObjectStore::new. A store publishes metrics for all of its IO or for none of it. Stores built by calling a provider's new_store directly (ObjectStore::local / memory) are the "none" case, unchanged from before this PR and now documented in the module docs.
Side effect: the label is materialized once per store instead of on every open / create / copy / remove_dir_all, which was CodeRabbit's other point on the tracking_store.rs thread.
New test test_store_built_from_new_is_metered covers the ObjectStore::new path and asserts put / get (the bypass paths) and delete (which goes through inner), so it's the all-or-nothing invariant under test rather than just "some metric appeared". The two tests that used ObjectStore::local() now build through from_uri.
`local_io_tracker()` labelled the local bypass metrics from `store_prefix` unconditionally, but `inner` was only wrapped with `MeteredObjectStore` in the registry path and in the caller-supplied branch of `from_uri_and_params`. `ObjectStore::new` never wrapped it, and `DatasetBuilder::build_object_store` routes a caller-supplied `object_store` through `ObjectStore::new` — so the branch metered here was bypassed by the main consumer of that option, and such stores published local reads and writes but nothing for `list` / `delete` / `rename`. Labelling the tracker and wrapping the store now happen together in `meter_store`, called from all three constructors, so a store publishes metrics for all of its IO or for none of it. Stores built by calling a provider's `new_store` directly (`ObjectStore::local` / `memory`) are in the "none" case, as before. This also materializes the label once per store instead of on every `open` / `create` / `copy` / `remove_dir_all`. `get_range`, `get_all` and `stream_local_range` awaited their blocking task with `?`, so a `JoinError` — the read panicked — returned before the operation was recorded, losing exactly the failure worth counting. The shared `join_local_io` flattens the join and IO errors instead, matching what `size()` already did. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The optimized local paths —
LocalWriter,LocalObjectReader, the io_uring readers, and the localcopy/ recursive-delete shortcuts — go straight to the filesystem, so they never reachMeteredObjectStoreand published no metrics at all. Writing a dataset to a local path produced zero object store metrics even though the IO tracker recorded it.IOTrackernow carries the metricsbaselabel of the store it belongs to (its store prefix, the same valueMeteredObjectStoreis given) and hands out anIoMetricsGuardfor IO that bypasses theobject_storelayer. Each bypassing path now records requests, bytes, latency, errors and the in-flight gauge under the same labels, so local IO aggregates with cloud IO:LocalWriter(oneputper file, from open to durable at its final path)putLocalObjectReader::get_range/get_all/ streamed chunksgetLocalObjectReader::size(the local equivalent of a HEAD)headUringReader/UringCurrentThreadReaderget_range/get_allgetObjectStore::copycopyObjectStore::remove_dir_all(one request, likedelete_stream)deleteTwo things worth a second opinion:
innerand labelling itsIOTrackernow happen together inmeter_store, called from all three constructors — the registry,from_uri_and_params, andObjectStore::new. Previously only the first two meteredinner, so a caller-supplied store routed throughObjectStore::new(which is whatDatasetBuilder::build_object_storedoes) would have published local reads and writes but nothing forlist/delete/rename. A store now publishes for all of its IO or none of it; stores built by calling a provider'snew_storedirectly (ObjectStore::local/memory) are the "none" case. This would still double-count if a caller passed in a store Lance had already metered.IoStatsis left alone. The localcopy, recursive delete and size lookup were never counted there and still aren't, so existing IO assertions are unaffected; adding them would shift IO counts across the test suite.The io_uring readers have no coverage here — they need Linux plus a working ring, so the new tests exercise the non-uring local paths only.
Fixes #7993