diff --git a/SPICE_PATCHES.md b/SPICE_PATCHES.md index 5b8e38c6fc2..b06f30c33ba 100644 --- a/SPICE_PATCHES.md +++ b/SPICE_PATCHES.md @@ -62,6 +62,8 @@ exists: a grep proves presence, a test proves behaviour. | 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 | Byte-capped data blocks for string/binary columns, with encodings chosen once per row block | this change | A string or binary column is written in blocks of 8,192 rows whatever its row width, because `effective_block_len` sizes blocks only for fixed-width types; a ~3 KB/row JSON column gets ~12.5 MiB blocks, and a one-row lookup reads and copies a whole block (spiceai/spiceai#15323). `WriteStrategyBuilder::with_data_block_max_bytes` splits such a block into power-of-two row counts under the cap. Searching encodings on each small block costs several full compressions per block (a block of at most 1,024 rows is its own sample), so `CompressingStrategy::with_group_rows` and `CascadingCompressor::plan_for`/`compress_with_plan` choose encodings once per row block, on the sample the compressor would draw from the unsplit block, and apply them to its split blocks (a choice made on the first split block alone fits only it: incompressible rows ahead of JSON picked FSST for the JSON, +15% size and 2.4x `LIKE` CPU) | `cargo test -p vortex-layout --lib repartition::tests::a_block_size_maximum_splits_wide_strings_into_power_of_two_blocks` **and** `cargo test -p vortex-btrblocks --lib canonical_compressor::tests::a_replayed_plan_roundtrips_other_data_of_the_same_dtype` **and** `cargo test -p vortex-layout --lib compressed::tests::a_group_is_encoded_as_its_unsplit_data_would_be` **and** `cargo test -p vortex-file --test test_write_table test_data_block_max_bytes_roundtrip` | Proposed — additive; off unless the cap is set | | 18 | Chunked reads skip chunks a known mask excludes | this change | A read of a row range given an already-resolved mask (`MaskFuture::ready`, as a point read passes) requested every chunk in the range, even those the mask selects no row of, because the flat reader fetches its segment before it looks at the mask. Once string blocks are split under a byte cap (row 17), a point read of one row of an 8,192-row range fetched all of its split blocks, so a lookup got slower rather than faster (spiceai/spiceai#15323). `MaskFuture::known` exposes a mask known up front, and `ChunkedReader` skips chunks it excludes | `cargo test -p vortex-layout --lib chunked::reader::test::a_known_mask_reads_only_the_chunks_it_selects` | Proposed — upstream defect, not Spice behaviour | +| 19 | 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 | +| 20 | 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-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-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(()) + } } diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index 99d8c605941..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"] } +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/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/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..ba4be88eec5 --- /dev/null +++ b/vortex-layout/src/scan/split_gate.rs @@ -0,0 +1,482 @@ +// 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); + } + // 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), + }); + } + 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::task::Wake; + use std::task::Waker; + 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; + use crate::segments::mark_pending_wait; + + /// 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) + .map(|result| { + let (output, permit) = result.expect("split task joins"); + drop(permit); + output + }) + .collect() + .await; + for output in done { + 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) + .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); + } + + /// 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); + } + + #[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 a7ecf26de3b..e51b54b1ae6 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,8 @@ pub use source::*; #[cfg(any(test, feature = "_test-harness"))] pub use test::*; use vortex_error::VortexError; +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 new file mode 100644 index 00000000000..7f9826a2f24 --- /dev/null +++ b/vortex-layout/src/segments/wait_marked.rs @@ -0,0 +1,90 @@ +// 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 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 { + 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 { + #[pin] + inner: F, + // 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 = F::Output; + + 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 + } +}