diff --git a/Cargo.lock b/Cargo.lock index cd2662d2372..4c0c45a8b37 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10117,6 +10117,7 @@ dependencies = [ name = "vortex-file" version = "0.1.0" dependencies = [ + "async-stream", "async-trait", "bytes", "codspeed-divan-compat", diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml index fc406f7133e..72a47d88f50 100644 --- a/vortex-file/Cargo.toml +++ b/vortex-file/Cargo.toml @@ -17,6 +17,7 @@ version = { workspace = true } all-features = true [dependencies] +async-stream = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } flatbuffers = { workspace = true } diff --git a/vortex-file/src/read/driver.rs b/vortex-file/src/read/driver.rs index 7f6dc3b2f7c..69d3851d079 100644 --- a/vortex-file/src/read/driver.rs +++ b/vortex-file/src/read/driver.rs @@ -23,20 +23,19 @@ use crate::segments::ReadEvent; use crate::segments::RequestMetrics; pin_project! { - /// A stream that performs coalescing and prioritization of I/O requests. + /// Converts request lifecycle events into batches of physical reads. /// - /// Takes an input stream of [`ReadRequest`]s and buffers all ready requests into local state. - /// When polled for the next request, this stream will choose the next best request based on - /// an ordering of `(has_been_polled, insertion_order)`, skipping any canceled requests, and - /// then coalescing with other nearby requests within the configured `window`. - /// - /// The output of this stream is expected to be buffered by the desired I/O concurrency, and - /// driven to completion. + /// Polled requests become eligible in registration order. An eligible request may absorb nearby + /// registered requests according to `coalesce_window`. Each poll emits every physical read + /// currently available, up to `batch_size`; it never waits for a full batch. pub(crate) struct IoRequestStream { #[pin] events: S, + // True after the event source closes; buffered requests may still remain. inner_done: bool, coalesce_window: Option, + // Maximum physical reads returned by one stream item. + batch_size: usize, state: State, } } @@ -48,15 +47,18 @@ impl IoRequestStream { events: S, coalesce_window: Option, coalesced_buffer_alignment: Alignment, + batch_size: usize, metrics: RequestMetrics, ) -> Self where S: Stream + Unpin + Send + 'static, { + assert!(batch_size > 0, "I/O request batch size must be non-zero"); IoRequestStream { events, inner_done: false, coalesce_window, + batch_size, state: State::new(metrics, coalesced_buffer_alignment), } } @@ -66,12 +68,13 @@ impl Stream for IoRequestStream where S: Stream + Unpin + Send + 'static, { - type Item = IoRequest; + type Item = Vec; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let mut this = self.project(); - // First, try to drain all immediately available requests from the inner stream + // Apply all events already available before choosing work. This gives coalescing visibility + // into registered neighbors without delaying emission for future events. loop { match this.events.as_mut().poll_next(cx) { Poll::Ready(Some(event)) => { @@ -87,17 +90,24 @@ where } } - // Try to get a coalesced request - if let Some(coalesced) = this.state.next(this.coalesce_window.as_ref()) { - return Poll::Ready(Some(coalesced)); + // Emit a partial batch immediately so the downstream driver can fill free I/O slots. + let mut batch = Vec::with_capacity(*this.batch_size); + while batch.len() < *this.batch_size { + let Some(request) = this.state.next(this.coalesce_window.as_ref()) else { + break; + }; + batch.push(request); + } + if !batch.is_empty() { + return Poll::Ready(Some(batch)); } - // If the inner stream is done, and we have no more _polled_ requests, we're done + // Unpolled requests cannot initiate I/O, so a closed source is done once none are eligible. if *this.inner_done && this.state.polled_requests.is_empty() { return Poll::Ready(None); } - // Otherwise, we need more data from the inner stream + // A new poll/drop/register event will wake us. Poll::Pending } } @@ -326,10 +336,12 @@ impl State { #[cfg(test)] mod tests { use futures::StreamExt; + use futures::channel::mpsc; use futures::stream; use vortex_array::buffer::BufferHandle; use vortex_buffer::Alignment; use vortex_error::VortexResult; + use vortex_error::vortex_panic; use vortex_metrics::DefaultMetricsRegistry; use vortex_metrics::MetricValue; use vortex_metrics::MetricsRegistry; @@ -374,9 +386,10 @@ mod tests { event_stream, coalesce_window, coalesced_buffer_alignment, + 1024, metrics, ); - io_stream.collect().await + io_stream.concat().await } #[tokio::test] @@ -415,6 +428,65 @@ mod tests { assert_eq!(offsets, vec![0, 100, 200]); // req1, req2, req3 } + #[tokio::test] + async fn test_bounded_request_batches() { + let mut events = Vec::new(); + let mut receivers = Vec::new(); + for id in 0..5 { + let (request, recv) = create_request(id, id as u64 * 10, 10); + events.push(ReadEvent::Request(request)); + events.push(ReadEvent::Polled(id)); + receivers.push(recv); + } + + let metrics_registry = DefaultMetricsRegistry::default(); + let metrics = RequestMetrics::new(&metrics_registry, vec![]); + let batches = + IoRequestStream::new(stream::iter(events), None, Alignment::none(), 2, metrics) + .collect::>() + .await; + + assert_eq!(receivers.len(), 5); + assert_eq!(batches.iter().map(Vec::len).collect::>(), [2, 2, 1]); + assert_eq!( + batches + .into_iter() + .flatten() + .map(|request| request.offset()) + .collect::>(), + [0, 10, 20, 30, 40] + ); + } + + #[test] + fn test_partial_batch_emits_without_waiting_for_more_events() { + let (sender, receiver) = mpsc::unbounded(); + let (request, _recv) = create_request(1, 0, 10); + assert!(sender.unbounded_send(ReadEvent::Request(request)).is_ok()); + assert!(sender.unbounded_send(ReadEvent::Polled(1)).is_ok()); + + let metrics_registry = DefaultMetricsRegistry::default(); + let metrics = RequestMetrics::new(&metrics_registry, vec![]); + let mut batches = Box::pin(IoRequestStream::new( + receiver, + None, + Alignment::none(), + 32, + metrics, + )); + + // Keep `sender` alive: the input is pending, not finished, and the partial batch must still + // be returned by the current poll rather than waiting for 31 more requests. + let waker = futures::task::noop_waker(); + let mut context = Context::from_waker(&waker); + let Poll::Ready(Some(batch)) = batches.as_mut().poll_next(&mut context) else { + vortex_panic!("partial batch was not emitted by the current poll"); + }; + assert_eq!(batch.len(), 1); + assert_eq!(batch[0].offset(), 0); + drop(sender); + } + #[tokio::test] async fn test_coalesce_adjacent() { let (req1, _rx1) = create_request(1, 0, 10); @@ -734,10 +806,11 @@ mod tests { max_size: 1024, }), Alignment::none(), + 1024, metrics, ); - let outputs: Vec = io_stream.collect().await; + let outputs: Vec = io_stream.concat().await; assert_eq!(outputs.len(), 2); let snapshot = metrics_registry.snapshot(); @@ -788,9 +861,9 @@ mod tests { let metrics_registry = DefaultMetricsRegistry::default(); let metrics = RequestMetrics::new(&metrics_registry, vec![]); // No coalescing window - should be individual requests - let io_stream = IoRequestStream::new(event_stream, None, Alignment::none(), metrics); + let io_stream = IoRequestStream::new(event_stream, None, Alignment::none(), 1024, metrics); - let outputs: Vec = io_stream.collect().await; + let outputs: Vec = io_stream.concat().await; assert_eq!(outputs.len(), 2); // Check metrics diff --git a/vortex-file/src/read/mod.rs b/vortex-file/src/read/mod.rs index a812b81f63b..f1b18e9a5b1 100644 --- a/vortex-file/src/read/mod.rs +++ b/vortex-file/src/read/mod.rs @@ -5,5 +5,6 @@ mod driver; mod request; pub(crate) use driver::IoRequestStream; +pub(crate) use request::IoRequest; pub(crate) use request::ReadRequest; pub(crate) use request::RequestId; diff --git a/vortex-file/src/segments/source.rs b/vortex-file/src/segments/source.rs index 3af33362b05..33aed2fa9d2 100644 --- a/vortex-file/src/segments/source.rs +++ b/vortex-file/src/segments/source.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::any::Any; +use std::collections::VecDeque; use std::future::Future; use std::pin::Pin; use std::sync::Arc; @@ -11,19 +12,25 @@ use std::task::Context; use std::task::Poll; use futures::FutureExt; +use futures::Stream; use futures::StreamExt; use futures::channel::mpsc; use futures::future; use futures::future::BoxFuture; use futures::future::Shared; +use futures::stream::BoxStream; +use futures::stream::Fuse; +use futures::stream::SelectAll; use parking_lot::Mutex; use vortex_array::buffer::BufferHandle; use vortex_buffer::Alignment; use vortex_buffer::ByteBuffer; +use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_error::vortex_panic; +use vortex_io::ReadAtRequest; +use vortex_io::ReadAtStream; use vortex_io::VortexReadAt; use vortex_io::runtime::Handle; use vortex_io::runtime::JoinOutcome; @@ -37,6 +44,7 @@ use vortex_metrics::MetricBuilder; use vortex_metrics::MetricsRegistry; use crate::SegmentSpec; +use crate::read::IoRequest; use crate::read::IoRequestStream; use crate::read::ReadRequest; use crate::read::RequestId; @@ -85,6 +93,188 @@ type SharedDriver = Shared>; /// observe completion takes the payload and re-raises it; later readers report a graceful error. type DriverPanic = Arc>>>; +fn validate_read_result( + request: &IoRequest, + result: VortexResult, +) -> VortexResult { + result.and_then(|buffer| { + if request.len() != buffer.len() { + return Err(vortex_err!( + "FileSegmentSource: expected buffer of length {} but received {}. {:?}", + request.len(), + buffer.len(), + request + )); + } + Ok(buffer) + }) +} + +type IoBatchStream = Fuse>>; + +enum ReadRangeResultsState { + Reading(ReadAtStream), + Missing, +} + +/// Matches streamed range results back to their logical requests. +struct ReadRangeResults { + state: ReadRangeResultsState, + remaining: Vec>, +} + +impl ReadRangeResults { + fn new(results: ReadAtStream, requests: Vec) -> Self { + Self { + state: ReadRangeResultsState::Reading(results), + remaining: requests.into_iter().map(Some).collect(), + } + } +} + +impl Stream for ReadRangeResults { + type Item = (IoRequest, VortexResult); + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + loop { + match &mut self.state { + ReadRangeResultsState::Reading(results) => match results.poll_next_unpin(cx) { + Poll::Ready(Some((request, result))) => { + let Some(position) = self.remaining.iter().position(|req| { + req.as_ref().is_some_and(|req| { + req.offset() == request.offset + && req.len() == request.length + && req.alignment() == request.alignment + }) + }) else { + tracing::warn!(?request, "reader returned an unknown range"); + continue; + }; + let req = self.remaining[position] + .take() + .vortex_expect("matched request is present"); + return Poll::Ready(Some((req, result))); + } + Poll::Ready(None) => self.state = ReadRangeResultsState::Missing, + Poll::Pending => return Poll::Pending, + }, + ReadRangeResultsState::Missing => { + let Some(req) = self.remaining.iter_mut().find_map(Option::take) else { + return Poll::Ready(None); + }; + let error = vortex_err!( + "FileSegmentSource: read_ranges ended before resolving request. {:?}", + req + ); + return Poll::Ready(Some((req, Err(error)))); + } + } + } + } +} + +/// Drives request batches while keeping the reader's concurrency slots occupied. +struct ReadDriver { + reader: Arc, + batches: IoBatchStream, + pending: VecDeque, + reads: SelectAll, + num_active: usize, + batches_done: bool, + concurrency: usize, + metrics: RequestMetrics, +} + +impl ReadDriver { + fn new( + reader: R, + batches: BoxStream<'static, Vec>, + concurrency: usize, + metrics: RequestMetrics, + ) -> Self { + Self { + reader: Arc::new(reader), + batches: batches.fuse(), + pending: VecDeque::new(), + reads: SelectAll::new(), + num_active: 0, + batches_done: false, + concurrency, + metrics, + } + } + + fn submit_pending(&mut self) { + while self.num_active < self.concurrency && !self.pending.is_empty() { + let batch_len = (self.concurrency - self.num_active).min(self.pending.len()); + let reqs = self.pending.drain(..batch_len).collect::>(); + self.num_active += batch_len; + + self.metrics.read_ranges_calls.add(1); + self.metrics.read_ranges_num_ranges.update(batch_len as f64); + if batch_len > 1 { + self.metrics.read_ranges_multi.add(1); + } + tracing::trace!( + target: "vortex_file::read_ranges", + num_ranges = batch_len, + num_active = self.num_active, + "submitting positional read batch" + ); + + let requests = reqs + .iter() + .map(|req| ReadAtRequest::new(req.offset(), req.len(), req.alignment())) + .collect::>() + .into(); + let results = self.reader.read_ranges(requests); + self.reads.push(ReadRangeResults::new(results, reqs)); + } + } +} + +impl Stream for ReadDriver { + type Item = (); + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.as_mut().get_mut(); + + // Observe every batch already available so submission can fill all free slots at once. + if !this.batches_done { + loop { + match this.batches.poll_next_unpin(cx) { + Poll::Ready(Some(batch)) => this.pending.extend(batch), + Poll::Ready(None) => { + this.batches_done = true; + break; + } + Poll::Pending => break, + } + } + } + + this.submit_pending(); + + if this.batches_done && this.num_active == 0 { + return Poll::Ready(None); + } + + match this.reads.poll_next_unpin(cx) { + Poll::Ready(Some((req, result))) => { + this.num_active -= 1; + let result = validate_read_result(&req, result); + req.resolve(result); + Poll::Ready(Some(())) + } + Poll::Ready(None) if this.num_active == 0 => Poll::Pending, + Poll::Ready(None) => { + vortex_panic!("read result streams ended with active requests") + } + Poll::Pending => Poll::Pending, + } + } +} + pub struct FileSegmentSource { segments: Arc<[SegmentSpec]>, /// A queue for sending read request events to the I/O stream. @@ -134,37 +324,12 @@ impl FileSegmentSource { StreamExt::boxed(recv), coalesce_config, max_alignment, - metrics, + concurrency, + metrics.clone(), ) .boxed(); - let drive_fut = async move { - stream - .map(move |req| { - let reader = reader.clone(); - async move { - let result = reader - .read_at(req.offset(), req.len(), req.alignment()) - .await; - let result = result.and_then(|buffer| { - if req.len() != buffer.len() { - vortex_bail!( - "FileSegmentSource: expected buffer of length {} but received {}. {:?}", - req.len(), - buffer.len(), - req - ) - } - Ok(buffer) - }); - - req.resolve(result); - } - }) - .buffer_unordered(concurrency) - .collect::<()>() - .await - }; + let drive_fut = ReadDriver::new(reader, stream, concurrency, metrics).collect::<()>(); // Spawn the driver so the runtime makes I/O progress independently of any reader. Readers // join it (below) only to surface a panic raised while driving reads. @@ -309,6 +474,7 @@ impl Drop for ReadFuture { } /// Metrics emitted by the file segment request driver. +#[derive(Clone)] pub struct RequestMetrics { /// Number of individual segment requests observed by the driver. pub individual_requests: Counter, @@ -316,6 +482,12 @@ pub struct RequestMetrics { pub coalesced_requests: Counter, /// Distribution of how many segment requests were merged into each physical read. pub num_requests_coalesced: Histogram, + /// Number of calls made to [`VortexReadAt::read_ranges`]. + pub read_ranges_calls: Counter, + /// Number of `read_ranges` calls containing more than one physical range. + pub read_ranges_multi: Counter, + /// Distribution of physical range counts submitted per `read_ranges` call. + pub read_ranges_num_ranges: Histogram, } impl RequestMetrics { @@ -329,8 +501,17 @@ impl RequestMetrics { .add_labels(labels.clone()) .counter("io.requests.coalesced"), num_requests_coalesced: MetricBuilder::new(metrics_registry) - .add_labels(labels) + .add_labels(labels.clone()) .histogram("io.requests.coalesced.num_coalesced"), + read_ranges_calls: MetricBuilder::new(metrics_registry) + .add_labels(labels.clone()) + .counter("io.read_ranges.calls"), + read_ranges_multi: MetricBuilder::new(metrics_registry) + .add_labels(labels.clone()) + .counter("io.read_ranges.multi_range_calls"), + read_ranges_num_ranges: MetricBuilder::new(metrics_registry) + .add_labels(labels) + .histogram("io.read_ranges.num_ranges"), } } } @@ -383,12 +564,44 @@ mod tests { use std::panic::AssertUnwindSafe; use futures::future::BoxFuture; + use vortex_error::vortex_bail; use vortex_io::runtime::tokio::TokioRuntime; use vortex_layout::segments::SegmentSource; use vortex_metrics::DefaultMetricsRegistry; use super::*; + fn io_request(id: RequestId, offset: u64, length: usize) -> IoRequest { + let (callback, _receiver) = oneshot::channel(); + IoRequest::new_single(ReadRequest { + id, + offset, + length, + alignment: Alignment::none(), + callback, + }) + } + + #[tokio::test] + async fn read_range_results_matches_results_and_reports_missing_requests() { + let requests = vec![io_request(0, 0, 4), io_request(1, 4, 4)]; + let unknown = ReadAtRequest::new(8, 4, Alignment::none()); + let returned = ReadAtRequest::new(4, 4, Alignment::none()); + let buffer = BufferHandle::new_host(ByteBuffer::from(vec![0; 4])); + let results = + futures::stream::iter([(unknown, Ok(buffer.clone())), (returned, Ok(buffer))]).boxed(); + + let resolved = ReadRangeResults::new(results, requests) + .collect::>() + .await; + + assert_eq!(resolved.len(), 2); + assert_eq!(resolved[0].0.offset(), 4); + assert!(resolved[0].1.is_ok()); + assert_eq!(resolved[1].0.offset(), 0); + assert!(resolved[1].1.is_err()); + } + #[derive(Clone)] struct PanickingReadAt; @@ -510,6 +723,265 @@ mod tests { ); } + #[derive(Clone)] + struct ReadRangesOnly { + calls: Arc, + } + + impl VortexReadAt for ReadRangesOnly { + fn concurrency(&self) -> usize { + 4 + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + async { Ok(16) }.boxed() + } + + fn read_at( + &self, + _offset: u64, + _length: usize, + _alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + async { panic!("read_at should not be called") }.boxed() + } + + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { + self.calls.fetch_add(1, Ordering::Relaxed); + let results = requests + .iter() + .copied() + .map(|request| { + let buffer = BufferHandle::new_host( + ByteBuffer::from(vec![0; request.length]).aligned(request.alignment), + ); + (request, Ok(buffer)) + }) + .collect::>(); + futures::stream::iter(results).boxed() + } + } + + #[tokio::test] + async fn read_driver_batches_ready_requests() -> VortexResult<()> { + let calls = Arc::new(AtomicUsize::new(0)); + let segments: Arc<[SegmentSpec]> = (0..4) + .map(|i| SegmentSpec { + offset: i * 4, + length: 4, + alignment: Alignment::none(), + }) + .collect(); + let metrics = DefaultMetricsRegistry::default(); + let request_metrics = RequestMetrics::new(&metrics, vec![]); + let source = FileSegmentSource::open( + segments, + ReadRangesOnly { + calls: Arc::clone(&calls), + }, + TokioRuntime::current(), + request_metrics.clone(), + ); + + let results = future::join_all((0..4).map(|i| source.request(SegmentId::from(i)))).await; + + for result in results { + assert_eq!(result?.len(), 4); + } + assert_eq!(calls.load(Ordering::Relaxed), 1); + assert_eq!(request_metrics.read_ranges_calls.value(), 1); + assert_eq!(request_metrics.read_ranges_multi.value(), 1); + assert_eq!(request_metrics.read_ranges_num_ranges.count(), 1); + assert_eq!(request_metrics.read_ranges_num_ranges.total(), 4.0); + Ok(()) + } + + #[derive(Clone)] + struct ControlledReadRanges { + active: Arc, + max_active: Arc, + batch_sizes: Arc>>, + permits: Arc, + } + + impl VortexReadAt for ControlledReadRanges { + fn concurrency(&self) -> usize { + 4 + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + async { Ok(24) }.boxed() + } + + fn read_at( + &self, + _offset: u64, + _length: usize, + _alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + async { panic!("read_at should not be called") }.boxed() + } + + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { + self.batch_sizes.lock().push(requests.len()); + let active = self.active.fetch_add(requests.len(), Ordering::SeqCst) + requests.len(); + self.max_active.fetch_max(active, Ordering::SeqCst); + + let reads = requests + .iter() + .copied() + .map(|request| { + let active = Arc::clone(&self.active); + let permits = Arc::clone(&self.permits); + async move { + let Ok(permit) = permits.acquire_owned().await else { + vortex_panic!("test semaphore unexpectedly closed"); + }; + permit.forget(); + active.fetch_sub(1, Ordering::SeqCst); + let buffer = BufferHandle::new_host( + ByteBuffer::from(vec![0; request.length]).aligned(request.alignment), + ); + (request, Ok(buffer)) + } + }) + .collect::>(); + futures::stream::iter(reads).buffer_unordered(4).boxed() + } + } + + #[tokio::test] + async fn read_driver_refills_global_concurrency_across_batches() -> VortexResult<()> { + let active = Arc::new(AtomicUsize::new(0)); + let max_active = Arc::new(AtomicUsize::new(0)); + let batch_sizes = Arc::new(Mutex::new(Vec::new())); + let permits = Arc::new(tokio::sync::Semaphore::new(0)); + let segments: Arc<[SegmentSpec]> = (0..6) + .map(|i| SegmentSpec { + offset: i * 4, + length: 4, + alignment: Alignment::none(), + }) + .collect(); + let metrics = DefaultMetricsRegistry::default(); + let source = FileSegmentSource::open( + segments, + ControlledReadRanges { + active: Arc::clone(&active), + max_active: Arc::clone(&max_active), + batch_sizes: Arc::clone(&batch_sizes), + permits: Arc::clone(&permits), + }, + TokioRuntime::current(), + RequestMetrics::new(&metrics, vec![]), + ); + let reads = TokioRuntime::current().spawn(async move { + future::join_all((0..6).map(|i| source.request(SegmentId::from(i)))).await + }); + + assert!( + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while active.load(Ordering::SeqCst) != 4 { + tokio::task::yield_now().await; + } + }) + .await + .is_ok() + ); + + permits.add_permits(1); + assert!( + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while batch_sizes.lock().len() < 2 || active.load(Ordering::SeqCst) != 4 { + tokio::task::yield_now().await; + } + }) + .await + .is_ok() + ); + assert_eq!(batch_sizes.lock().as_slice(), [4, 1]); + assert_eq!(max_active.load(Ordering::SeqCst), 4); + + permits.add_permits(5); + for result in reads.await { + assert_eq!(result?.len(), 4); + } + assert_eq!(max_active.load(Ordering::SeqCst), 4); + Ok(()) + } + + #[tokio::test] + async fn read_driver_keeps_slots_full_while_a_straggler_is_in_flight() -> VortexResult<()> { + let active = Arc::new(AtomicUsize::new(0)); + let max_active = Arc::new(AtomicUsize::new(0)); + let batch_sizes = Arc::new(Mutex::new(Vec::new())); + let permits = Arc::new(tokio::sync::Semaphore::new(0)); + let segments: Arc<[SegmentSpec]> = (0..8) + .map(|i| SegmentSpec { + offset: i * 4, + length: 4, + alignment: Alignment::none(), + }) + .collect(); + let metrics = DefaultMetricsRegistry::default(); + let source = FileSegmentSource::open( + segments, + ControlledReadRanges { + active: Arc::clone(&active), + max_active: Arc::clone(&max_active), + batch_sizes: Arc::clone(&batch_sizes), + permits: Arc::clone(&permits), + }, + TokioRuntime::current(), + RequestMetrics::new(&metrics, vec![]), + ); + let reads = TokioRuntime::current().spawn(async move { + future::join_all((0..8).map(|i| source.request(SegmentId::from(i)))).await + }); + + wait_for_active_reads(&active, 4).await; + + // Complete three reads while leaving one original read blocked as a straggler. Each freed + // slot must be refilled before the next completion; a batch-barrier implementation would + // instead fall from four active reads to one and submit no replacement work. + for expected_calls in 2..=4 { + permits.add_permits(1); + assert!( + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while batch_sizes.lock().len() < expected_calls + || active.load(Ordering::SeqCst) != 4 + { + tokio::task::yield_now().await; + } + }) + .await + .is_ok() + ); + } + + assert_eq!(batch_sizes.lock().as_slice(), [4, 1, 1, 1]); + assert_eq!(active.load(Ordering::SeqCst), 4); + assert_eq!(max_active.load(Ordering::SeqCst), 4); + + permits.add_permits(5); + for result in reads.await { + assert_eq!(result?.len(), 4); + } + Ok(()) + } + + async fn wait_for_active_reads(active: &AtomicUsize, expected: usize) { + assert!( + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while active.load(Ordering::SeqCst) != expected { + tokio::task::yield_now().await; + } + }) + .await + .is_ok() + ); + } + #[derive(Clone)] struct SlowErrReadAt; diff --git a/vortex-io/src/compat/read_at.rs b/vortex-io/src/compat/read_at.rs index 4fc49785d28..3d9cc93b1a6 100644 --- a/vortex-io/src/compat/read_at.rs +++ b/vortex-io/src/compat/read_at.rs @@ -4,12 +4,15 @@ use std::sync::Arc; use futures::FutureExt; +use futures::StreamExt; use futures::future::BoxFuture; use vortex_array::buffer::BufferHandle; use vortex_buffer::Alignment; use vortex_error::VortexResult; use crate::CoalesceConfig; +use crate::ReadAtRequest; +use crate::ReadAtStream; use crate::VortexReadAt; use crate::compat::Compat; @@ -40,4 +43,8 @@ impl VortexReadAt for Compat { ) -> BoxFuture<'static, VortexResult> { Compat::new(self.inner().read_at(offset, length, alignment)).boxed() } + + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { + Compat::new(self.inner().read_ranges(requests)).boxed() + } } diff --git a/vortex-io/src/object_store/read_at.rs b/vortex-io/src/object_store/read_at.rs index 086d70c1bcf..462bc498f82 100644 --- a/vortex-io/src/object_store/read_at.rs +++ b/vortex-io/src/object_store/read_at.rs @@ -5,8 +5,11 @@ use std::io; use std::sync::Arc; use futures::FutureExt; +use futures::SinkExt; use futures::StreamExt; +use futures::channel::mpsc; use futures::future::BoxFuture; +use futures::stream; use object_store::GetOptions; use object_store::GetRange; use object_store::GetResultPayload; @@ -22,6 +25,8 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use crate::CoalesceConfig; +use crate::ReadAtRequest; +use crate::ReadAtStream; use crate::VortexReadAt; use crate::runtime::Handle; #[cfg(not(target_arch = "wasm32"))] @@ -79,6 +84,75 @@ impl ObjectStoreReadAt { } } +async fn read_object_store_range( + store: Arc, + path: ObjectPath, + io_handle: Handle, + allocator: HostAllocatorRef, + request: ReadAtRequest, +) -> VortexResult { + let ReadAtRequest { + offset, + length, + alignment, + } = request; + let range = offset..(offset + length as u64); + let mut buffer = allocator.allocate(length, alignment)?; + + let response = store + .get_opts( + &path, + GetOptions { + range: Some(GetRange::Bounded(range.clone())), + ..Default::default() + }, + ) + .await?; + + let buffer = match response.payload { + #[cfg(not(target_arch = "wasm32"))] + GetResultPayload::File(file, _) => io_handle + .spawn_blocking(move || { + read_exact_at(&file, buffer.as_mut_slice(), range.start)?; + Ok::<_, io::Error>(buffer) + }) + .await + .map_err(io::Error::other)?, + #[cfg(target_arch = "wasm32")] + GetResultPayload::File(..) => { + unreachable!("File payload not supported on wasm32") + } + GetResultPayload::Stream(mut byte_stream) => { + let mut written = 0usize; + while let Some(bytes) = byte_stream.next().await { + let bytes = bytes?; + let end = written + bytes.len(); + vortex_ensure!( + end <= length, + "Object store stream returned too many bytes: {} > expected {} (range: {:?})", + end, + length, + range + ); + buffer.as_mut_slice()[written..end].copy_from_slice(&bytes); + written = end; + } + + vortex_ensure!( + written == length, + "Object store stream returned {} bytes but expected {} bytes (range: {:?})", + written, + length, + range + ); + + buffer + } + }; + + Ok(BufferHandle::new_host(buffer.freeze())) +} + impl VortexReadAt for ObjectStoreReadAt { fn uri(&self) -> Option<&Arc> { Some(&self.uri) @@ -115,70 +189,62 @@ impl VortexReadAt for ObjectStoreReadAt { let path = self.path.clone(); let handle = self.handle.clone(); let allocator = Arc::clone(&self.allocator); - let range = offset..(offset + length as u64); + let io_handle = handle.clone(); + handle + .spawn_io(read_object_store_range( + store, + path, + io_handle, + allocator, + ReadAtRequest::new(offset, length, alignment), + )) + .boxed() + } - // Requires to deal with borrowed lifetimes + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { + if requests.is_empty() { + return stream::empty().boxed(); + } + + let store = Arc::clone(&self.store); + let path = self.path.clone(); + let handle = self.handle.clone(); + let allocator = Arc::clone(&self.allocator); + let concurrency = self.concurrency.max(1); + let (mut send, recv) = mpsc::channel(concurrency); let io_handle = handle.clone(); - handle - .spawn_io(async move { - let mut buffer = allocator.allocate(length, alignment)?; - - let response = store - .get_opts( - &path, - GetOptions { - range: Some(GetRange::Bounded(range.clone())), - ..Default::default() - }, - ) - .await?; - - let buffer = match response.payload { - #[cfg(not(target_arch = "wasm32"))] - GetResultPayload::File(file, _) => { - io_handle - .spawn_blocking(move || { - read_exact_at(&file, buffer.as_mut_slice(), range.start)?; - Ok::<_, io::Error>(buffer) - }) - .await - .map_err(io::Error::other)? - } - #[cfg(target_arch = "wasm32")] - GetResultPayload::File(..) => { - unreachable!("File payload not supported on wasm32") - } - GetResultPayload::Stream(mut byte_stream) => { - let mut written = 0usize; - while let Some(bytes) = byte_stream.next().await { - let bytes = bytes?; - let end = written + bytes.len(); - vortex_ensure!( - end <= length, - "Object store stream returned too many bytes: {} > expected {} (range: {:?})", - end, - length, - range - ); - buffer.as_mut_slice()[written..end].copy_from_slice(&bytes); - written = end; - } - - vortex_ensure!( - written == length, - "Object store stream returned {} bytes but expected {} bytes (range: {:?})", - written, - length, - range - ); - - buffer - } - }; - - Ok(BufferHandle::new_host(buffer.freeze())) - }) + // A single runtime task drives all GETs, avoiding one spawn per range. Do not use + // ObjectStore::get_ranges here: it returns one Vec after every range completes, whereas + // VortexReadAt::read_ranges must expose each result as soon as it is ready. + let task = handle.spawn_io(async move { + let reads = requests.iter().copied().map(|request| { + let store = Arc::clone(&store); + let path = path.clone(); + let io_handle = io_handle.clone(); + let allocator = Arc::clone(&allocator); + async move { + let result = + read_object_store_range(store, path, io_handle, allocator, request).await; + (request, result) + } + }); + + let mut reads = stream::iter(reads).buffer_unordered(concurrency); + while let Some(result) = reads.next().await { + if send.send(result).await.is_err() { + break; + } + } + }); + + async_stream::stream! { + let mut recv = recv; + while let Some(result) = recv.next().await { + yield result; + } + task.await; + } .boxed() } } @@ -258,4 +324,38 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn read_ranges_uses_one_io_task() -> 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 reader = ObjectStoreReadAt::new(store, path, handle); + let requests: Arc<[ReadAtRequest]> = Arc::from([ + ReadAtRequest::new(0, 6, Alignment::new(1)), + ReadAtRequest::new(7, 5, Alignment::new(1)), + ReadAtRequest::new(18, 4, Alignment::new(1)), + ]); + let results = reader.read_ranges(requests).collect::>().await; + + assert_eq!(results.len(), 3); + for (request, result) in results { + let buffer = result?; + let offset = usize::try_from(request.offset)?; + assert_eq!(buffer.len(), request.length); + assert_eq!( + buffer.to_host().await.as_slice(), + &TEST_DATA[offset..offset + request.length] + ); + } + assert_eq!(executor.spawn_io_count.load(Ordering::SeqCst), 1); + assert_eq!(executor.spawn_count.load(Ordering::SeqCst), 0); + + Ok(()) + } } diff --git a/vortex-io/src/read_at.rs b/vortex-io/src/read_at.rs index aa9a8a03abf..82190bdd699 100644 --- a/vortex-io/src/read_at.rs +++ b/vortex-io/src/read_at.rs @@ -2,9 +2,13 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::sync::Arc; +use std::time::Instant; use futures::FutureExt; +use futures::StreamExt; use futures::future::BoxFuture; +use futures::stream; +use futures::stream::BoxStream; use vortex_array::buffer::BufferHandle; use vortex_buffer::Alignment; use vortex_buffer::ByteBuffer; @@ -27,6 +31,31 @@ pub struct CoalesceConfig { pub max_size: u64, } +/// A positional read request used by [`VortexReadAt::read_ranges`]. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ReadAtRequest { + /// The byte offset at which to start reading. + pub offset: u64, + /// The exact number of bytes to read. + pub length: usize, + /// The required alignment of the returned buffer. + pub alignment: Alignment, +} + +impl ReadAtRequest { + /// Creates a positional read request. + pub const fn new(offset: u64, length: usize, alignment: Alignment) -> Self { + Self { + offset, + length, + alignment, + } + } +} + +/// A stream of positional read results, yielded as each request completes. +pub type ReadAtStream = BoxStream<'static, (ReadAtRequest, VortexResult)>; + impl CoalesceConfig { /// Creates a new coalesce configuration. pub const fn new(distance: u64, max_size: u64) -> Self { @@ -89,6 +118,25 @@ pub trait VortexReadAt: Send + Sync + 'static { length: usize, alignment: Alignment, ) -> BoxFuture<'static, VortexResult>; + + /// Request multiple asynchronous positional reads. + /// + /// Each item includes its request and result, and is yielded as soon as that read completes. + /// A failed request does not prevent other results from being yielded. Callers should submit + /// batches no larger than [`VortexReadAt::concurrency`]. + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { + let reads = requests + .iter() + .copied() + .map(|request| { + let read = self.read_at(request.offset, request.length, request.alignment); + async move { (request, read.await) } + }) + .collect::>(); + stream::iter(reads) + .buffer_unordered(self.concurrency().max(1)) + .boxed() + } } impl VortexReadAt for Arc { @@ -116,6 +164,10 @@ impl VortexReadAt for Arc { ) -> BoxFuture<'static, VortexResult> { self.as_ref().read_at(offset, length, alignment) } + + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { + self.as_ref().read_ranges(requests) + } } impl VortexReadAt for Arc { @@ -143,6 +195,10 @@ impl VortexReadAt for Arc { ) -> BoxFuture<'static, VortexResult> { self.as_ref().read_at(offset, length, alignment) } + + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { + self.as_ref().read_ranges(requests) + } } impl VortexReadAt for ByteBuffer { @@ -316,17 +372,63 @@ impl VortexReadAt for InstrumentedReadAt { } .boxed() } + + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { + let durations = self.metrics.durations.clone(); + let sizes = self.metrics.sizes.clone(); + let total_size = self.metrics.total_size.clone(); + let start = Instant::now(); + self.read + .read_ranges(requests) + .map(move |(request, result)| { + durations.update(start.elapsed()); + sizes.update(request.length as f64); + total_size.add(request.length as u64); + (request, result) + }) + .boxed() + } } #[cfg(test)] mod tests { use std::sync::Arc; + use std::time::Duration; use vortex_buffer::Alignment; use vortex_buffer::ByteBuffer; use super::*; + struct DelayedReadAt; + + impl VortexReadAt for DelayedReadAt { + fn concurrency(&self) -> usize { + 2 + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + async { Ok(2) }.boxed() + } + + fn read_at( + &self, + offset: u64, + _length: usize, + _alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + async move { + if offset == 0 { + tokio::time::sleep(Duration::from_millis(50)).await; + } + Ok(BufferHandle::new_host(ByteBuffer::from(vec![ + u8::try_from(offset).vortex_expect("test offset fits in u8"), + ]))) + } + .boxed() + } + } + #[test] fn test_coalesce_config_in_memory() { let config = CoalesceConfig::in_memory(); @@ -356,6 +458,67 @@ mod tests { assert_eq!(result.to_host().await.as_ref(), &[2, 3, 4]); } + #[tokio::test] + async fn test_byte_buffer_read_ranges() -> VortexResult<()> { + let data = ByteBuffer::from(vec![1, 2, 3, 4, 5, 6]); + let requests = Arc::from([ + ReadAtRequest::new(4, 2, Alignment::none()), + ReadAtRequest::new(0, 1, Alignment::none()), + ReadAtRequest::new(2, 3, Alignment::none()), + ]); + + let results = data.read_ranges(requests).collect::>().await; + for (request, result) in results { + let expected: &[u8] = match request.offset { + 0 => &[1], + 2 => &[3, 4, 5], + 4 => &[5, 6], + offset => panic!("unexpected offset: {offset}"), + }; + assert_eq!(result?.to_host().await.as_ref(), expected); + } + Ok(()) + } + + #[tokio::test] + async fn test_read_ranges_keeps_streaming_after_an_error() -> VortexResult<()> { + let data = ByteBuffer::from(vec![1, 2, 3]); + let requests = Arc::from([ + ReadAtRequest::new(100, 1, Alignment::none()), + ReadAtRequest::new(1, 2, Alignment::none()), + ]); + + let results = data.read_ranges(requests).collect::>().await; + + assert_eq!(results.len(), 2); + assert!(results.iter().any(|(_, result)| result.is_err())); + let (_, valid) = results + .into_iter() + .find(|(request, _)| request.offset == 1) + .vortex_expect("valid request result is present"); + assert_eq!(valid?.to_host().await.as_ref(), &[2, 3]); + Ok(()) + } + + #[tokio::test] + async fn test_read_ranges_yields_in_completion_order() -> VortexResult<()> { + let requests = Arc::from([ + ReadAtRequest::new(0, 1, Alignment::none()), + ReadAtRequest::new(1, 1, Alignment::none()), + ]); + let mut results = DelayedReadAt.read_ranges(requests); + + let (first_request, first_result) = results + .next() + .await + .vortex_expect("first result is present"); + + assert_eq!(first_request.offset, 1); + assert_eq!(first_result?.to_host().await.as_ref(), &[1]); + assert_eq!(results.count().await, 1); + Ok(()) + } + #[tokio::test] async fn test_byte_buffer_read_out_of_bounds() { let data = ByteBuffer::from(vec![1, 2, 3]);