From 92ee83ad13e4e19c992ecfa1e4f00c2f257ce7d5 Mon Sep 17 00:00:00 2001 From: peasee <98815791+peasee@users.noreply.github.com> Date: Fri, 23 Oct 2026 21:27:14 +1000 Subject: [PATCH 1/5] feat(scan): overlap split I/O without increasing decode concurrency Signed-off-by: peasee <98815791+peasee@users.noreply.github.com> --- vortex-file/src/open.rs | 8 +- vortex-layout/Cargo.toml | 2 +- vortex-layout/src/scan/mod.rs | 1 + vortex-layout/src/scan/repeated_scan.rs | 48 ++- vortex-layout/src/scan/scan_builder.rs | 24 ++ vortex-layout/src/scan/split_gate.rs | 382 ++++++++++++++++++++++ vortex-layout/src/segments/mod.rs | 2 + vortex-layout/src/segments/wait_marked.rs | 89 +++++ 8 files changed, 546 insertions(+), 10 deletions(-) create mode 100644 vortex-layout/src/scan/split_gate.rs create mode 100644 vortex-layout/src/segments/wait_marked.rs diff --git a/vortex-file/src/open.rs b/vortex-file/src/open.rs index b54da20e192..7ba48f8248a 100644 --- a/vortex-file/src/open.rs +++ b/vortex-file/src/open.rs @@ -22,6 +22,7 @@ use vortex_layout::segments::SegmentCacheSourceAdapter; use vortex_layout::segments::SegmentId; use vortex_layout::segments::SegmentSource; use vortex_layout::segments::SharedSegmentSource; +use vortex_layout::segments::WaitMarkedSegmentSource; use vortex_layout::session::LayoutSessionExt; use vortex_metrics::DefaultMetricsRegistry; use vortex_metrics::Label; @@ -326,9 +327,10 @@ impl VortexOpenOptions { ))); // Wrap up the segment source to first resolve segments from the initial read cache. - let segment_source: Arc = Arc::new(SegmentCacheSourceAdapter::new( - segment_cache, - segment_source, + // The outermost source, so each request pending on a cache fill or a shared read counts + // as a wait for the scan split that made it. + let segment_source: Arc = Arc::new(WaitMarkedSegmentSource::new( + SegmentCacheSourceAdapter::new(segment_cache, segment_source), )); let metadata = if self.include_metadata { diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index 99d8c605941..56a8e39f202 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -62,7 +62,7 @@ insta = { workspace = true } rand = { workspace = true } rstest = { workspace = true } temp-env = { workspace = true } -tokio = { workspace = true, features = ["rt", "macros"] } +tokio = { workspace = true, features = ["rt", "macros", "rt-multi-thread", "time"] } vortex-array = { path = "../vortex-array", features = ["_test-harness"] } vortex-io = { path = "../vortex-io", features = ["tokio"] } diff --git a/vortex-layout/src/scan/mod.rs b/vortex-layout/src/scan/mod.rs index 98fd1918a42..0de1ff7ca7a 100644 --- a/vortex-layout/src/scan/mod.rs +++ b/vortex-layout/src/scan/mod.rs @@ -8,6 +8,7 @@ pub mod multi; pub mod repeated_scan; pub mod scan_builder; pub mod split_by; +pub(crate) mod split_gate; mod splits; mod tasks; #[cfg(test)] diff --git a/vortex-layout/src/scan/repeated_scan.rs b/vortex-layout/src/scan/repeated_scan.rs index 06506c1b711..f8cc69d621d 100644 --- a/vortex-layout/src/scan/repeated_scan.rs +++ b/vortex-layout/src/scan/repeated_scan.rs @@ -6,7 +6,9 @@ use std::iter; use std::ops::Range; use std::sync::Arc; +use futures::FutureExt; use futures::Stream; +use futures::StreamExt; use futures::future::BoxFuture; use futures::future::Either as EitherFuture; use itertools::Either; @@ -28,6 +30,9 @@ use vortex_session::VortexSession; use crate::LayoutReaderRef; use crate::scan::filter::FilterExpr; use crate::scan::scan_builder::SplitConcurrency; +use crate::scan::split_gate::DecodePermit; +use crate::scan::split_gate::GatedSplit; +use crate::scan::split_gate::SplitGate; use crate::scan::splits::Splits; use crate::scan::tasks::TaskContext; use crate::scan::tasks::TaskFuture; @@ -51,6 +56,9 @@ pub struct RepeatedScan { splits: Splits, /// How many splits to make progress on concurrently. concurrency: SplitConcurrency, + /// Splits allowed in flight beyond `concurrency`, see + /// [`ScanBuilder::with_split_io_overlap`](crate::scan::scan_builder::ScanBuilder::with_split_io_overlap). + max_splits_in_flight: Option, /// Function to apply to each [`ArrayRef`] within the spawned split tasks. map_fn: Arc VortexResult + Send + Sync>, /// Maximal number of rows to read (after filtering) @@ -101,6 +109,7 @@ impl RepeatedScan { selection: Selection, splits: Splits, concurrency: SplitConcurrency, + max_splits_in_flight: Option, map_fn: Arc VortexResult + Send + Sync>, limit: Option, dtype: DType, @@ -115,6 +124,7 @@ impl RepeatedScan { selection, splits, concurrency, + max_splits_in_flight, map_fn, limit, dtype, @@ -235,7 +245,8 @@ impl RepeatedScan { }) } - /// Streams the scan, holding at most [`SplitConcurrency::effective`] splits in flight. + /// Streams the scan, holding at most [`SplitConcurrency::effective`] splits in flight, or up + /// to the overlap set with `ScanBuilder::with_split_io_overlap` when the scan is unordered. /// /// Splits are constructed lazily under that bound, so a split's reads are not registered /// until it is admitted. @@ -243,22 +254,47 @@ impl RepeatedScan { &self, row_range: Option>, ) -> VortexResult> + Send + 'static + use> { - use futures::StreamExt; let concurrency = self.concurrency.effective(); let handle = self.session.handle(); + // Only an unordered scan can overlap I/O: see `ScanBuilder::with_split_io_overlap`. + let gate = match self.max_splits_in_flight { + Some(max) if !self.ordered && max > concurrency => Some(SplitGate::new(concurrency)), + _ => None, + }; + let in_flight = match (&gate, self.max_splits_in_flight) { + (Some(_), Some(max)) => max, + _ => concurrency, + }; + + // The permit of a finished split travels with its output and is dropped when the output + // leaves the stream below. let stream = futures::stream::iter(self.split_tasks(row_range)).map(move |task| match task { - Ok(task) => EitherFuture::Left(handle.spawn(task)), - Err(err) => EitherFuture::Right(futures::future::ready(Err(err))), + Ok(task) => match &gate { + Some(gate) => EitherFuture::Left(EitherFuture::Left( + handle.spawn(GatedSplit::new(task, Arc::clone(gate))), + )), + None => EitherFuture::Left(EitherFuture::Right( + handle.spawn(task.map(|output| (output, None::))), + )), + }, + Err(err) => { + EitherFuture::Right(futures::future::ready((Err(err), None::))) + } }); let stream = if self.ordered { - stream.buffered(concurrency).boxed() + stream.buffered(in_flight).boxed() } else { - stream.buffer_unordered(concurrency).boxed() + stream.buffer_unordered(in_flight).boxed() }; + let stream = stream.map(|(output, permit)| { + drop(permit); + output + }); + Ok(stream.filter_map(|chunk| async move { chunk.transpose() })) } } diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index c4931243c90..637f9b348d5 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -102,6 +102,9 @@ pub struct ScanBuilder { natural_splits: Option>, /// How many splits to make progress on concurrently. concurrency: SplitConcurrency, + /// Splits allowed in flight beyond the decoding concurrency, see + /// [`Self::with_split_io_overlap`]. + max_splits_in_flight: Option, /// Function to apply to each [`ArrayRef`] within the spawned split tasks. map_fn: Arc VortexResult + Send + Sync>, metrics_registry: Option>, @@ -131,6 +134,7 @@ impl ScanBuilder { // We default to four tasks per worker thread, which allows for some I/O lookahead // without too much impact on work-stealing. concurrency: SplitConcurrency::PerWorker(4), + max_splits_in_flight: None, map_fn: Arc::new(Ok), metrics_registry: None, file_stats: None, @@ -296,6 +300,24 @@ impl ScanBuilder { self } + /// Lets the scan keep up to `max_splits_in_flight` splits in flight while it still decodes + /// only as many at once as [`Self::concurrency`]. + /// + /// A split holds one of `concurrency` decode permits while it runs and gives it back while it + /// waits for a segment from a source wrapped in + /// [`WaitMarkedSegmentSource`](crate::segments::WaitMarkedSegmentSource), so a scan over a + /// high-latency source can have many reads outstanding without decoding many splits at once. + /// A split's permit is held until its output leaves the stream, so the decoded batches a scan + /// holds stay bounded by `concurrency`. + /// + /// Has no effect unless it exceeds the concurrency, and none on an ordered scan: an ordered + /// stream releases a split's output only in file order, so a split that finished early could + /// hold the permit the earlier split needs to finish. + pub fn with_split_io_overlap(mut self, max_splits_in_flight: usize) -> Self { + self.max_splits_in_flight = Some(max_splits_in_flight); + self + } + /// Add or clear the metrics registry used by scan execution. pub fn with_some_metrics_registry(mut self, metrics: Option>) -> Self { self.metrics_registry = metrics; @@ -347,6 +369,7 @@ impl ScanBuilder { split_by: self.split_by, natural_splits: self.natural_splits, concurrency: self.concurrency, + max_splits_in_flight: self.max_splits_in_flight, metrics_registry: self.metrics_registry, file_stats: self.file_stats, limit: self.limit, @@ -415,6 +438,7 @@ impl ScanBuilder { self.selection, splits, self.concurrency, + self.max_splits_in_flight, self.map_fn, self.limit, dtype, diff --git a/vortex-layout/src/scan/split_gate.rs b/vortex-layout/src/scan/split_gate.rs new file mode 100644 index 00000000000..001e946c8af --- /dev/null +++ b/vortex-layout/src/scan/split_gate.rs @@ -0,0 +1,382 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Lets a scan keep more splits in flight than it decodes at once. +//! +//! A split spends most of its life waiting for segment bytes, during which it uses no CPU. When +//! the scan's concurrency is also the number of splits with reads outstanding, a high-latency +//! source (an object store) is left idle between a split's reads: the scan has to choose between +//! few reads in flight and many concurrent decodes. +//! +//! [`SplitGate`] separates the two. Every split takes a *decode permit* before it is polled and +//! gives it back while it waits on segment I/O marked by +//! [`WaitMarkedSegmentSource`](crate::segments::WaitMarkedSegmentSource). At most as many splits +//! decode as there are permits, however many are in flight, and a finished split keeps its +//! permit until its output leaves the scan stream, so the decoded batches a scan holds stay +//! bounded by the permit count as well. + +use std::cell::RefCell; +use std::collections::VecDeque; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; + +use parking_lot::Mutex; +use pin_project_lite::pin_project; + +/// Counts the marked waits that are in progress within one split. +#[derive(Debug, Default)] +pub(crate) struct WaitScope { + waits: AtomicUsize, +} + +impl WaitScope { + pub(crate) fn begin(&self) { + self.waits.fetch_add(1, Ordering::AcqRel); + } + + pub(crate) fn end(&self) { + self.waits.fetch_sub(1, Ordering::AcqRel); + } + + fn waiting(&self) -> bool { + self.waits.load(Ordering::Acquire) > 0 + } +} + +thread_local! { + /// The split being polled on this thread, if any. + static CURRENT: RefCell>> = const { RefCell::new(None) }; +} + +/// The split being polled on this thread, if there is one. +pub(crate) fn current_scope() -> Option> { + CURRENT.with(|current| current.borrow().clone()) +} + +/// Restores the previous scope when dropped, so a panic in the polled future cannot leak one. +struct ScopeGuard(Option>); + +impl Drop for ScopeGuard { + fn drop(&mut self) { + CURRENT.with(|current| *current.borrow_mut() = self.0.take()); + } +} + +#[derive(Debug)] +struct GateState { + available: usize, + waiters: VecDeque<(u64, Waker)>, + next_waiter: u64, +} + +/// A counting semaphore handing out decode permits. +#[derive(Debug)] +pub(crate) struct SplitGate { + state: Mutex, +} + +impl SplitGate { + pub(crate) fn new(permits: usize) -> Arc { + Arc::new(Self { + state: Mutex::new(GateState { + available: permits.max(1), + waiters: VecDeque::new(), + next_waiter: 0, + }), + }) + } + + /// Takes a permit if one is free, otherwise queues `waiter` to be woken when one is. + fn poll_acquire( + self: &Arc, + waiter: &mut Option, + cx: &Context<'_>, + ) -> Poll { + let mut state = self.state.lock(); + if state.available > 0 { + state.available -= 1; + if let Some(id) = waiter.take() { + state.waiters.retain(|(queued, _)| *queued != id); + } + return Poll::Ready(DecodePermit { + gate: Arc::clone(self), + }); + } + match *waiter { + Some(id) => { + if let Some((_, waker)) = state.waiters.iter_mut().find(|(queued, _)| *queued == id) + { + waker.clone_from(cx.waker()); + } + } + None => { + let id = state.next_waiter; + state.next_waiter += 1; + state.waiters.push_back((id, cx.waker().clone())); + *waiter = Some(id); + } + } + Poll::Pending + } + + /// Removes a waiter that will not poll again, passing on a wake-up it may have consumed. + fn cancel(&self, id: u64) { + let next = { + let mut state = self.state.lock(); + state.waiters.retain(|(queued, _)| *queued != id); + if state.available > 0 { + state.waiters.front().map(|(_, waker)| waker.clone()) + } else { + None + } + }; + if let Some(waker) = next { + waker.wake(); + } + } + + fn release(&self) { + let next = { + let mut state = self.state.lock(); + state.available += 1; + state.waiters.front().map(|(_, waker)| waker.clone()) + }; + if let Some(waker) = next { + waker.wake(); + } + } +} + +/// A decode permit; returned to its gate when dropped. +#[derive(Debug)] +pub(crate) struct DecodePermit { + gate: Arc, +} + +impl Drop for DecodePermit { + fn drop(&mut self) { + self.gate.release(); + } +} + +pin_project! { + /// Runs one split future under a [`SplitGate`]. + /// + /// The split is polled only while it holds a permit. It gives the permit back when a poll + /// leaves it waiting on marked segment I/O, and takes one again before its next poll. When it + /// finishes it hands its permit to the caller together with the output, to be dropped when + /// the output leaves the scan. + pub(crate) struct GatedSplit { + #[pin] + inner: F, + gate: Arc, + scope: Arc, + permit: Option, + waiter: Option, + } + + impl PinnedDrop for GatedSplit { + fn drop(this: Pin<&mut Self>) { + let this = this.project(); + if let Some(id) = this.waiter.take() { + this.gate.cancel(id); + } + } + } +} + +impl GatedSplit { + pub(crate) fn new(inner: F, gate: Arc) -> Self { + Self { + inner, + gate, + scope: Arc::new(WaitScope::default()), + permit: None, + waiter: None, + } + } +} + +impl Future for GatedSplit { + type Output = (F::Output, Option); + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let mut this = self.project(); + if this.permit.is_none() { + match this.gate.poll_acquire(this.waiter, cx) { + Poll::Ready(permit) => *this.permit = Some(permit), + Poll::Pending => return Poll::Pending, + } + } + + let previous = CURRENT.with(|current| current.replace(Some(Arc::clone(this.scope)))); + let polled = { + let _restore = ScopeGuard(previous); + this.inner.as_mut().poll(cx) + }; + + match polled { + Poll::Ready(output) => Poll::Ready((output, this.permit.take())), + Poll::Pending => { + if this.scope.waiting() { + // Waiting on I/O uses no CPU: let another split decode meanwhile. + drop(this.permit.take()); + } + Poll::Pending + } + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use std::time::Duration; + + use futures::FutureExt; + use futures::StreamExt; + use futures::future::BoxFuture; + use futures::stream; + use vortex_array::buffer::BufferHandle; + use vortex_buffer::ByteBuffer; + use vortex_error::VortexResult; + + use super::*; + use crate::segments::SegmentFuture; + use crate::segments::SegmentId; + use crate::segments::SegmentSource; + use crate::segments::WaitMarkedSegmentSource; + + /// Resolves each request after a pause, standing in for a remote read. + struct SlowSource; + + impl SegmentSource for SlowSource { + fn request(&self, _id: SegmentId) -> SegmentFuture { + async { + tokio::time::sleep(Duration::from_millis(20)).await; + Ok(BufferHandle::new_host(ByteBuffer::empty())) + } + .boxed() + } + } + + /// Tracks how many splits are inside a CPU section and how many are in flight. + #[derive(Default)] + struct Probe { + decoding: AtomicUsize, + max_decoding: AtomicUsize, + reading: AtomicUsize, + max_reading: AtomicUsize, + } + + fn bump(now: &AtomicUsize, max: &AtomicUsize) { + let value = now.fetch_add(1, Ordering::SeqCst) + 1; + max.fetch_max(value, Ordering::SeqCst); + } + + fn split( + source: Arc>, + probe: Arc, + ) -> BoxFuture<'static, VortexResult<()>> { + async move { + // CPU section before the read: the permit is held. + bump(&probe.decoding, &probe.max_decoding); + probe.decoding.fetch_sub(1, Ordering::SeqCst); + + bump(&probe.reading, &probe.max_reading); + source.request(SegmentId::from(0)).await?; + probe.reading.fetch_sub(1, Ordering::SeqCst); + + // CPU section after the read: the permit has been retaken. Yield first so another + // split can be inside its own CPU section when this one is. + bump(&probe.decoding, &probe.max_decoding); + std::thread::sleep(Duration::from_millis(2)); + probe.decoding.fetch_sub(1, Ordering::SeqCst); + Ok(()) + } + .boxed() + } + + async fn run(permits: usize, in_flight: usize, splits: usize) -> Arc { + let probe = Arc::new(Probe::default()); + let source = Arc::new(WaitMarkedSegmentSource::new(SlowSource)); + let gate = SplitGate::new(permits); + let tasks = (0..splits).map(|_| { + let gated = GatedSplit::new( + split(Arc::clone(&source), Arc::clone(&probe)), + Arc::clone(&gate), + ); + tokio::spawn(gated) + }); + let done: Vec<_> = stream::iter(tasks) + .buffer_unordered(in_flight) + .collect() + .await; + for result in done { + let (output, _permit) = result.expect("split task joins"); + output.expect("split succeeds"); + } + probe + } + + /// Splits waiting on a segment read hand their permit on, so more splits have reads + /// outstanding than there are permits, while decoding stays within the permit count. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn reads_overlap_beyond_the_permit_count() { + let probe = run(2, 16, 64).await; + assert!( + probe.max_decoding.load(Ordering::SeqCst) <= 2, + "decoding exceeded the permits: {}", + probe.max_decoding.load(Ordering::SeqCst) + ); + assert!( + probe.max_reading.load(Ordering::SeqCst) > 2, + "reads never overlapped beyond the permits: {}", + probe.max_reading.load(Ordering::SeqCst) + ); + } + + /// Without the gate's hand-off a split holds its permit through its read, so the reads in + /// flight never exceed the permit count; this is the behaviour the hand-off removes. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn unmarked_waits_keep_the_permit() { + let probe = Arc::new(Probe::default()); + let gate = SplitGate::new(2); + let tasks = (0..32).map(|_| { + let probe = Arc::clone(&probe); + let unmarked = async move { + bump(&probe.reading, &probe.max_reading); + tokio::time::sleep(Duration::from_millis(5)).await; + probe.reading.fetch_sub(1, Ordering::SeqCst); + }; + tokio::spawn(GatedSplit::new(unmarked, Arc::clone(&gate))) + }); + let done: Vec<_> = stream::iter(tasks).buffer_unordered(16).collect().await; + assert_eq!(done.len(), 32); + assert_eq!(probe.max_reading.load(Ordering::SeqCst), 2); + } + + /// A finished split keeps its permit until the caller drops it. + #[tokio::test] + async fn a_finished_split_holds_its_permit_until_dropped() { + let gate = SplitGate::new(1); + let (first, permit) = GatedSplit::new(async { 1 }, Arc::clone(&gate)).await; + assert_eq!(first, 1); + assert!(permit.is_some(), "the finished split hands its permit over"); + let mut second = Box::pin(GatedSplit::new(async { 2 }, Arc::clone(&gate))); + assert!( + futures::poll!(second.as_mut()).is_pending(), + "the second split must wait for the first one's permit" + ); + drop(permit); + let (value, _permit) = second.await; + assert_eq!(value, 2); + } +} diff --git a/vortex-layout/src/segments/mod.rs b/vortex-layout/src/segments/mod.rs index a7ecf26de3b..43031a93031 100644 --- a/vortex-layout/src/segments/mod.rs +++ b/vortex-layout/src/segments/mod.rs @@ -12,6 +12,7 @@ mod cache; mod shared; mod sink; mod source; +mod wait_marked; #[cfg(any(test, feature = "_test-harness"))] mod test; @@ -26,6 +27,7 @@ pub use source::*; #[cfg(any(test, feature = "_test-harness"))] pub use test::*; use vortex_error::VortexError; +pub use wait_marked::*; /// Identifier for a single physical segment referenced by a layout. /// diff --git a/vortex-layout/src/segments/wait_marked.rs b/vortex-layout/src/segments/wait_marked.rs new file mode 100644 index 00000000000..ce70d62978e --- /dev/null +++ b/vortex-layout/src/segments/wait_marked.rs @@ -0,0 +1,89 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; + +use futures::FutureExt; +use pin_project_lite::pin_project; +use vortex_array::buffer::BufferHandle; +use vortex_error::VortexResult; + +use crate::scan::split_gate::WaitScope; +use crate::scan::split_gate::current_scope; +use crate::segments::SegmentFuture; +use crate::segments::SegmentId; +use crate::segments::SegmentSource; + +/// A [`SegmentSource`] that reports a scan split's waits for segment bytes. +/// +/// A split that is polled under a scan's split gate and finds a requested segment not yet +/// available is waiting on I/O, which uses no CPU. The gate uses that to let other splits +/// decode in the meantime. Requests made outside a gated split are passed through untouched. +/// +/// Wrap the outermost source a reader sees, so that a segment served from a cache or shared +/// with another request still counts as a wait for each request that is pending on it. +pub struct WaitMarkedSegmentSource { + inner: S, +} + +impl WaitMarkedSegmentSource { + /// Wraps `inner`. + pub fn new(inner: S) -> Self { + Self { inner } + } +} + +impl SegmentSource for WaitMarkedSegmentSource { + fn request(&self, id: SegmentId) -> SegmentFuture { + WaitMarked { + inner: self.inner.request(id), + scope: None, + } + .boxed() + } +} + +pin_project! { + struct WaitMarked { + #[pin] + inner: SegmentFuture, + // The split this request is making wait, while it is pending. + scope: Option>, + } + + impl PinnedDrop for WaitMarked { + fn drop(this: Pin<&mut Self>) { + if let Some(scope) = this.project().scope.take() { + scope.end(); + } + } + } +} + +impl Future for WaitMarked { + type Output = VortexResult; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.project(); + let polled = this.inner.poll(cx); + match (&polled, this.scope.is_some()) { + (Poll::Pending, false) => { + if let Some(scope) = current_scope() { + scope.begin(); + *this.scope = Some(scope); + } + } + (Poll::Ready(_), true) => { + if let Some(scope) = this.scope.take() { + scope.end(); + } + } + _ => {} + } + polled + } +} From 8ce85c5cd392628e5ff05718900fe52bfad984bf Mon Sep 17 00:00:00 2001 From: peasee <98815791+peasee@users.noreply.github.com> Date: Fri, 23 Oct 2026 21:36:11 +1000 Subject: [PATCH 2/5] test(scan): release permits while collecting overlap results Signed-off-by: peasee <98815791+peasee@users.noreply.github.com> --- vortex-layout/src/scan/split_gate.rs | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/vortex-layout/src/scan/split_gate.rs b/vortex-layout/src/scan/split_gate.rs index 001e946c8af..7b530b6e6e6 100644 --- a/vortex-layout/src/scan/split_gate.rs +++ b/vortex-layout/src/scan/split_gate.rs @@ -317,10 +317,14 @@ mod tests { }); let done: Vec<_> = stream::iter(tasks) .buffer_unordered(in_flight) + .map(|result| { + let (output, permit) = result.expect("split task joins"); + drop(permit); + output + }) .collect() .await; - for result in done { - let (output, _permit) = result.expect("split task joins"); + for output in done { output.expect("split succeeds"); } probe @@ -358,7 +362,15 @@ mod tests { }; tokio::spawn(GatedSplit::new(unmarked, Arc::clone(&gate))) }); - let done: Vec<_> = stream::iter(tasks).buffer_unordered(16).collect().await; + let done: Vec<_> = stream::iter(tasks) + .buffer_unordered(16) + .map(|result| { + let (output, permit) = result.expect("split task joins"); + drop(permit); + output + }) + .collect() + .await; assert_eq!(done.len(), 32); assert_eq!(probe.max_reading.load(Ordering::SeqCst), 2); } From b0df3cdede8cb51507f970cbbf5bfee98d227b16 Mon Sep 17 00:00:00 2001 From: peasee <98815791+peasee@users.noreply.github.com> Date: Sat, 24 Oct 2026 12:12:33 +1000 Subject: [PATCH 3/5] style(scan): format layout test dependencies Signed-off-by: peasee <98815791+peasee@users.noreply.github.com> --- vortex-layout/Cargo.toml | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index 56a8e39f202..605b46dfe04 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -62,7 +62,12 @@ insta = { workspace = true } rand = { workspace = true } rstest = { workspace = true } temp-env = { workspace = true } -tokio = { workspace = true, features = ["rt", "macros", "rt-multi-thread", "time"] } +tokio = { workspace = true, features = [ + "rt", + "macros", + "rt-multi-thread", + "time", +] } vortex-array = { path = "../vortex-array", features = ["_test-harness"] } vortex-io = { path = "../vortex-io", features = ["tokio"] } From 2b1f2ed92cd534638341c99a8377e3f416219921 Mon Sep 17 00:00:00 2001 From: peasee <98815791+peasee@users.noreply.github.com> Date: Sat, 24 Oct 2026 12:56:04 +1000 Subject: [PATCH 4/5] fix(vortex-layout): wake queued splits and mark shared waits --- SPICE_PATCHES.md | 1 + vortex-layout/src/layouts/dict/reader.rs | 11 ++- vortex-layout/src/layouts/flat/reader.rs | 5 +- vortex-layout/src/scan/split_gate.rs | 88 +++++++++++++++++++++++ vortex-layout/src/segments/mod.rs | 3 +- vortex-layout/src/segments/wait_marked.rs | 25 +++---- 6 files changed, 116 insertions(+), 17 deletions(-) diff --git a/SPICE_PATCHES.md b/SPICE_PATCHES.md index cef260fd871..9f38db971b2 100644 --- a/SPICE_PATCHES.md +++ b/SPICE_PATCHES.md @@ -60,6 +60,7 @@ exists: a grep proves presence, a test proves behaviour. | 12 | Restore lint checks on forks | `bb80c537b` | Fork CI not running lints | — (CI config) | No | | 13 | `set_available_parallelism` | this change | Scan and writer fan-out sized from the machine's core count rather than what the host process is entitled to (spiceai/spiceai#12328) | `cargo test -p vortex-utils --test parallelism_declared --test parallelism_declared_too_late` | Proposed — additive, detection unchanged | | 14 | Arrow `Map` alias | `1a6dc54f1` (`lukim/map`) | `Map` columns unwritable: `Array encoding not implemented for Arrow data type Map(...)` on every flush (spiceai/spiceai#13524) | `cargo test -p vortex-arrow --lib executor::map` | Yes — 0.86.1 native `Map` (`#9111`) | +| 17 | Overlap remote split I/O without increasing concurrent decodes | spiceai/vortex#107 | Cold object-store scans hold their limited split slots while awaiting segment bytes, leaving read capacity idle; shared pending arrays must hand off a permit for every waiting split | `cargo test -p vortex-layout --lib scan::split_gate::tests` and `grep -q 'mark_pending_wait(array)' vortex-layout/src/layouts/flat/reader.rs` and `grep -q 'mark_pending_wait(values_eval)' vortex-layout/src/layouts/dict/reader.rs` | No | **Confidence:** rows 1–4, 13 and 14 are verified against the branches. Rows 5–12 are seeded from an audit of non-merge commits on `spiceai-54` authored by Spice engineers; their descriptions diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 10984c9d9e7..f3d83ea0ebf 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -42,6 +42,7 @@ use crate::RowSplits; use crate::SplitRange; use crate::layouts::SharedArrayFuture; use crate::segments::SegmentSource; +use crate::segments::mark_pending_wait; pub struct DictReader { layout: DictLayout, @@ -273,7 +274,10 @@ impl LayoutReader for DictReader { Ok(MaskFuture::new(mask.len(), async move { // Join on the I/O futures first, before the mask. - let (codes, values) = try_join!(codes_eval, values_eval.map_err(VortexError::from))?; + let (codes, values) = try_join!( + codes_eval, + mark_pending_wait(values_eval).map_err(VortexError::from) + )?; let mask = mask.await?; let mut ctx = session.create_execution_ctx(); @@ -325,7 +329,10 @@ impl LayoutReader for DictReader { }; let all_values_referenced = self.layout.has_all_values_referenced(); Ok(async move { - let (values, codes) = try_join!(values_eval.map_err(VortexError::from), codes_eval)?; + let (values, codes) = try_join!( + mark_pending_wait(values_eval).map_err(VortexError::from), + codes_eval + )?; // SAFETY: Layout was validated at write time. // * The codes dtype is guaranteed to be an integer type from the layout diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index f6b3d10da86..64ee5a24b1b 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -27,6 +27,7 @@ use crate::reader::LayoutReader; use crate::reader::RowSplits; use crate::reader::SplitRange; use crate::segments::SegmentSource; +use crate::segments::mark_pending_wait; /// The threshold of mask density below which we will evaluate the expression only over the /// selected rows, and above which we evaluate the expression over all rows and then select @@ -144,7 +145,7 @@ impl LayoutReader for FlatReader { // TODO(ngates): if the mask density is low enough, or if the mask is dense within a range // (as often happens with zone map pruning), then we could slice/filter the array prior // to evaluating the expression. - let mut array = array.clone().await?; + let mut array = mark_pending_wait(array).await?; let mask = mask.await?; // Slice the array based on the row mask. @@ -201,7 +202,7 @@ impl LayoutReader for FlatReader { Ok(async move { trace!("Flat array evaluation {} - {}", name, expr); - let mut array = array.clone().await?; + let mut array = mark_pending_wait(array).await?; let mask = mask.await?; // Slice the array based on the row mask. diff --git a/vortex-layout/src/scan/split_gate.rs b/vortex-layout/src/scan/split_gate.rs index 7b530b6e6e6..ba4be88eec5 100644 --- a/vortex-layout/src/scan/split_gate.rs +++ b/vortex-layout/src/scan/split_gate.rs @@ -104,6 +104,15 @@ impl SplitGate { if let Some(id) = waiter.take() { state.waiters.retain(|(queued, _)| *queued != id); } + // Several releases can wake the same front waiter before it runs. Pass the wake-up + // on if another permit is still free after this waiter acquires one. + let next = (state.available > 0) + .then(|| state.waiters.front().map(|(_, waker)| waker.clone())) + .flatten(); + drop(state); + if let Some(waker) = next { + waker.wake(); + } return Poll::Ready(DecodePermit { gate: Arc::clone(self), }); @@ -238,6 +247,8 @@ impl Future for GatedSplit { mod tests { use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; + use std::task::Wake; + use std::task::Waker; use std::time::Duration; use futures::FutureExt; @@ -253,6 +264,7 @@ mod tests { use crate::segments::SegmentId; use crate::segments::SegmentSource; use crate::segments::WaitMarkedSegmentSource; + use crate::segments::mark_pending_wait; /// Resolves each request after a pause, standing in for a remote read. struct SlowSource; @@ -391,4 +403,80 @@ mod tests { let (value, _permit) = second.await; assert_eq!(value, 2); } + + #[tokio::test] + async fn consumers_of_one_pending_array_each_release_their_permit() { + let gate = SplitGate::new(1); + // One marked source future is shared by two splits, like a FlatReader array future. + let shared = mark_pending_wait(futures::future::pending::()) + .boxed() + .shared(); + let mut first = Box::pin(GatedSplit::new( + mark_pending_wait(shared.clone()), + Arc::clone(&gate), + )); + let mut second = Box::pin(GatedSplit::new( + mark_pending_wait(shared), + Arc::clone(&gate), + )); + assert!(futures::poll!(first.as_mut()).is_pending()); + assert!(futures::poll!(second.as_mut()).is_pending()); + + // The second consumer did not poll the shared source. It must still hand off its + // permit so another split can start a read while the shared array is pending. + let mut third = Box::pin(GatedSplit::new(async { 3 }, gate)); + assert!(futures::poll!(third.as_mut()).is_ready()); + } + + #[test] + fn two_releases_wake_two_queued_splits() { + struct CountWakes(AtomicUsize); + impl Wake for CountWakes { + fn wake(self: Arc) { + self.wake_by_ref(); + } + + fn wake_by_ref(self: &Arc) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + let gate = SplitGate::new(2); + let first_waker: Waker = Arc::new(CountWakes(AtomicUsize::new(0))).into(); + let second_wakes = Arc::new(CountWakes(AtomicUsize::new(0))); + let second_waker: Waker = Arc::clone(&second_wakes).into(); + let first_context = Context::from_waker(&first_waker); + let second_context = Context::from_waker(&second_waker); + + let mut first_id = None; + let mut second_id = None; + let first = match gate.poll_acquire(&mut first_id, &first_context) { + Poll::Ready(permit) => permit, + Poll::Pending => panic!("first permit must be available"), + }; + let second = match gate.poll_acquire(&mut second_id, &second_context) { + Poll::Ready(permit) => permit, + Poll::Pending => panic!("second permit must be available"), + }; + + assert!( + gate.poll_acquire(&mut first_id, &first_context) + .is_pending() + ); + assert!( + gate.poll_acquire(&mut second_id, &second_context) + .is_pending() + ); + drop(first); + drop(second); + + let _first_reacquired = match gate.poll_acquire(&mut first_id, &first_context) { + Poll::Ready(permit) => permit, + Poll::Pending => panic!("released permit must be available"), + }; + assert!( + second_wakes.0.load(Ordering::SeqCst) > 0, + "the second queued split must be woken for the remaining permit" + ); + } } diff --git a/vortex-layout/src/segments/mod.rs b/vortex-layout/src/segments/mod.rs index 43031a93031..e51b54b1ae6 100644 --- a/vortex-layout/src/segments/mod.rs +++ b/vortex-layout/src/segments/mod.rs @@ -27,7 +27,8 @@ pub use source::*; #[cfg(any(test, feature = "_test-harness"))] pub use test::*; use vortex_error::VortexError; -pub use wait_marked::*; +pub use wait_marked::WaitMarkedSegmentSource; +pub(crate) use wait_marked::mark_pending_wait; /// Identifier for a single physical segment referenced by a layout. /// diff --git a/vortex-layout/src/segments/wait_marked.rs b/vortex-layout/src/segments/wait_marked.rs index ce70d62978e..7f9826a2f24 100644 --- a/vortex-layout/src/segments/wait_marked.rs +++ b/vortex-layout/src/segments/wait_marked.rs @@ -9,8 +9,6 @@ use std::task::Poll; use futures::FutureExt; use pin_project_lite::pin_project; -use vortex_array::buffer::BufferHandle; -use vortex_error::VortexResult; use crate::scan::split_gate::WaitScope; use crate::scan::split_gate::current_scope; @@ -39,23 +37,26 @@ impl WaitMarkedSegmentSource { impl SegmentSource for WaitMarkedSegmentSource { fn request(&self, id: SegmentId) -> SegmentFuture { - WaitMarked { - inner: self.inner.request(id), - scope: None, - } - .boxed() + mark_pending_wait(self.inner.request(id)).boxed() } } +/// Mark each consumer's pending wait, including consumers of a shared layout future. +/// The underlying segment future may be polled by just one consumer, so marking only that +/// future would leave the other consumers holding decode permits while they wait for it. +pub(crate) fn mark_pending_wait(inner: F) -> impl Future { + WaitMarked { inner, scope: None } +} + pin_project! { - struct WaitMarked { + struct WaitMarked { #[pin] - inner: SegmentFuture, + inner: F, // The split this request is making wait, while it is pending. scope: Option>, } - impl PinnedDrop for WaitMarked { + impl PinnedDrop for WaitMarked { fn drop(this: Pin<&mut Self>) { if let Some(scope) = this.project().scope.take() { scope.end(); @@ -64,8 +65,8 @@ pin_project! { } } -impl Future for WaitMarked { - type Output = VortexResult; +impl Future for WaitMarked { + type Output = F::Output; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.project(); From f4bd807f5324f8db72a6f08f0453d46352010de4 Mon Sep 17 00:00:00 2001 From: peasee <98815791+peasee@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:35:45 +1000 Subject: [PATCH 5/5] feat(vortex-io): mark remote object-store waits through callback --- SPICE_PATCHES.md | 1 + vortex-io/src/object_store/read_at.rs | 157 +++++++++++++++++++++++--- 2 files changed, 140 insertions(+), 18 deletions(-) diff --git a/SPICE_PATCHES.md b/SPICE_PATCHES.md index 9f38db971b2..a89b72d08d9 100644 --- a/SPICE_PATCHES.md +++ b/SPICE_PATCHES.md @@ -61,6 +61,7 @@ exists: a grep proves presence, a test proves behaviour. | 13 | `set_available_parallelism` | this change | Scan and writer fan-out sized from the machine's core count rather than what the host process is entitled to (spiceai/spiceai#12328) | `cargo test -p vortex-utils --test parallelism_declared --test parallelism_declared_too_late` | Proposed — additive, detection unchanged | | 14 | Arrow `Map` alias | `1a6dc54f1` (`lukim/map`) | `Map` columns unwritable: `Array encoding not implemented for Arrow data type Map(...)` on every flush (spiceai/spiceai#13524) | `cargo test -p vortex-arrow --lib executor::map` | Yes — 0.86.1 native `Map` (`#9111`) | | 17 | Overlap remote split I/O without increasing concurrent decodes | spiceai/vortex#107 | Cold object-store scans hold their limited split slots while awaiting segment bytes, leaving read capacity idle; shared pending arrays must hand off a permit for every waiting split | `cargo test -p vortex-layout --lib scan::split_gate::tests` and `grep -q 'mark_pending_wait(array)' vortex-layout/src/layouts/flat/reader.rs` and `grep -q 'mark_pending_wait(values_eval)' vortex-layout/src/layouts/dict/reader.rs` | No | +| 18 | Callback for object-store I/O waits | spiceai/vortex#107 | The embedding query executor must be able to mark metadata, GET, and response-stream waits without classifying Vortex decode work as remote I/O | `cargo test -p vortex-io --features object_store,tokio --lib io_wait_marker_covers_head_and_both_range_read_paths` | No | **Confidence:** rows 1–4, 13 and 14 are verified against the branches. Rows 5–12 are seeded from an audit of non-merge commits on `spiceai-54` authored by Spice engineers; their descriptions diff --git a/vortex-io/src/object_store/read_at.rs b/vortex-io/src/object_store/read_at.rs index ed36c2a02a1..a9f4f8c1e29 100644 --- a/vortex-io/src/object_store/read_at.rs +++ b/vortex-io/src/object_store/read_at.rs @@ -4,22 +4,25 @@ use std::io; use std::sync::Arc; +use bytes::Bytes; use futures::FutureExt; use futures::SinkExt; use futures::StreamExt; use futures::channel::mpsc; use futures::future::BoxFuture; use futures::stream; +use futures::stream::BoxStream; use object_store::GetOptions; use object_store::GetRange; +use object_store::GetResult; use object_store::GetResultPayload; +use object_store::ObjectMeta; use object_store::ObjectStore; use object_store::ObjectStoreExt; use object_store::path::Path as ObjectPath; use vortex_array::buffer::BufferHandle; use vortex_array::memory::BufferAllocatorRef; use vortex_buffer::Alignment; -use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -34,6 +37,32 @@ use crate::std_file::read_exact_at; /// Default number of concurrent requests to allow. pub const DEFAULT_CONCURRENCY: usize = 192; +/// Marks waits on a remote object store without coupling Vortex to a query engine. +/// +/// The callback receives the network request and its response byte stream separately. It must +/// preserve their outputs. In particular, it should not mark work that decodes a Vortex array or +/// copies bytes after a stream item is ready. Applications can leave the callback unset for local +/// and in-memory stores. +pub trait ObjectStoreIoWait: Send + Sync + 'static { + /// Wrap a remote object metadata request. + fn head( + &self, + future: BoxFuture<'static, object_store::Result>, + ) -> BoxFuture<'static, object_store::Result>; + + /// Wrap a remote object range request. + fn get( + &self, + future: BoxFuture<'static, object_store::Result>, + ) -> BoxFuture<'static, object_store::Result>; + + /// Wrap each pending wait for the next bytes from a remote response. + fn stream( + &self, + stream: BoxStream<'static, object_store::Result>, + ) -> BoxStream<'static, object_store::Result>; +} + /// An object store backed I/O source. pub struct ObjectStoreReadAt { store: Arc, @@ -43,6 +72,7 @@ pub struct ObjectStoreReadAt { allocator: BufferAllocatorRef, concurrency: usize, coalesce_config: Option, + io_wait: Option>, } impl ObjectStoreReadAt { @@ -72,9 +102,17 @@ impl ObjectStoreReadAt { allocator, concurrency: DEFAULT_CONCURRENCY, coalesce_config: Some(CoalesceConfig::object_storage()), + io_wait: None, } } + /// Mark remote object-store waits using the application's I/O-wait contract. + /// Leave this unset for local and in-memory object stores. + pub fn with_io_wait(mut self, io_wait: Arc) -> Self { + self.io_wait = Some(io_wait); + self + } + /// Set the concurrency for this source. pub fn with_concurrency(mut self, concurrency: usize) -> Self { self.concurrency = concurrency; @@ -94,6 +132,7 @@ async fn read_object_store_range( io_handle: Handle, allocator: BufferAllocatorRef, request: ReadAtRequest, + io_wait: Option>, ) -> VortexResult { let ReadAtRequest { offset, @@ -105,15 +144,22 @@ async fn read_object_store_range( // SAFETY: each return path checks that every byte was initialized. unsafe { buffer.set_len(length) }; - let response = store - .get_opts( - &path, - GetOptions { - range: Some(GetRange::Bounded(range.clone())), - ..Default::default() - }, - ) - .await?; + let get_range = range.clone(); + let get = async move { + store + .get_opts( + &path, + GetOptions { + range: Some(GetRange::Bounded(get_range)), + ..Default::default() + }, + ) + .await + }; + let response = match &io_wait { + Some(marker) => marker.get(get.boxed()).await?, + None => get.await?, + }; let buffer = match response.payload { #[cfg(not(target_arch = "wasm32"))] @@ -128,7 +174,11 @@ async fn read_object_store_range( GetResultPayload::File(..) => { unreachable!("File payload not supported on wasm32") } - GetResultPayload::Stream(mut byte_stream) => { + GetResultPayload::Stream(byte_stream) => { + let mut byte_stream = match &io_wait { + Some(marker) => marker.stream(byte_stream), + None => byte_stream, + }; let mut written = 0usize; while let Some(bytes) = byte_stream.next().await { let bytes = bytes?; @@ -175,12 +225,14 @@ impl VortexReadAt for ObjectStoreReadAt { fn size(&self) -> BoxFuture<'static, VortexResult> { let store = Arc::clone(&self.store); let path = self.path.clone(); + let io_wait = self.io_wait.clone(); async move { - store - .head(&path) - .await - .map(|h| h.size) - .map_err(VortexError::from) + let head = async move { store.head(&path).await }; + let meta = match io_wait { + Some(marker) => marker.head(head.boxed()).await?, + None => head.await?, + }; + Ok(meta.size) } .boxed() } @@ -196,6 +248,7 @@ impl VortexReadAt for ObjectStoreReadAt { let handle = self.handle.clone(); let allocator = self.allocator.clone(); let io_handle = handle.clone(); + let io_wait = self.io_wait.clone(); handle .spawn_io(read_object_store_range( store, @@ -203,6 +256,7 @@ impl VortexReadAt for ObjectStoreReadAt { io_handle, allocator, ReadAtRequest::new(offset, length, alignment), + io_wait, )) .boxed() } @@ -216,6 +270,7 @@ impl VortexReadAt for ObjectStoreReadAt { let path = self.path.clone(); let handle = self.handle.clone(); let allocator = self.allocator.clone(); + let io_wait = self.io_wait.clone(); let concurrency = self.concurrency.max(1); let (mut send, recv) = mpsc::channel(concurrency); let io_handle = handle.clone(); @@ -229,9 +284,12 @@ impl VortexReadAt for ObjectStoreReadAt { let path = path.clone(); let io_handle = io_handle.clone(); let allocator = allocator.clone(); + let io_wait = io_wait.clone(); async move { - let result = - read_object_store_range(store, path, io_handle, allocator, request).await; + let result = read_object_store_range( + store, path, io_handle, allocator, request, io_wait, + ) + .await; (request, result) } }); @@ -277,6 +335,39 @@ mod tests { spawn_io_count: AtomicUsize, } + #[derive(Default)] + struct CountingIoWait { + heads: AtomicUsize, + gets: AtomicUsize, + streams: AtomicUsize, + } + + impl ObjectStoreIoWait for CountingIoWait { + fn head( + &self, + future: BoxFuture<'static, object_store::Result>, + ) -> BoxFuture<'static, object_store::Result> { + self.heads.fetch_add(1, Ordering::SeqCst); + future + } + + fn get( + &self, + future: BoxFuture<'static, object_store::Result>, + ) -> BoxFuture<'static, object_store::Result> { + self.gets.fetch_add(1, Ordering::SeqCst); + future + } + + fn stream( + &self, + stream: BoxStream<'static, object_store::Result>, + ) -> BoxStream<'static, object_store::Result> { + self.streams.fetch_add(1, Ordering::SeqCst); + stream + } + } + impl Executor for CountingExecutor { fn spawn(&self, fut: BoxFuture<'static, ()>) -> AbortHandleRef { self.spawn_count.fetch_add(1, Ordering::SeqCst); @@ -364,4 +455,34 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn io_wait_marker_covers_head_and_both_range_read_paths() -> anyhow::Result<()> { + let executor = Arc::new(CountingExecutor::default()); + let runtime = Arc::clone(&executor) as Arc; + let handle = Handle::new(Arc::downgrade(&runtime)); + let store = Arc::new(InMemory::new()) as Arc; + let path = ObjectPath::from("test.bin"); + store.put(&path, PutPayload::from_static(TEST_DATA)).await?; + + let marker = Arc::new(CountingIoWait::default()); + let reader = ObjectStoreReadAt::new(store, path, handle) + .with_io_wait(Arc::clone(&marker) as Arc); + assert_eq!(reader.size().await?, TEST_DATA.len() as u64); + assert_eq!(reader.read_at(7, 5, Alignment::new(1)).await?.len(), 5); + let requests: Arc<[ReadAtRequest]> = Arc::from([ + ReadAtRequest::new(0, 6, Alignment::new(1)), + ReadAtRequest::new(7, 5, Alignment::new(1)), + ]); + let results = reader.read_ranges(requests).collect::>().await; + assert_eq!(results.len(), 2); + for (_, result) in results { + result?; + } + + assert_eq!(marker.heads.load(Ordering::SeqCst), 1); + assert_eq!(marker.gets.load(Ordering::SeqCst), 3); + assert_eq!(marker.streams.load(Ordering::SeqCst), 3); + Ok(()) + } }