diff --git a/Cargo.lock b/Cargo.lock index 009891d5e94..411e87b3d8c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12067,6 +12067,7 @@ dependencies = [ "parking_lot", "pin-project-lite", "rstest", + "rustix", "smol", "tempfile", "tokio", diff --git a/vortex-file/src/file.rs b/vortex-file/src/file.rs index 72a5e8787e0..fefe824ecf4 100644 --- a/vortex-file/src/file.rs +++ b/vortex-file/src/file.rs @@ -17,6 +17,7 @@ use vortex_array::dtype::FieldMask; use vortex_array::expr::Expression; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; +use vortex_io::request::IoService; use vortex_layout::LayoutReader; use vortex_layout::scan::layout::LayoutReaderDataSource; use vortex_layout::scan::scan_builder::ScanBuilder; @@ -48,6 +49,8 @@ pub struct VortexFile { metadata: Arc>, /// None id LayoutReader caching is turned off layout_reader_cache: Option>>, + /// Serves the planning protocol's reads, when the file was opened over a reader. + scan_io: Option>, } fn layout_reader( @@ -84,9 +87,21 @@ impl VortexFile { session, metadata: Arc::new(HashMap::new()), layout_reader_cache: None, + scan_io: None, } } + pub(crate) fn with_scan_io(mut self, scan_io: Arc) -> Self { + self.scan_io = Some(scan_io); + self + } + + /// The service that serves the planning protocol's reads of this file, when the file was + /// opened over a reader. + pub fn scan_io(&self) -> Option<&Arc> { + self.scan_io.as_ref() + } + pub(crate) fn with_metadata(mut self, metadata: Arc>) -> Self { self.metadata = metadata; self @@ -103,6 +118,7 @@ impl VortexFile { session: self.session, metadata: self.metadata, layout_reader_cache: Some(OnceLock::new()), + scan_io: self.scan_io, } } @@ -158,6 +174,8 @@ impl VortexFile { /// replacement source. pub fn with_segment_source(mut self, segment_source: Arc) -> Self { self.segment_source = segment_source; + // Reads now go to the replacement, which the old read driver knows nothing of. + self.scan_io = None; if self.layout_reader_cache.is_some() { self.layout_reader_cache = Some(OnceLock::new()); } diff --git a/vortex-file/src/open.rs b/vortex-file/src/open.rs index bbb63e6195d..aa1f0fff3c5 100644 --- a/vortex-file/src/open.rs +++ b/vortex-file/src/open.rs @@ -14,6 +14,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_io::VortexReadAt; +use vortex_io::request::IoService; use vortex_io::session::RuntimeSessionExt; use vortex_layout::segments::InstrumentedSegmentCache; use vortex_layout::segments::NoOpSegmentCache; @@ -35,6 +36,7 @@ use crate::MAX_POSTSCRIPT_SIZE; use crate::VortexFile; use crate::footer::Footer; use crate::segments::BufferSegmentSource; +use crate::segments::FileIoService; use crate::segments::FileSegmentSource; use crate::segments::InitialReadSegmentCache; use crate::segments::RequestMetrics; @@ -281,6 +283,7 @@ impl VortexOpenOptions { /// /// This is the common path for files, object stores, and custom random-access sources. pub async fn open_read(self, reader: R) -> VortexResult { + let has_segment_cache = self.segment_cache.is_some(); let segment_cache = self .segment_cache .clone() @@ -327,13 +330,26 @@ impl VortexOpenOptions { let metrics = RequestMetrics::new(metrics_registry.as_ref(), self.labels); + // Scans that drive the planning protocol read through their own service when it is + // selected, and through the segment source's read driver otherwise. + let batch_io = FileIoService::enabled().then(|| { + Arc::new(FileIoService::open( + footer.segment_specs_with_metadata(), + reader.clone(), + self.session.handle(), + metrics.clone(), + )) as Arc + }); + // Create a segment source backed by the VortexRead implementation. - let segment_source = Arc::new(SharedSegmentSource::new(FileSegmentSource::open( + let file_source = FileSegmentSource::open( footer.segment_specs_with_metadata(), reader, self.session.handle(), metrics, - ))); + ); + let scan_io = batch_io.unwrap_or_else(|| Arc::new(file_source.scan_io())); + let segment_source = Arc::new(SharedSegmentSource::new(file_source)); // Wrap up the segment source to first resolve segments from the initial read cache. let segment_source: Arc = Arc::new(SegmentCacheSourceAdapter::new( @@ -348,6 +364,13 @@ impl VortexOpenOptions { }; let file = VortexFile::new(footer, segment_source, self.session.clone()).with_metadata(metadata); + // The direct range service bypasses SegmentSource and its cache. Let V2 use the + // segment adapter when a caller explicitly configured a cache. + let file = if has_segment_cache { + file + } else { + file.with_scan_io(scan_io) + }; Ok(if self.cache_layout_reader { file.with_caching() } else { diff --git a/vortex-file/src/read/driver.rs b/vortex-file/src/read/driver.rs index 9f8b7aedbd3..3124effdfd9 100644 --- a/vortex-file/src/read/driver.rs +++ b/vortex-file/src/read/driver.rs @@ -4,6 +4,7 @@ use std::collections::BTreeMap; use std::collections::BTreeSet; use std::pin::Pin; +use std::sync::LazyLock; use std::task::Context; use std::task::Poll; @@ -13,7 +14,6 @@ use tracing::trace; use vortex_buffer::Alignment; use vortex_error::VortexExpect; use vortex_io::CoalesceConfig; -use vortex_utils::aliases::hash_set::HashSet; use crate::read::ReadRequest; use crate::read::RequestId; @@ -41,6 +41,11 @@ pin_project! { } impl IoRequestStream { + pub(crate) fn set_batch_size(&mut self, batch_size: usize) { + assert!(batch_size > 0, "I/O request batch size must be non-zero"); + self.batch_size = batch_size; + } + // FIXME(ngates): split this into coalesce_distance and max_read_size. We should keep // expanding the request by coalesce_distance, but stop if we hit max_read_size. pub(crate) fn new( @@ -90,6 +95,14 @@ where } } + if this.state.polled_requests.is_empty() { + return if *this.inner_done { + Poll::Ready(None) + } else { + Poll::Pending + }; + } + // 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 { @@ -129,16 +142,24 @@ struct State { // Metrics for tracking I/O request patterns metrics: RequestMetrics, coalesced_buffer_alignment: Alignment, + // Reuse removal storage across physical reads; the requests themselves move into the result. + coalesce_keys: Vec<(u64, RequestId)>, + coalesce_optional: bool, } impl State { fn new(metrics: RequestMetrics, coalesced_buffer_alignment: Alignment) -> Self { + static COALESCE_OPTIONAL: LazyLock = LazyLock::new(|| { + std::env::var("VORTEX_SCAN_IO_COALESCE_OPTIONAL").map_or(true, |value| value != "0") + }); Self { requests: BTreeMap::new(), polled_requests: BTreeMap::new(), requests_by_offset: BTreeSet::new(), metrics, coalesced_buffer_alignment, + coalesce_keys: Vec::new(), + coalesce_optional: *COALESCE_OPTIONAL, } } @@ -230,10 +251,8 @@ impl State { let mut current_end = requests[0].offset + requests[0].length as u64; let align = self.coalesced_buffer_alignment.as_usize() as u64; - // Track requests that we've already decided to remove (or that were cancelled) so that - // we don't repeatedly process them during range scans. - let mut keys_to_remove: Vec<(u64, RequestId)> = Vec::new(); - let mut ids_to_remove: HashSet = HashSet::new(); + // Retire chosen and cancelled offsets after each pass so expansion never revisits them. + self.coalesce_keys.clear(); let mut found_new_requests = true; // Keep expanding the window while we can find new requests within constraints @@ -249,22 +268,19 @@ impl State { .requests_by_offset .range((scan_start, RequestId::MIN)..=(scan_end, RequestId::MAX)) { - // Skip if we've already marked this request for removal - if ids_to_remove.contains(&req_id) { - continue; - } - - let req = self - .polled_requests - .get(&req_id) - .or_else(|| self.requests.get(&req_id)) - .vortex_expect("Missing request in requests_by_offset"); + let (req, polled) = match self.polled_requests.get(&req_id) { + Some(req) => (req, true), + None => ( + self.requests + .get(&req_id) + .vortex_expect("Missing request in requests_by_offset"), + false, + ), + }; // Skip any cancelled requests if req.callback.is_canceled() { - if ids_to_remove.insert(req_id) { - keys_to_remove.push((req_offset, req_id)); - } + self.coalesce_keys.push((req_offset, req_id)); continue; } @@ -277,6 +293,15 @@ impl State { let new_start = current_start.min(req_offset); let new_end = current_end.max(req_end); let aligned_start = new_start - (new_start % align); + // Optional members already covered by the range can share its result, but + // need not cause speculative bytes to be read before filtering finishes. + if !self.coalesce_optional + && !polled + && (aligned_start < current_start - (current_start % align) + || new_end != current_end) + { + continue; + } let new_total_size = new_end - aligned_start; if new_total_size > window.max_size { @@ -293,20 +318,16 @@ impl State { .vortex_expect("Missing request in requests_by_offset"); requests.push(req); - if ids_to_remove.insert(req_id) { - keys_to_remove.push((req_offset, req_id)); - } + self.coalesce_keys.push((req_offset, req_id)); found_new_requests = true; } } - } - - // Remove any dropped requests - for (req_offset, req_id) in keys_to_remove { - self.requests_by_offset.remove(&(req_offset, req_id)); - self.polled_requests - .remove(&req_id) - .or_else(|| self.requests.remove(&req_id)); + for (req_offset, req_id) in self.coalesce_keys.drain(..) { + self.requests_by_offset.remove(&(req_offset, req_id)); + self.polled_requests + .remove(&req_id) + .or_else(|| self.requests.remove(&req_id)); + } } // Sort requests by offset for correct slicing in resolve @@ -339,9 +360,11 @@ mod tests { use futures::channel::mpsc; use futures::channel::oneshot; use futures::stream; + use rstest::rstest; use vortex_array::buffer::BufferHandle; use vortex_buffer::Alignment; use vortex_error::VortexResult; + use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_metrics::DefaultMetricsRegistry; use vortex_metrics::MetricValue; @@ -751,6 +774,105 @@ mod tests { assert_eq!(outputs.len(), 0); } + #[test] + fn coalescing_expands_in_multiple_passes_and_removes_cancelled_neighbors() -> VortexResult<()> { + let registry = DefaultMetricsRegistry::default(); + let mut state = State::new(RequestMetrics::new(®istry, vec![]), Alignment::none()); + let mut receivers = Vec::new(); + for (id, offset) in (0..=100).step_by(10).enumerate() { + let (request, receiver) = create_request(id, offset, 1); + state.on_event(ReadEvent::Request(request)); + receivers.push(receiver); + } + let (cancelled, receiver) = create_request(11, 45, 1); + state.on_event(ReadEvent::Request(cancelled)); + drop(receiver); + let (separate, receiver) = create_request(12, 150, 1); + state.on_event(ReadEvent::Request(separate)); + receivers.push(receiver); + state.on_event(ReadEvent::Polled(5)); + state.on_event(ReadEvent::Polled(12)); + let window = CoalesceConfig::new(10, 512); + let first = state + .next_coalesced(&window) + .ok_or_else(|| vortex_err!("missing first read"))?; + assert_eq!(*first.range(), 0..101); + assert_eq!( + first + .requests() + .iter() + .map(|request| request.offset) + .collect::>(), + (0u64..=100).step_by(10).collect::>() + ); + let second = state + .next_coalesced(&window) + .ok_or_else(|| vortex_err!("missing second read"))?; + assert_eq!(*second.range(), 150..151); + assert!(state.requests_by_offset.is_empty()); + assert!(state.requests.is_empty()); + assert!(state.polled_requests.is_empty()); + drop(receivers); + Ok(()) + } + + #[rstest] + #[case(true, Alignment::none())] + #[case(false, Alignment::none())] + #[case(true, Alignment::new(8))] + #[case(false, Alignment::new(8))] + fn optional_coalescing_preserves_unread_interests( + #[case] optional: bool, + #[case] alignment: Alignment, + ) -> VortexResult<()> { + let registry = DefaultMetricsRegistry::default(); + let mut state = State::new(RequestMetrics::new(®istry, vec![]), alignment); + state.coalesce_optional = optional; + let mut receivers = Vec::new(); + for (id, offset, length) in [ + (1, 0, 10), + (2, 10, 10), + (3, 12, 4), + (4, 20, 10), + (5, 30, 10), + (6, 9, 1), + ] { + let (request, receiver) = create_request(id, offset, length); + state.on_event(ReadEvent::Request(request)); + receivers.push(receiver); + } + state.on_event(ReadEvent::Polled(2)); + state.on_event(ReadEvent::Polled(4)); + let window = CoalesceConfig::new(100, 512); + let first = state + .next_coalesced(&window) + .ok_or_else(|| vortex_err!("missing demanded read"))?; + let start = 10 - (10 % alignment.as_usize() as u64); + assert_eq!(*first.range(), if optional { 0..40 } else { start..30 }); + assert!(first.requests().iter().any(|request| request.id == 3)); + if alignment.as_usize() == 8 || optional { + assert!(first.requests().iter().any(|request| request.id == 6)); + } + assert!(state.next_coalesced(&window).is_none()); + + state.on_event(ReadEvent::Polled(1)); + if optional { + assert!(state.next_coalesced(&window).is_none()); + } else { + let next = state + .next_coalesced(&window) + .ok_or_else(|| vortex_err!("missing later demand"))?; + assert_eq!(*next.range(), 0..10); + assert!(state.requests.contains_key(&5)); + state.on_event(ReadEvent::Dropped(5)); + } + assert!(state.requests.is_empty()); + assert!(state.polled_requests.is_empty()); + assert!(state.requests_by_offset.is_empty()); + drop(receivers); + Ok(()) + } + #[tokio::test] async fn test_mixed_coalesced_and_single() { let (req1, _rx1) = create_request(1, 0, 10); diff --git a/vortex-file/src/read/request.rs b/vortex-file/src/read/request.rs index 1ef5beb3757..11a24ae3e95 100644 --- a/vortex-file/src/read/request.rs +++ b/vortex-file/src/read/request.rs @@ -21,6 +21,13 @@ use vortex_error::vortex_ensure; pub(crate) struct IoRequest(IoRequestInner); impl IoRequest { + pub(crate) fn requests(&self) -> &[ReadRequest] { + match &self.0 { + IoRequestInner::Single(request) => std::slice::from_ref(request), + IoRequestInner::Coalesced(request) => request.requests(), + } + } + pub(crate) fn new_single(request: ReadRequest) -> Self { IoRequest(IoRequestInner::Single(request)) } diff --git a/vortex-file/src/segments/cache.rs b/vortex-file/src/segments/cache.rs index c022751edcb..4b151bdb632 100644 --- a/vortex-file/src/segments/cache.rs +++ b/vortex-file/src/segments/cache.rs @@ -20,6 +20,13 @@ pub struct InitialReadSegmentCache { #[async_trait] impl SegmentCache for InitialReadSegmentCache { + fn get_if_ready(&self, id: SegmentId) -> Option { + self.initial + .get(&id) + .cloned() + .or_else(|| self.fallback.get_if_ready(id)) + } + async fn get(&self, id: SegmentId) -> VortexResult> { if let Some(buffer) = self.initial.get(&id) { return Ok(Some(buffer.clone())); diff --git a/vortex-file/src/segments/io_service.rs b/vortex-file/src/segments/io_service.rs new file mode 100644 index 00000000000..886cadb84c1 --- /dev/null +++ b/vortex-file/src/segments/io_service.rs @@ -0,0 +1,1032 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Serves the planning protocol's requests straight from a file's reader. +//! +//! The service keeps one table of byte ranges for the file, shared by every session. A request +//! only updates the table: an announcement registers a range, a prefetch or fetch also marks it +//! wanted, and a fetch waits for the bytes. A range that several sessions ask for is one entry, +//! read once, and its bytes are kept until every registration is consumed or cleared. +//! +//! Physical reads are chosen from the table rather than from one request: a wanted range is +//! extended over the registered ranges around it, whichever session registered them and whether +//! or not they are wanted yet, while the gaps stay within the reader's coalescing distance and +//! the read within its maximum size. Sessions that register the ranges they will need well ahead +//! of asking for them therefore get large reads that carry their later requests. +//! +//! Enable `RUST_LOG=vortex_file::io_lock=trace` to measure table-lock acquisition and hold times +//! by operation. These diagnostics exclude the session-local claim and delivery locks. + +use std::collections::BTreeSet; +use std::collections::VecDeque; +use std::env; +use std::ops::Deref; +use std::ops::DerefMut; +use std::sync::Arc; +use std::sync::LazyLock; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::time::Duration; +use std::time::Instant; + +use futures::StreamExt; +use futures::task::AtomicWaker; +use parking_lot::Mutex; +use parking_lot::MutexGuard; +use tracing::trace; +use vortex_array::buffer::BufferHandle; +use vortex_buffer::Alignment; +use vortex_error::VortexError; +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::CoalesceConfig; +use vortex_io::ReadAtRequest; +use vortex_io::VortexReadAt; +use vortex_io::request::Completion; +use vortex_io::request::IoBatch; +use vortex_io::request::IoIntent; +use vortex_io::request::IoOwnerId; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_io::request::IoService; +use vortex_io::request::IoSource; +use vortex_io::request::IoTarget; +use vortex_io::runtime::Handle; +use vortex_utils::aliases::hash_map::Entry as MapEntry; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::SegmentSpec; +use crate::segments::RequestMetrics; + +/// The environment variable that selects this service for a file's scans when set to `1`. +pub const ENV_VAR: &str = "VORTEX_SCAN_BATCH_IO"; + +static ENABLED: LazyLock = LazyLock::new(|| env::var(ENV_VAR).is_ok_and(|v| v == "1")); + +/// A byte range of the file: its offset and length. +type Range = (u64, usize); + +fn end(range: Range) -> u64 { + range.0 + range.1 as u64 +} + +/// Serves the scans of one file from its reader, one session per planning run root. +#[derive(Clone)] +pub struct FileIoService(Arc); + +struct Inner { + reader: Arc, + handle: Handle, + coalesce: Option, + /// The alignment every physical read starts at, so each range in it keeps its own. + max_alignment: Alignment, + /// The most physical reads in flight at once. + concurrency: usize, + metrics: RequestMetrics, + table: Mutex, +} + +#[derive(Default)] +struct Table { + entries: HashMap, + /// Ranges no physical read has been chosen for, in file order. + unread: BTreeSet, + /// Wanted ranges with a fetch waiting, oldest first. May name ranges already chosen. + fetches: VecDeque, + /// Wanted ranges nothing waits for yet, oldest first. May name ranges already chosen. + prefetches: VecDeque, + /// Physical reads in flight. + active: usize, +} + +/// Log after unlocking so the trace subscriber does not extend the critical section. +struct TableGuard<'a> { + guard: Option>, + timing: Option<(Instant, Duration, bool)>, + operation: &'static str, +} + +impl Deref for TableGuard<'_> { + type Target = Table; + + fn deref(&self) -> &Table { + self.guard.as_deref().vortex_expect("table guard is held") + } +} + +impl DerefMut for TableGuard<'_> { + fn deref_mut(&mut self) -> &mut Table { + self.guard + .as_deref_mut() + .vortex_expect("table guard is held") + } +} + +impl Drop for TableGuard<'_> { + fn drop(&mut self) { + let hold = self.timing.map(|(acquired, ..)| acquired.elapsed()); + drop(self.guard.take()); + if let Some((_, wait, contended)) = self.timing { + trace!( + target: "vortex_file::io_lock", + operation = self.operation, + contended, + wait_ns = u64::try_from(wait.as_nanos()).unwrap_or(u64::MAX), + hold_ns = u64::try_from(hold.vortex_expect("timed guard has hold duration").as_nanos()) + .unwrap_or(u64::MAX), + "file IO table lock" + ); + } + } +} + +struct Entry { + alignment: Alignment, + state: State, + /// Sessions that registered the range and have not been cleared. + claims: usize, + /// Fetches to complete when the bytes arrive. + waiters: Vec, +} + +enum State { + /// Announced: no read has been chosen and nothing wants the bytes yet. + Registered, + /// Wanted by a prefetch or fetch; no read has been chosen. + Wanted, + /// Part of a physical read in flight. + Reading, + Ready(BufferHandle), + Failed(Arc), +} + +struct Waiter { + session: Arc, + owner: IoOwnerId, + request: IoRequestId, +} + +/// A session's completed fetches, and the waker of the run waiting for the next. +#[derive(Default)] +struct Delivery { + ready: Mutex>, + waker: AtomicWaker, + /// Fetches submitted and not yet taken by the run. + outstanding: AtomicUsize, +} + +impl Delivery { + fn push( + &self, + range: Range, + owner: IoOwnerId, + request: IoRequestId, + result: VortexResult, + ) { + self.ready.lock().push_back(( + range, + Completion { + owner, + request, + result: result.map(IoResult::Bytes), + }, + )); + } + + fn take(&self) -> Option<(Range, Completion)> { + let completion = self.ready.lock().pop_front()?; + self.outstanding.fetch_sub(1, Ordering::Relaxed); + Some(completion) + } +} + +/// One physical read: a coalesced range and the registered ranges it carries. +struct PhysicalRead { + request: ReadAtRequest, + members: Vec, +} + +impl FileIoService { + /// Whether `VORTEX_SCAN_BATCH_IO=1` selects this service. Read once per process. + pub fn enabled() -> bool { + *ENABLED + } + + /// Open a file-backed IO service over `reader`. + /// + /// Unlike the segment source, the service spawns nothing until a session wants bytes. The + /// reads chosen together run as one task on `handle`, at most the reader's concurrency in + /// flight at once. + pub fn open( + segments: Arc<[SegmentSpec]>, + reader: R, + handle: Handle, + metrics: RequestMetrics, + ) -> Self { + let max_alignment = segments + .iter() + .map(|segment| segment.alignment) + .max() + .unwrap_or_else(Alignment::none); + let coalesce = reader.coalesce_config().map(|mut config| { + // Aligning the coalesced start down can add up to (alignment - 1) bytes. + // Increase max_size to keep the effective payload window consistent. + let extra = (max_alignment.as_usize() as u64).saturating_sub(1); + config.max_size = config.max_size.saturating_add(extra); + config + }); + let concurrency = reader.concurrency(); + if concurrency == 0 { + vortex_panic!( + "VortexReadAt::concurrency returned 0 (uri={:?}); this would stall I/O", + reader.uri() + ); + } + + Self(Arc::new(Inner { + reader: Arc::new(reader), + handle, + coalesce, + max_alignment, + concurrency, + metrics, + table: Mutex::default(), + })) + } +} + +impl IoService for FileIoService { + fn session(&self) -> Arc { + Arc::new(FileIoSession { + service: self.clone(), + delivery: Arc::default(), + claimed: Mutex::default(), + }) + } +} + +impl Table { + /// The next wanted range no read has been chosen for: one a fetch waits for, if any. + fn next_wanted(&mut self) -> Option { + loop { + let range = self + .fetches + .pop_front() + .or_else(|| self.prefetches.pop_front())?; + if matches!( + self.entries.get(&range), + Some(Entry { + state: State::Wanted, + .. + }) + ) { + return Some(range); + } + } + } + + /// Moves `range` into a read being built. + fn take(&mut self, range: Range, members: &mut Vec) { + self.unread.remove(&range); + if let Some(entry) = self.entries.get_mut(&range) { + entry.state = State::Reading; + } + members.push(range); + } + + /// Chooses the physical read for the next wanted range, extended over its registered + /// neighbours in file order while the gaps and the total stay within `coalesce`. + fn choose( + &mut self, + coalesce: Option<&CoalesceConfig>, + max_alignment: Alignment, + ) -> Option { + let head = self.next_wanted()?; + let Some(window) = coalesce else { + let alignment = self.entries.get(&head).map(|entry| entry.alignment); + let mut members = Vec::with_capacity(1); + self.take(head, &mut members); + return Some(PhysicalRead { + request: ReadAtRequest::new( + head.0, + head.1, + alignment.unwrap_or_else(Alignment::none), + ), + members, + }); + }; + + let align = max_alignment.as_usize() as u64; + let fits = |start: u64, end: u64| end - (start - start % align) <= window.max_size; + let (mut start, mut stop) = (head.0, end(head)); + let mut members = Vec::new(); + self.take(head, &mut members); + loop { + let mut grew = false; + // Forwards first: ranges after a wanted one are what the scan asks for next. + if let Some(&next) = self.unread.range((start, 0)..).next() + && next.0 <= stop.saturating_add(window.distance) + && fits(start, stop.max(end(next))) + { + stop = stop.max(end(next)); + self.take(next, &mut members); + grew = true; + } + if let Some(&previous) = self.unread.range(..(start, 0)).next_back() + && end(previous).saturating_add(window.distance) >= start + && fits(previous.0, stop.max(end(previous))) + { + start = previous.0; + stop = stop.max(end(previous)); + self.take(previous, &mut members); + grew = true; + } + if !grew { + break; + } + } + + let aligned_start = start - start % align; + Some(PhysicalRead { + request: ReadAtRequest::new( + aligned_start, + usize::try_from(stop - aligned_start).unwrap_or(usize::MAX), + max_alignment, + ), + members, + }) + } +} + +impl Inner { + fn lock_table(&self, operation: &'static str) -> TableGuard<'_> { + let (guard, timing) = if tracing::enabled!(target: "vortex_file::io_lock", tracing::Level::TRACE) + { + let start = Instant::now(); + let (guard, contended) = match self.table.try_lock() { + Some(guard) => (guard, false), + None => (self.table.lock(), true), + }; + let acquired = Instant::now(); + ( + guard, + Some((acquired, acquired.duration_since(start), contended)), + ) + } else { + (self.table.lock(), None) + }; + TableGuard { + guard: Some(guard), + timing, + operation, + } + } + + /// Starts physical reads for wanted ranges while the reader has concurrency to spare. + fn dispatch(self: &Arc) { + let mut reads = Vec::new(); + { + let mut table = self.lock_table("choose"); + while table.active < self.concurrency { + let Some(read) = table.choose(self.coalesce.as_ref(), self.max_alignment) else { + break; + }; + table.active += 1; + reads.push(read); + } + } + if reads.is_empty() { + return; + } + + self.metrics.read_ranges_calls.add(1); + self.metrics + .read_ranges_num_ranges + .update(reads.len() as f64); + if reads.len() > 1 { + self.metrics.read_ranges_multi.add(1); + } + for read in &reads { + trace!( + "Coalesced {} requests into range {}..{} (len={})", + read.members.len(), + read.request.offset, + read.request.offset + read.request.length as u64, + read.request.length, + ); + match read.members.len() { + 1 => self.metrics.individual_requests.add(1), + members => { + self.metrics.coalesced_requests.add(1); + self.metrics.num_requests_coalesced.update(members as f64); + } + } + } + + // One task drives every read chosen together, as the reader's `read_ranges` allows. + let inner = Arc::clone(self); + self.handle + .spawn_io(async move { + let requests: Arc<[ReadAtRequest]> = + reads.iter().map(|read| read.request).collect(); + let mut pending: Vec> = reads.into_iter().map(Some).collect(); + let mut results = inner.reader.read_ranges(requests); + while let Some((request, result)) = results.next().await { + let Some(read) = pending + .iter_mut() + .find(|read| read.as_ref().is_some_and(|read| read.request == request)) + .and_then(Option::take) + else { + tracing::warn!(?request, "reader returned an unknown range"); + continue; + }; + inner.complete(read, result); + } + // A reader that ends early leaves reads unanswered; fail them rather than hang. + for read in pending.into_iter().flatten() { + let error = vortex_err!("the reader ended before answering {:?}", read.request); + inner.complete(read, Err(error)); + } + }) + .detach(); + } + + /// Hands the bytes of a finished read to its ranges, wakes the fetches waiting for them, and + /// starts whatever became startable. + fn complete(self: &Arc, read: PhysicalRead, result: VortexResult) { + let result = result + .and_then(|buffer| { + if buffer.len() != read.request.length { + vortex_bail!( + "expected {} bytes at {} but the reader returned {}", + read.request.length, + read.request.offset, + buffer.len() + ); + } + buffer.ensure_aligned(Alignment::none()) + }) + .map_err(Arc::new); + + let mut woken = Vec::new(); + { + let mut table = self.lock_table("complete"); + table.active -= 1; + for range in read.members { + let MapEntry::Occupied(mut occupied) = table.entries.entry(range) else { + continue; + }; + let entry = occupied.get_mut(); + let bytes = result.clone().and_then(|base| { + let start = usize::try_from(range.0 - read.request.offset) + .map_err(|err| Arc::new(vortex_err!("{err}")))?; + base.slice(start..start + range.1) + .ensure_aligned(entry.alignment) + .map_err(Arc::new) + }); + for waiter in entry.waiters.drain(..) { + let delivered = match &bytes { + Ok(bytes) => Ok(bytes.clone()), + Err(error) => Err(VortexError::from(Arc::clone(error))), + }; + waiter + .session + .push(range, waiter.owner, waiter.request, delivered); + woken.push(waiter.session); + } + if entry.claims == 0 { + // Every session that registered the range was cleared while it was read. + occupied.remove(); + continue; + } + entry.state = match bytes { + Ok(bytes) => State::Ready(bytes), + Err(error) => State::Failed(error), + }; + } + } + for session in woken { + session.waker.wake(); + } + self.dispatch(); + } +} + +/// One planning run root's reads through a [`FileIoService`]. +struct FileIoSession { + service: FileIoService, + delivery: Arc, + /// One claim per range while registered consumers or undelivered fetches remain. + claimed: Mutex>, +} + +#[derive(Default)] +struct Claim { + unfetched: usize, + pending: usize, +} + +impl FileIoSession { + fn take(&self) -> Option { + let mut claimed = self.claimed.lock(); + let (range, completion) = self.delivery.take()?; + let claim = claimed + .get_mut(&range) + .vortex_expect("a delivered fetch has a pending range"); + claim.pending -= 1; + if claim.pending == 0 && claim.unfetched == 0 { + claimed.remove(&range); + let mut table = self.service.0.lock_table("consume"); + if let MapEntry::Occupied(mut occupied) = table.entries.entry(range) { + let entry = occupied.get_mut(); + entry.claims = entry.claims.saturating_sub(1); + if entry.claims == 0 && !matches!(entry.state, State::Reading) { + occupied.remove(); + table.unread.remove(&range); + } + } + } + Some(completion) + } + + /// Gives up the session's claims, and with them any bytes nothing else has registered. + fn release_claims(&self) { + let mut claimed = self.claimed.lock(); + let ranges = std::mem::take(&mut *claimed); + let mut table = self.service.0.lock_table("cleanup"); + for (range, _) in ranges { + let MapEntry::Occupied(mut occupied) = table.entries.entry(range) else { + continue; + }; + let entry = occupied.get_mut(); + let pending = entry.waiters.len(); + entry + .waiters + .retain(|waiter| !Arc::ptr_eq(&waiter.session, &self.delivery)); + let cancelled = pending - entry.waiters.len(); + if cancelled > 0 { + self.delivery + .outstanding + .fetch_sub(cancelled, Ordering::Relaxed); + } + entry.claims = entry.claims.saturating_sub(1); + // A read in flight keeps its entry until it completes, which then removes it. + if entry.claims == 0 && !matches!(entry.state, State::Reading) { + occupied.remove(); + table.unread.remove(&range); + } + } + drop(table); + let dropped = std::mem::take(&mut *self.delivery.ready.lock()).len(); + self.delivery + .outstanding + .fetch_sub(dropped, Ordering::Relaxed); + } +} + +impl IoSource for FileIoSession { + fn submit(&self, owner: IoOwnerId, batch: IoBatch) -> VortexResult<()> { + let mut wanted = false; + { + let mut claimed = self.claimed.lock(); + let mut table = self.service.0.lock_table("submit"); + let Table { + entries, + unread, + fetches, + prefetches, + .. + } = &mut *table; + for request in batch { + if request.intent == IoIntent::Forget { + continue; + } + let IoTarget::Range { + offset, + len, + alignment, + } = request.target + else { + vortex_bail!("a file's IO serves byte ranges, not {:?}", request.target); + }; + let range = (offset, len); + let entry = entries.entry(range).or_insert_with(|| { + unread.insert(range); + Entry { + alignment, + state: State::Registered, + claims: 0, + waiters: Vec::new(), + } + }); + let claim = claimed.entry(range).or_insert_with(|| { + entry.claims += 1; + Claim::default() + }); + entry.alignment = entry.alignment.max(alignment); + if request.intent == IoIntent::Announce { + claim.unfetched += 1; + continue; + } + + let fetch = request.intent == IoIntent::Fetch; + if fetch { + claim.unfetched = claim.unfetched.saturating_sub(1); + claim.pending += 1; + } else { + claim.unfetched = claim.unfetched.max(1); + } + match &entry.state { + State::Ready(bytes) if fetch => { + let bytes = bytes.clone().ensure_aligned(alignment); + self.delivery.outstanding.fetch_add(1, Ordering::Relaxed); + self.delivery.push(range, owner, request.request, bytes); + continue; + } + State::Failed(error) if fetch => { + let error = VortexError::from(Arc::clone(error)); + self.delivery.outstanding.fetch_add(1, Ordering::Relaxed); + self.delivery + .push(range, owner, request.request, Err(error)); + continue; + } + State::Ready(_) | State::Failed(_) | State::Reading => {} + State::Registered | State::Wanted => { + entry.state = State::Wanted; + if fetch { + fetches.push_back(range); + } else { + prefetches.push_back(range); + } + wanted = true; + } + } + if fetch { + entry.waiters.push(Waiter { + session: Arc::clone(&self.delivery), + owner, + request: request.request, + }); + self.delivery.outstanding.fetch_add(1, Ordering::Relaxed); + } + } + } + if wanted { + self.service.0.dispatch(); + } + Ok(()) + } + + fn poll(&self) -> VortexResult> { + Ok(self.take()) + } + + fn poll_completion(&self, cx: &mut Context<'_>) -> Poll> { + // Register before looking, so a completion pushed in between still wakes the run. + self.delivery.waker.register(cx.waker()); + match self.take() { + Some(completion) => Poll::Ready(Ok(completion)), + None if self.delivery.outstanding.load(Ordering::Relaxed) == 0 => Poll::Ready(Err( + vortex_err!("the run is waiting with no fetch in flight"), + )), + None => Poll::Pending, + } + } + + fn wait(&self) -> VortexResult { + vortex_bail!("a file's IO is awaited through poll_completion, not waited on") + } + + fn release(&self, _owner: IoOwnerId) {} + + fn clear(&self) { + self.release_claims(); + } +} + +impl Drop for FileIoSession { + fn drop(&mut self) { + self.release_claims(); + } +} + +#[cfg(test)] +mod tests { + use futures::future::BoxFuture; + use vortex_buffer::ByteBuffer; + use vortex_io::request::IoRequest; + use vortex_io::session::RuntimeSession; + use vortex_io::session::RuntimeSessionExt; + use vortex_layout::session::LayoutSession; + use vortex_metrics::DefaultMetricsRegistry; + use vortex_session::VortexSession; + + use super::*; + + fn session() -> VortexSession { + let session = vortex_array::array_session() + .with::() + .with::() + .with_tokio(); + crate::register_default_encodings(&session); + crate::enable_all_registered_array_encodings(&session); + session + } + + fn metrics() -> RequestMetrics { + RequestMetrics::new(&DefaultMetricsRegistry::default(), Vec::new()) + } + + /// A buffer that coalesces within 8 bytes, up to 64, and records every physical read. + #[derive(Clone)] + struct RecordingReader { + bytes: ByteBuffer, + reads: Arc>>, + } + + impl RecordingReader { + fn new(bytes: ByteBuffer) -> Self { + Self { + bytes, + reads: Arc::default(), + } + } + + fn reads(&self) -> Vec<(u64, usize)> { + let mut reads = self.reads.lock().clone(); + reads.sort_unstable(); + reads + } + } + + impl VortexReadAt for RecordingReader { + fn coalesce_config(&self) -> Option { + Some(CoalesceConfig { + distance: 8, + max_size: 64, + }) + } + + fn concurrency(&self) -> usize { + 4 + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + self.bytes.size() + } + + fn read_at( + &self, + offset: u64, + length: usize, + alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + self.reads.lock().push((offset, length)); + self.bytes.read_at(offset, length, alignment) + } + } + + fn request(intent: IoIntent, id: u32, offset: u64, len: usize) -> IoRequest { + IoRequest { + intent, + request: IoRequestId(id), + target: IoTarget::range(offset, len), + } + } + + fn service(session: &VortexSession) -> (FileIoService, RecordingReader) { + let reader = RecordingReader::new(ByteBuffer::from((0..=255u8).collect::>())); + let service = + FileIoService::open(Arc::from([]), reader.clone(), session.handle(), metrics()); + (service, reader) + } + + /// Takes `count` completions from `io`, as `(request, bytes)` in request order. + async fn completions( + io: &Arc, + count: usize, + ) -> VortexResult)>> { + let mut delivered = Vec::new(); + for _ in 0..count { + let completion = std::future::poll_fn(|cx| io.poll_completion(cx)).await?; + let IoResult::Bytes(bytes) = completion.result? else { + vortex_bail!("expected bytes"); + }; + delivered.push(( + completion.request.0, + bytes.try_into_host_sync()?.as_slice().to_vec(), + )); + } + delivered.sort(); + Ok(delivered) + } + + /// The fetches of one batch that lie within the coalescing window become one physical read, + /// whatever order the batch lists them in, and a fetch outside it gets its own. + #[tokio::test(flavor = "multi_thread")] + async fn a_batch_is_coalesced_as_a_whole() -> VortexResult<()> { + let session = session(); + let (service, reader) = service(&session); + let io = service.session(); + io.submit( + IoOwnerId(0), + vec![ + request(IoIntent::Fetch, 0, 20, 4), + request(IoIntent::Fetch, 1, 0, 4), + request(IoIntent::Fetch, 2, 10, 4), + request(IoIntent::Fetch, 3, 200, 4), + ], + )?; + + assert_eq!( + completions(&io, 4).await?, + vec![ + (0, vec![20, 21, 22, 23]), + (1, vec![0, 1, 2, 3]), + (2, vec![10, 11, 12, 13]), + (3, vec![200, 201, 202, 203]), + ] + ); + assert_eq!(reader.reads(), vec![(0, 24), (200, 4)]); + Ok(()) + } + + /// An announcement reads nothing by itself, is carried by a nearby read of another session, + /// and is then served without a second read. + #[tokio::test(flavor = "multi_thread")] + async fn an_announcement_joins_another_sessions_read() -> VortexResult<()> { + let session = session(); + let (service, reader) = service(&session); + let (announcer, fetcher) = (service.session(), service.session()); + + announcer.submit(IoOwnerId(0), vec![request(IoIntent::Announce, 0, 8, 4)])?; + assert!(reader.reads().is_empty()); + + fetcher.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 0, 0, 4)])?; + assert_eq!(completions(&fetcher, 1).await?, vec![(0, vec![0, 1, 2, 3])]); + assert_eq!(reader.reads(), vec![(0, 12)]); + + announcer.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 1, 8, 4)])?; + assert_eq!( + completions(&announcer, 1).await?, + vec![(1, vec![8, 9, 10, 11])] + ); + assert_eq!(reader.reads(), vec![(0, 12)]); + Ok(()) + } + + /// Two sessions that ask for the same range share one read, and the bytes stay for the + /// session with an unconsumed announcement after the fetches complete. + #[tokio::test(flavor = "multi_thread")] + async fn a_range_is_read_once_for_every_session() -> VortexResult<()> { + let session = session(); + let (service, reader) = service(&session); + let (first, second, late) = (service.session(), service.session(), service.session()); + late.submit(IoOwnerId(0), vec![request(IoIntent::Announce, 0, 100, 4)])?; + + first.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 0, 100, 4)])?; + second.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 0, 100, 4)])?; + let bytes = vec![(0, vec![100, 101, 102, 103])]; + assert_eq!(completions(&first, 1).await?, bytes); + assert_eq!(completions(&second, 1).await?, bytes); + first.clear(); + second.clear(); + + late.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 0, 100, 4)])?; + assert_eq!(completions(&late, 1).await?, bytes); + assert_eq!(reader.reads(), vec![(100, 4)]); + Ok(()) + } + + #[tokio::test(flavor = "multi_thread")] + async fn a_repeated_fetch_reads_again_with_stricter_alignment() -> VortexResult<()> { + let session = session(); + let (service, reader) = service(&session); + let io = service.session(); + io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 0, 1, 8)])?; + completions(&io, 1).await?; + + io.submit( + IoOwnerId(0), + vec![IoRequest { + intent: IoIntent::Fetch, + request: IoRequestId(1), + target: IoTarget::Range { + offset: 1, + len: 8, + alignment: Alignment::new(64), + }, + }], + )?; + let completion = std::future::poll_fn(|cx| io.poll_completion(cx)).await?; + let IoResult::Bytes(bytes) = completion.result? else { + vortex_bail!("expected bytes"); + }; + let bytes = bytes.try_into_host_sync()?; + assert_eq!(bytes.as_slice(), &[1, 2, 3, 4, 5, 6, 7, 8]); + assert!(bytes.alignment() >= Alignment::new(64)); + assert_eq!(reader.reads(), vec![(1, 8), (1, 8)]); + Ok(()) + } + + #[tokio::test(flavor = "multi_thread")] + async fn concurrent_fetches_consume_their_registration_together() -> VortexResult<()> { + let session = session(); + let (service, reader) = service(&session); + let io = service.session(); + io.submit( + IoOwnerId(0), + vec![ + request(IoIntent::Fetch, 0, 100, 4), + request(IoIntent::Fetch, 1, 100, 4), + ], + )?; + assert_eq!( + completions(&io, 2).await?, + vec![(0, vec![100, 101, 102, 103]), (1, vec![100, 101, 102, 103])] + ); + assert_eq!(reader.reads(), vec![(100, 4)]); + + io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 2, 100, 4)])?; + assert_eq!( + completions(&io, 1).await?, + vec![(2, vec![100, 101, 102, 103])] + ); + assert_eq!(reader.reads(), vec![(100, 4), (100, 4)]); + Ok(()) + } + + #[tokio::test] + async fn clearing_in_flight_fetches_does_not_leave_a_waiter() -> VortexResult<()> { + let session = session(); + let (service, _) = service(&session); + let range = (0, 4); + let io = FileIoSession { + service: service.clone(), + delivery: Arc::default(), + claimed: Mutex::new(HashMap::from_iter([( + range, + Claim { + unfetched: 0, + pending: 1, + }, + )])), + }; + io.delivery.outstanding.store(1, Ordering::Relaxed); + { + let mut table = service.0.table.lock(); + table.active = 1; + table.entries.insert( + range, + Entry { + alignment: Alignment::none(), + state: State::Reading, + claims: 1, + waiters: vec![Waiter { + session: Arc::clone(&io.delivery), + owner: IoOwnerId(0), + request: IoRequestId(0), + }], + }, + ); + } + + io.clear(); + let mut cx = Context::from_waker(std::task::Waker::noop()); + assert!(matches!(io.poll_completion(&mut cx), Poll::Ready(Err(_)))); + service.0.complete( + PhysicalRead { + request: ReadAtRequest::new(0, 4, Alignment::none()), + members: vec![range], + }, + Ok(BufferHandle::new_host(ByteBuffer::from(vec![0, 1, 2, 3]))), + ); + assert!(io.poll()?.is_none()); + assert!(service.0.table.lock().entries.is_empty()); + Ok(()) + } + + /// A session that is dropped withdraws its registrations, so later reads do not carry them, + /// and its bytes go with it. + #[tokio::test(flavor = "multi_thread")] + async fn a_dropped_session_gives_up_its_ranges() -> VortexResult<()> { + let session = session(); + let (service, reader) = service(&session); + let dropped = service.session(); + dropped.submit(IoOwnerId(0), vec![request(IoIntent::Announce, 0, 8, 4)])?; + drop(dropped); + + let io = service.session(); + io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, 0, 0, 4)])?; + assert_eq!(completions(&io, 1).await?, vec![(0, vec![0, 1, 2, 3])]); + assert_eq!(reader.reads(), vec![(0, 4)]); + drop(io); + assert!(service.0.table.lock().entries.is_empty()); + Ok(()) + } +} diff --git a/vortex-file/src/segments/mod.rs b/vortex-file/src/segments/mod.rs index d96620fca23..18ac95489c6 100644 --- a/vortex-file/src/segments/mod.rs +++ b/vortex-file/src/segments/mod.rs @@ -2,8 +2,12 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod cache; -mod source; +mod io_service; +mod scan_io; +pub(crate) mod source; pub(crate) mod writer; pub use cache::*; +pub use io_service::FileIoService; +pub use scan_io::FileScanIo; pub use source::*; diff --git a/vortex-file/src/segments/scan_io.rs b/vortex-file/src/segments/scan_io.rs new file mode 100644 index 00000000000..61a178ca574 --- /dev/null +++ b/vortex-file/src/segments/scan_io.rs @@ -0,0 +1,1151 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Serves the planning protocol's requests with a file's coalescing reads. +//! +//! The file's read driver tracks each read through three events, and each intent of the protocol +//! maps onto them. An announcement registers a read, which does no IO itself but may be folded +//! into a nearby read. A prefetch registers a read and marks it wanted, so the driver starts it +//! and the service keeps the bytes. A fetch does the same, or marks an earlier announcement or +//! prefetch of the same range wanted, and consumes one registered consumer. Pending consumers +//! share the read, as V1's eagerly constructed projection futures do. Once consumed, later fetches +//! register new reads. Forgetting a range releases its unused consumers; pending fetches and +//! other splits retain their own references until they finish. + +use std::collections::VecDeque; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::LazyLock; +use std::sync::OnceLock; +use std::sync::Weak; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::task::Wake; +use std::task::Waker; +use std::time::Instant; + +use futures::FutureExt; +use futures::StreamExt; +use futures::channel::mpsc; +use futures::channel::oneshot; +use futures::future::BoxFuture; +use futures::future::Shared; +use futures::stream::FuturesUnordered; +use parking_lot::Mutex; +use vortex_array::buffer::BufferHandle; +use vortex_buffer::Alignment; +use vortex_error::VortexError; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_io::request::Completion; +use vortex_io::request::IoBatch; +use vortex_io::request::IoIntent; +use vortex_io::request::IoOwnerId; +use vortex_io::request::IoRequest; +use vortex_io::request::IoRequestId; +use vortex_io::request::IoResult; +use vortex_io::request::IoService; +use vortex_io::request::IoSource; +use vortex_io::request::IoTarget; +use vortex_io::request::trace::next_id as next_trace_id; +use vortex_io::request::trace::timestamp_ns; +use vortex_utils::aliases::dash_map::DashMap; +use vortex_utils::aliases::hash_map::Entry; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::read::ReadRequest; +use crate::read::RequestId; +use crate::segments::ReadEvent; +use crate::segments::source::SharedDriver; + +/// A read's bytes, shared by every fetch of its range. +type SharedRead = Shared; +type ReadKey = (u64, usize, Alignment); + +/// Stored directly in the shared future's allocation. +struct ReadBytes { + receiver: oneshot::Receiver>, +} + +impl Future for ReadBytes { + type Output = Result>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + self.get_mut() + .receiver + .poll_unpin(cx) + .map(|result| match result { + Ok(result) => result.map_err(Arc::new), + Err(_) => Err(Arc::new(vortex_err!( + "the file's read driver dropped the read" + ))), + }) + } +} + +/// Serves a scan's splits from a file's coalescing read driver, one IO session per split. +/// +/// Every split's reads go to the same driver, so reads of different splits coalesce with each +/// other as the default scan's do. +#[derive(Clone)] +pub struct FileScanIo { + trace_source: u64, + uri: Option>, + events: mpsc::UnboundedSender, + /// The read driver's task, which runs while any handle to it is held: the file's segment + /// source may be dropped while the scan's splits still read. + _driver: SharedDriver, + next_id: Arc, + /// Live registrations held by announcements, prefetches, or pending fetches. Weak entries + /// deduplicate overlapping reads without retaining consumed bytes. + /// Initialized only for direct reads; scans with a segment cache use another source. + reads: Arc>>>, +} + +impl FileScanIo { + pub(crate) fn new( + events: mpsc::UnboundedSender, + driver: SharedDriver, + next_id: Arc, + uri: Option>, + trace_source: u64, + ) -> Self { + Self { + uri, + trace_source, + events, + _driver: driver, + next_id, + reads: Arc::default(), + } + } +} + +impl IoService for FileScanIo { + fn session(&self) -> Arc { + static INLINE_FETCH: LazyLock = LazyLock::new(|| { + !std::env::var("VORTEX_SCAN_IO_INLINE_FETCH").is_ok_and(|value| value == "0") + }); + if *INLINE_FETCH { + Arc::new(FileSplitIo::::new(self.clone())) + } else { + Arc::new(FileSplitIo::>::new( + self.clone(), + )) + } + } +} + +trait FetchFuture: Future + Unpin + Send + 'static { + const INLINE: bool; + + fn new(read: Arc, owner: IoOwnerId, request: IoRequestId) -> Self; +} + +impl FetchFuture for BoxFuture<'static, Completion> { + const INLINE: bool = false; + + fn new(read: Arc, owner: IoOwnerId, request: IoRequestId) -> Self { + async move { + let bytes = read.bytes.clone().await; + drop(read); + fetch_completion(owner, request, bytes) + } + .boxed() + } +} + +/// Lives inside `FuturesUnordered`'s task allocation, avoiding a second allocation per Fetch. +struct InlineFetch { + bytes: SharedRead, + _read: Arc, + owner: IoOwnerId, + request: IoRequestId, +} + +impl FetchFuture for InlineFetch { + const INLINE: bool = true; + + fn new(read: Arc, owner: IoOwnerId, request: IoRequestId) -> Self { + Self { + bytes: read.bytes.clone(), + _read: read, + owner, + request, + } + } +} + +impl Future for InlineFetch { + type Output = Completion; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.get_mut(); + let owner = this.owner; + let request = this.request; + this.bytes + .poll_unpin(cx) + .map(|bytes| fetch_completion(owner, request, bytes)) + } +} + +fn fetch_completion( + owner: IoOwnerId, + request: IoRequestId, + bytes: Result>, +) -> Completion { + Completion { + owner, + request, + result: bytes + .map(IoResult::Bytes) + .map_err(|err| vortex_err!("{err}")), + } +} + +/// One split's reads through a [`FileScanIo`]. +struct FileSplitIo { + trace_session: u64, + ready_fetch: bool, + io: FileScanIo, + state: Mutex>, +} + +struct SplitState { + /// Registered consumers not yet consumed by a fetch, by byte range. + reads: HashMap, + /// Fetches waiting for their read, in the order they complete. + fetches: Option>, + ready: VecDeque, + // Register under the queue lock; submission takes the waker and wakes after unlocking. + ready_waker: Option, + stats: ScopeStats, +} + +impl Default for SplitState { + fn default() -> Self { + Self { + reads: HashMap::default(), + fetches: None, + ready: VecDeque::new(), + ready_waker: None, + stats: ScopeStats::default(), + } + } +} + +struct Interest { + read: Arc, + fetched: bool, + consumers: usize, +} + +#[derive(Default)] +struct ScopeStats { + announcements: usize, + prefetches: usize, + fetches: usize, + inline_fetches: usize, + ready_fetches: usize, + future_queues: usize, + new_registrations: usize, + reused_registrations: usize, + peak_interests: usize, + polls: usize, + completion_polls: usize, + completions: usize, + forgets: usize, + forgotten: usize, + forgotten_bytes: u64, + forget_fetch_pins: usize, +} + +struct Read { + id: RequestId, + trace_source: u64, + /// Whether the driver was told the read is wanted. + wanted: AtomicBool, + bytes: SharedRead, + events: mpsc::UnboundedSender, +} + +impl Drop for Read { + fn drop(&mut self) { + if tracing::enabled!(target: "vortex_file::scan_lifetime", tracing::Level::DEBUG) { + tracing::debug!(target: "vortex_file::scan_lifetime", + ts_ns = timestamp_ns(), event = "registration_drop", source = self.trace_source, + read = self.id, result_observed = self.bytes.peek().is_some(), "scan IO"); + } + if self.bytes.peek().is_none() { + // Only the last split can cancel a shared registration. + drop(self.events.unbounded_send(ReadEvent::Dropped(self.id))); + } + } +} + +impl FileSplitIo { + fn new(io: FileScanIo) -> Self { + static READY_FETCH: LazyLock = LazyLock::new(|| { + std::env::var("VORTEX_SCAN_IO_READY_FETCH").is_ok_and(|value| value == "1") + }); + Self { + trace_session: next_trace_id(), + ready_fetch: *READY_FETCH, + io, + state: Mutex::default(), + } + } + + /// Registers a read of `offset..offset + len`, aligned to `alignment`, with the driver. + fn register( + &self, + offset: u64, + len: usize, + alignment: Alignment, + ) -> VortexResult<(Arc, bool)> { + let key = (offset, len, alignment); + // Keep this range locked until Request is queued, so another split cannot send Polled + // before its registration. Other ranges can register through different shards. + let mut entry = self + .io + .reads + .get_or_init(DashMap::default) + .entry(key) + .or_default(); + if let Some(read) = entry.upgrade() { + return Ok((read, false)); + } + let id = self.io.next_id.fetch_add(1, Ordering::Relaxed); + let (callback, receiver) = oneshot::channel(); + self.io + .events + .unbounded_send(ReadEvent::Request(ReadRequest { + id, + offset, + length: len, + alignment, + callback, + })) + .map_err(|err| vortex_err!("the file's read driver has stopped: {err}"))?; + let bytes = ReadBytes { receiver }.shared(); + let read = Arc::new(Read { + id, + trace_source: self.io.trace_source, + wanted: AtomicBool::new(false), + bytes, + events: self.io.events.clone(), + }); + *entry = Arc::downgrade(&read); + Ok((read, true)) + } + + /// Tells the driver `read` is wanted, so it starts the read if nothing has yet. + fn want(&self, read: &Read) -> VortexResult { + if read.wanted.load(Ordering::Relaxed) { + return Ok(false); + } + let wanted = !read.wanted.swap(true, Ordering::Relaxed); + if wanted { + self.io + .events + .unbounded_send(ReadEvent::Polled(read.id)) + .map_err(|err| vortex_err!("the file's read driver has stopped: {err}"))?; + } + Ok(wanted) + } + + fn clear_scope(&self, reason: &str) { + let mut state = self.state.lock(); + if !state.reads.is_empty() + && tracing::enabled!(target: "vortex_file::scan_lifetime", tracing::Level::DEBUG) + { + let unfetched = state.reads.values().filter(|i| !i.fetched).count(); + let unfetched_bytes: usize = state + .reads + .iter() + .filter(|(_, i)| !i.fetched) + .map(|(key, _)| key.1) + .sum(); + tracing::debug!(target: "vortex_file::scan_lifetime", + ts_ns = timestamp_ns(), event = "scope_clear", session = self.trace_session, + source = self.io.trace_source, reason, + interests = state.reads.len(), unfetched, unfetched_bytes, + announcements = state.stats.announcements, prefetches = state.stats.prefetches, + fetches = state.stats.fetches, new_registrations = state.stats.new_registrations, + inline_fetches = state.stats.inline_fetches, + ready_fetches = state.stats.ready_fetches, + future_queues = state.stats.future_queues, + reused_registrations = state.stats.reused_registrations, + peak_interests = state.stats.peak_interests, polls = state.stats.polls, + completion_polls = state.stats.completion_polls, completions = state.stats.completions, + forgets = state.stats.forgets, forgotten = state.stats.forgotten, + forgotten_bytes = state.stats.forgotten_bytes, + forget_fetch_pins = state.stats.forget_fetch_pins, + "scan IO"); + } + state.fetches = None; + state.ready.clear(); + state.ready_waker = None; + state.reads.clear(); + state.stats = ScopeStats::default(); + } +} + +// Tracing lives in these helpers so the request handling in `submit` stays readable. +impl FileSplitIo { + fn trace_forget(&self, interest: &Interest, offset: u64, len: usize) { + if tracing::enabled!(target: "vortex_file::scan_lifetime", tracing::Level::DEBUG) { + tracing::debug!(target: "vortex_file::scan_lifetime", + ts_ns = timestamp_ns(), event = "scope_forget", session = self.trace_session, + source = self.io.trace_source, read = interest.read.id, + offset, length = len, references = Arc::strong_count(&interest.read), + wanted = interest.read.wanted.load(Ordering::Relaxed), + result_observed = interest.read.bytes.peek().is_some(), "scan IO"); + } + } + + fn trace_binding( + &self, + owner: IoOwnerId, + request: &IoRequest, + read: &Read, + offset: u64, + len: usize, + ) { + if tracing::enabled!(target: "vortex_scan::driver", tracing::Level::DEBUG) { + tracing::debug!(target: "vortex_scan::driver", + ts_ns = timestamp_ns(), event = "binding", + session = self.trace_session, owner = owner.0, + request = request.request.0, intent = ?request.intent, + source = self.io.trace_source, read = read.id, + uri = self.io.uri.as_deref().unwrap_or(""), offset, length = len, + "scan IO"); + } + } +} + +fn trace_submit(requests: usize, registered: usize, wanted: usize, start: Instant) { + tracing::debug!( + target: "vortex_file::io_submit", + requests, + registered, + wanted, + elapsed_ns = u64::try_from(start.elapsed().as_nanos()).unwrap_or(u64::MAX), + "scan IO submit" + ); +} + +impl Drop for FileSplitIo { + fn drop(&mut self) { + self.clear_scope("drop"); + } +} + +struct TraceWake { + session: u64, + waker: Waker, +} + +impl Wake for TraceWake { + fn wake(self: Arc) { + self.wake_by_ref(); + } + + fn wake_by_ref(self: &Arc) { + tracing::debug!(target: "vortex_file::scan_lifetime", + ts_ns = timestamp_ns(), event = "session_wake", session = self.session, "scan IO"); + self.waker.wake_by_ref(); + } +} + +impl IoSource for FileSplitIo { + fn trace_id(&self) -> Option { + Some(self.trace_session) + } + + fn submit(&self, owner: IoOwnerId, batch: IoBatch) -> VortexResult<()> { + let timing = tracing::enabled!(target: "vortex_file::io_submit", tracing::Level::DEBUG) + .then(Instant::now); + let requests = batch.len(); + let mut state = self.state.lock(); + let mut registered = 0usize; + let mut made_ready = false; + let SplitState { + reads, + fetches, + ready, + stats, + .. + } = &mut *state; + let mut wanted = 0usize; + for request in batch { + let IoTarget::Range { + offset, + len, + alignment, + } = request.target + else { + vortex_bail!("a split's IO serves byte ranges, not {:?}", request.target); + }; + let key = (offset, len, alignment); + if request.intent == IoIntent::Forget { + stats.forgets += 1; + if let Some(interest) = reads.remove(&key) { + stats.forget_fetch_pins += + usize::from(interest.fetched && interest.read.bytes.peek().is_none()); + stats.forgotten += 1; + stats.forgotten_bytes += len as u64; + self.trace_forget(&interest, offset, len); + // Pending fetches and other splits retain their own Arcs; withdrawal only + // removes consumers which have not yet fetched this scope's registration. + drop(interest); + } + continue; + } + let interest = match reads.entry(key) { + Entry::Occupied(entry) => { + stats.reused_registrations += 1; + entry.into_mut() + } + Entry::Vacant(entry) => { + registered += 1; + let (read, fresh) = self.register(offset, len, alignment)?; + stats.new_registrations += usize::from(fresh); + stats.reused_registrations += usize::from(!fresh); + entry.insert(Interest { + read, + fetched: false, + consumers: 0, + }) + } + }; + match request.intent { + IoIntent::Announce => interest.consumers += 1, + IoIntent::Prefetch => interest.consumers = interest.consumers.max(1), + IoIntent::Fetch => {} + IoIntent::Forget => unreachable!("handled before registration"), + } + let read = &interest.read; + self.trace_binding(owner, &request, read, offset, len); + if request.intent != IoIntent::Announce { + wanted += usize::from(self.want(read)?); + } + match request.intent { + IoIntent::Announce => stats.announcements += 1, + IoIntent::Prefetch => stats.prefetches += 1, + IoIntent::Fetch => stats.fetches += 1, + IoIntent::Forget => unreachable!("handled before registration"), + } + if request.intent == IoIntent::Fetch { + interest.fetched = true; + interest.consumers = interest.consumers.saturating_sub(1); + let read = if interest.consumers == 0 { + reads.remove(&key).vortex_expect("registered interest").read + } else { + Arc::clone(&interest.read) + }; + if self.ready_fetch + && let Some(bytes) = read.bytes.peek() + { + ready.push_back(fetch_completion(owner, request.request, bytes.clone())); + stats.ready_fetches += 1; + made_ready = true; + } else { + stats.inline_fetches += usize::from(F::INLINE); + stats.future_queues += usize::from(fetches.is_none()); + fetches + .get_or_insert_with(FuturesUnordered::new) + .push(F::new(read, owner, request.request)); + } + } + } + state.stats.peak_interests = state.stats.peak_interests.max(state.reads.len()); + let ready_waker = made_ready.then(|| state.ready_waker.take()).flatten(); + drop(state); + if let Some(waker) = ready_waker { + waker.wake(); + } + if let Some(start) = timing { + trace_submit(requests, registered, wanted, start); + } + Ok(()) + } + + fn poll(&self) -> VortexResult> { + // The split polls again with its own waker before it waits, so a fetch that finishes + // after this check is not missed. + let mut cx = Context::from_waker(Waker::noop()); + let tracing = + tracing::enabled!(target: "vortex_file::scan_lifetime", tracing::Level::DEBUG); + if tracing { + tracing::debug!(target: "vortex_file::scan_lifetime", + ts_ns = timestamp_ns(), event = "session_poll_begin", session = self.trace_session, + method = "sweep", "scan IO"); + } + let mut state = self.state.lock(); + state.stats.polls += 1; + let next = state.ready.pop_front().map_or_else( + || { + state.fetches.as_mut().map_or(Poll::Ready(None), |fetches| { + fetches.poll_next_unpin(&mut cx) + }) + }, + |completion| Poll::Ready(Some(completion)), + ); + let result = match next { + Poll::Ready(completion) => { + state.stats.completions += usize::from(completion.is_some()); + state.ready_waker = None; + Ok(completion) + } + Poll::Pending => Ok(None), + }; + if tracing { + tracing::debug!(target: "vortex_file::scan_lifetime", + ts_ns = timestamp_ns(), event = "session_poll", session = self.trace_session, + method = "sweep", outcome = if result.as_ref().is_ok_and(|r| r.is_some()) { "ready" } else { "pending" }, + "scan IO"); + } + result + } + + fn poll_completion(&self, cx: &mut Context<'_>) -> Poll> { + let tracing = + tracing::enabled!(target: "vortex_file::scan_lifetime", tracing::Level::DEBUG); + if tracing { + tracing::debug!(target: "vortex_file::scan_lifetime", + ts_ns = timestamp_ns(), event = "session_poll_begin", session = self.trace_session, + method = "completion", "scan IO"); + } + let waker = tracing.then(|| { + Waker::from(Arc::new(TraceWake { + session: self.trace_session, + waker: cx.waker().clone(), + })) + }); + let mut state = self.state.lock(); + state.stats.completion_polls += 1; + let next = if let Some(completion) = state.ready.pop_front() { + Poll::Ready(Some(completion)) + } else { + state.fetches.as_mut().map_or(Poll::Ready(None), |fetches| { + if let Some(waker) = waker.as_ref() { + fetches.poll_next_unpin(&mut Context::from_waker(waker)) + } else { + fetches.poll_next_unpin(cx) + } + }) + }; + let result = match next { + Poll::Ready(Some(completion)) => { + state.stats.completions += 1; + Poll::Ready(Ok(completion)) + } + Poll::Ready(None) => Poll::Ready(Err(vortex_err!( + "the split's driver is waiting with no fetch in flight" + ))), + Poll::Pending => { + if self.ready_fetch { + let ready_waker = waker.as_ref().unwrap_or_else(|| cx.waker()); + if !state + .ready_waker + .as_ref() + .is_some_and(|stored| stored.will_wake(ready_waker)) + { + state.ready_waker = Some(ready_waker.clone()); + } + } + Poll::Pending + } + }; + if result.is_ready() { + state.ready_waker = None; + } + if tracing { + let outcome = match &result { + Poll::Ready(Ok(_)) => "ready", + Poll::Ready(Err(_)) => "error", + Poll::Pending => "pending", + }; + tracing::debug!(target: "vortex_file::scan_lifetime", + ts_ns = timestamp_ns(), event = "session_poll", session = self.trace_session, + method = "completion", outcome, "scan IO"); + } + result + } + + fn wait(&self) -> VortexResult { + vortex_bail!("a split's IO is awaited through poll_completion, not waited on") + } + + fn release(&self, _owner: IoOwnerId) {} + + fn clear(&self) { + self.clear_scope("clear"); + } +} + +#[cfg(test)] +mod tests { + use std::future::poll_fn; + use std::sync::Arc; + use std::sync::atomic::AtomicBool; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use std::task::Context; + use std::task::Poll; + use std::task::Wake; + use std::task::Waker; + + use futures::FutureExt; + use futures::StreamExt; + use futures::channel::mpsc; + use futures::channel::oneshot; + use futures::future; + use futures::future::BoxFuture; + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::StructArray; + use vortex_array::buffer::BufferHandle; + use vortex_buffer::Alignment; + use vortex_buffer::Buffer; + use vortex_buffer::ByteBuffer; + use vortex_buffer::ByteBufferMut; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use vortex_io::VortexReadAt; + use vortex_io::request::Completion; + use vortex_io::request::IoIntent; + use vortex_io::request::IoOwnerId; + use vortex_io::request::IoRequest; + use vortex_io::request::IoRequestId; + use vortex_io::request::IoSource; + use vortex_io::request::IoTarget; + use vortex_io::session::RuntimeSession; + use vortex_io::session::RuntimeSessionExt; + use vortex_layout::session::LayoutSession; + use vortex_metrics::DefaultMetricsRegistry; + use vortex_session::VortexSession; + + use super::FetchFuture; + use super::FileScanIo; + use super::FileSplitIo; + use super::InlineFetch; + use super::Interest; + use super::Read; + use super::ReadBytes; + use super::SharedRead; + use crate::OpenOptionsSessionExt; + use crate::WriteOptionsSessionExt; + use crate::read::IoRequestStream; + use crate::segments::RequestMetrics; + + fn ready_source(failed: bool) -> FileSplitIo { + let (events, _) = mpsc::unbounded(); + let io = FileScanIo::new( + events, + future::pending().boxed().shared(), + Arc::new(AtomicUsize::new(0)), + None, + 0, + ); + let mut source = FileSplitIo::new(io); + source.ready_fetch = true; + let result = if failed { + Err(vortex_err!("read failed")) + } else { + Ok(BufferHandle::new_host(ByteBuffer::from_iter([ + 1u8, 2, 3, 4, + ]))) + }; + let (callback, receiver) = oneshot::channel(); + assert!(callback.send(result).is_ok()); + let bytes = ReadBytes { receiver }.shared(); + drop(futures::executor::block_on(bytes.clone())); + source.state.lock().reads.insert( + (0, 4, Alignment::none()), + Interest { + read: Arc::new(Read { + id: 0, + trace_source: 0, + wanted: AtomicBool::new(true), + bytes, + events: source.io.events.clone(), + }), + fetched: false, + consumers: 2, + }, + ); + source + } + + #[rstest] + #[case::bytes(false)] + #[case::error(true)] + fn ready_fetch_delivers_without_a_future_and_clear_discards_the_queue( + #[case] failed: bool, + ) -> VortexResult<()> { + let source = ready_source(failed); + source.submit( + IoOwnerId(5), + (10..12) + .map(|id| IoRequest { + intent: IoIntent::Fetch, + request: IoRequestId(id), + target: IoTarget::range(0, 4), + }) + .collect(), + )?; + { + let state = source.state.lock(); + assert!(state.fetches.is_none()); + assert_eq!(state.ready.len(), 2); + assert_eq!(state.stats.ready_fetches, 2); + assert_eq!(state.stats.future_queues, 0); + } + let completion = source + .poll()? + .ok_or_else(|| vortex_err!("missing ready fetch"))?; + assert_eq!(completion.owner, IoOwnerId(5)); + assert_eq!(completion.request, IoRequestId(10)); + assert_eq!(completion.result.is_err(), failed); + if let Err(error) = completion.result { + assert!(error.to_string().contains("read failed")); + } + source.clear(); + assert!(source.poll()?.is_none()); + assert!(source.state.lock().reads.is_empty()); + Ok(()) + } + + #[test] + fn ready_fetch_wakes_a_consumer_waiting_for_another_read() -> VortexResult<()> { + let source = ready_source(false); + let (_callback, receiver) = oneshot::channel(); + source.state.lock().reads.insert( + (4, 4, Alignment::none()), + Interest { + read: Arc::new(Read { + id: 1, + trace_source: 0, + wanted: AtomicBool::new(true), + bytes: ReadBytes { receiver }.shared(), + events: source.io.events.clone(), + }), + fetched: false, + consumers: 1, + }, + ); + let request = |offset| IoRequest { + intent: IoIntent::Fetch, + request: IoRequestId(0), + target: IoTarget::range(offset, 4), + }; + source.submit(IoOwnerId(1), vec![request(4)])?; + let count = Arc::new(CountingWake(AtomicUsize::new(0))); + let waker = Waker::from(Arc::clone(&count)); + let mut cx = Context::from_waker(&waker); + assert!(source.poll_completion(&mut cx).is_pending()); + assert_eq!(source.state.lock().stats.future_queues, 1); + let before = count.0.load(Ordering::Relaxed); + source.submit(IoOwnerId(2), vec![request(0)])?; + assert!(count.0.load(Ordering::Relaxed) > before); + let Poll::Ready(completion) = source.poll_completion(&mut cx) else { + return Err(vortex_err!("ready fetch did not become ready")); + }; + assert_eq!(completion?.owner, IoOwnerId(2)); + assert!(source.poll_completion(&mut cx).is_pending()); + source.clear(); + assert!(source.state.lock().fetches.is_none()); + Ok(()) + } + + struct CountingWake(AtomicUsize); + + impl Wake for CountingWake { + fn wake(self: Arc) { + self.wake_by_ref(); + } + + fn wake_by_ref(self: &Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + + #[test] + fn traced_wake_forwards_each_notification() { + let count = Arc::new(CountingWake(AtomicUsize::new(0))); + let traced = Waker::from(Arc::new(super::TraceWake { + session: 0, + waker: Waker::from(Arc::clone(&count)), + })); + traced.wake_by_ref(); + traced.wake(); + assert_eq!(count.0.load(Ordering::Relaxed), 2); + } + + fn fetch_wakes_and_forwards_errors(cancelled: bool) -> VortexResult<()> { + let (send, receiver) = oneshot::channel(); + let bytes: SharedRead = ReadBytes { receiver }.shared(); + let (events, _) = mpsc::unbounded(); + let read = Arc::new(Read { + id: 0, + trace_source: 0, + wanted: AtomicBool::new(true), + bytes, + events, + }); + let mut fetch = F::new(read, IoOwnerId(3), IoRequestId(7)); + let count = Arc::new(CountingWake(AtomicUsize::new(0))); + let waker = Waker::from(Arc::clone(&count)); + let mut cx = Context::from_waker(&waker); + assert!(fetch.poll_unpin(&mut cx).is_pending()); + if cancelled { + drop(send); + } else { + send.send(Err(vortex_err!("read failed"))) + .map_err(|_| vortex_err!("fetch receiver was dropped"))?; + } + assert!(count.0.load(Ordering::Relaxed) > 0); + let Poll::Ready(completion) = fetch.poll_unpin(&mut cx) else { + return Err(vortex_err!("fetch did not become ready")); + }; + assert_eq!(completion.owner, IoOwnerId(3)); + assert_eq!(completion.request, IoRequestId(7)); + let error = completion + .result + .err() + .ok_or_else(|| vortex_err!("fetch did not forward the read error"))?; + assert!(error.to_string().contains(if cancelled { + "the file's read driver dropped the read" + } else { + "read failed" + })); + Ok(()) + } + + #[rstest] + #[case::boxed_error(false, false)] + #[case::inline_error(true, false)] + #[case::boxed_cancelled(false, true)] + #[case::inline_cancelled(true, true)] + fn fetch_future_preserves_notifications_and_errors( + #[case] inline: bool, + #[case] cancelled: bool, + ) -> VortexResult<()> { + if inline { + fetch_wakes_and_forwards_errors::(cancelled) + } else { + fetch_wakes_and_forwards_errors::>(cancelled) + } + } + + fn session() -> VortexSession { + let session = vortex_array::array_session() + .with::() + .with::() + .with_tokio(); + crate::register_default_encodings(&session); + crate::enable_all_registered_array_encodings(&session); + session + } + + /// Four 50,000-row chunks of `{a, b}`, with `a` over `0..200_000` and `b = a % 97`. + async fn write(session: &VortexSession) -> VortexResult { + let chunks = (0..4) + .map(|chunk| { + let rows = chunk * 50_000..(chunk + 1) * 50_000; + let a = Buffer::from_iter(rows.clone()).into_array(); + let b = Buffer::from_iter(rows.map(|a| a % 97)).into_array(); + Ok(StructArray::from_fields(&[("a", a), ("b", b)])?.into_array()) + }) + .collect::>>()?; + let dtype = chunks[0].dtype().clone(); + let mut output = ByteBufferMut::empty(); + session + .write_options() + .write( + &mut output, + ChunkedArray::try_new(chunks, dtype)? + .into_array() + .to_array_stream(), + ) + .await?; + Ok(ByteBuffer::from(output)) + } + + #[derive(Clone)] + struct CountingRead { + bytes: ByteBuffer, + reads: Arc, + } + + impl VortexReadAt for CountingRead { + fn size(&self) -> BoxFuture<'static, VortexResult> { + self.bytes.size() + } + + fn concurrency(&self) -> usize { + self.bytes.concurrency() + } + + fn read_at( + &self, + offset: u64, + length: usize, + alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + self.reads.fetch_add(1, Ordering::Relaxed); + self.bytes.read_at(offset, length, alignment) + } + } + + /// A cancelled announcement cannot cancel pending fetches. Concurrent fetches share a read, + /// but later fetches cannot reuse its completed bytes. + #[tokio::test] + async fn splits_share_live_reads_and_cancel_only_the_last_registration() -> VortexResult<()> { + let session = session(); + let reads = Arc::new(AtomicUsize::new(0)); + let file = session + .open_options() + .open_read(CountingRead { + bytes: write(&session).await?, + reads: Arc::clone(&reads), + }) + .await?; + reads.store(0, Ordering::Relaxed); + let service = file + .scan_io() + .ok_or_else(|| vortex_err!("missing file IO"))?; + let spec = file.footer().segment_map()[0]; + let request = |intent, id| IoRequest { + intent, + request: IoRequestId(id), + target: IoTarget::Range { + offset: spec.offset, + len: spec.length as usize, + alignment: spec.alignment, + }, + }; + let announced = service.session(); + let fetched = service.session(); + let concurrent = service.session(); + announced.submit(IoOwnerId(0), vec![request(IoIntent::Announce, 0)])?; + fetched.submit(IoOwnerId(1), vec![request(IoIntent::Fetch, 1)])?; + concurrent.submit(IoOwnerId(4), vec![request(IoIntent::Fetch, 4)])?; + announced.submit(IoOwnerId(0), vec![request(IoIntent::Forget, 5)])?; + fetched.submit(IoOwnerId(1), vec![request(IoIntent::Forget, 6)])?; + poll_fn(|cx| fetched.poll_completion(cx)).await?.result?; + poll_fn(|cx| concurrent.poll_completion(cx)).await?.result?; + assert_eq!(reads.load(Ordering::Relaxed), 1); + + let later = service.session(); + later.submit(IoOwnerId(2), vec![request(IoIntent::Fetch, 2)])?; + poll_fn(|cx| later.poll_completion(cx)).await?.result?; + assert_eq!(reads.load(Ordering::Relaxed), 2); + fetched.clear(); + later.clear(); + + let fresh = service.session(); + fresh.submit(IoOwnerId(3), vec![request(IoIntent::Fetch, 3)])?; + poll_fn(|cx| fresh.poll_completion(cx)).await?.result?; + assert_eq!(reads.load(Ordering::Relaxed), 3); + Ok(()) + } + + #[tokio::test] + async fn registered_consumers_reuse_only_their_pending_read() -> VortexResult<()> { + let session = session(); + let reads = Arc::new(AtomicUsize::new(0)); + let file = session + .open_options() + .open_read(CountingRead { + bytes: write(&session).await?, + reads: Arc::clone(&reads), + }) + .await?; + reads.store(0, Ordering::Relaxed); + let service = file + .scan_io() + .ok_or_else(|| vortex_err!("missing file IO"))?; + let io = service.session(); + let spec = file.footer().segment_map()[0]; + let request = |intent, id| IoRequest { + intent, + request: IoRequestId(id), + target: IoTarget::Range { + offset: spec.offset, + len: spec.length as usize, + alignment: spec.alignment, + }, + }; + io.submit( + IoOwnerId(0), + vec![ + request(IoIntent::Announce, 0), + request(IoIntent::Announce, 1), + ], + )?; + for id in 2..5 { + io.submit(IoOwnerId(0), vec![request(IoIntent::Fetch, id)])?; + poll_fn(|cx| io.poll_completion(cx)).await?.result?; + assert_eq!(reads.load(Ordering::Relaxed), if id < 4 { 1 } else { 2 }); + } + Ok(()) + } + + #[tokio::test] + async fn forgotten_interest_is_absent_from_later_coalescing() -> VortexResult<()> { + let (events, receiver) = mpsc::unbounded(); + let io = FileScanIo::new( + events, + future::pending().boxed().shared(), + Arc::default(), + None, + 0, + ); + let source = FileSplitIo::::new(io); + let request = |intent, id, offset| IoRequest { + intent, + request: IoRequestId(id), + target: IoTarget::range(offset, 4), + }; + source.submit(IoOwnerId(0), vec![request(IoIntent::Announce, 0, 0)])?; + source.submit(IoOwnerId(1), vec![request(IoIntent::Forget, 0, 0)])?; + source.submit(IoOwnerId(2), vec![request(IoIntent::Fetch, 0, 8)])?; + let registry = DefaultMetricsRegistry::default(); + let mut reads = IoRequestStream::new( + receiver.boxed(), + Some(vortex_io::CoalesceConfig::new(16, 32)), + Alignment::none(), + 1, + RequestMetrics::new(®istry, vec![]), + ); + let mut batch = reads + .next() + .await + .ok_or_else(|| vortex_err!("missing read"))?; + assert_eq!(batch.len(), 1); + let physical = batch.pop().ok_or_else(|| vortex_err!("missing read"))?; + assert_eq!(physical.offset(), 8); + assert_eq!(physical.len(), 4); + assert_eq!(physical.requests().len(), 1); + physical.resolve(Ok(BufferHandle::new_host(ByteBuffer::copy_from(b"data")))); + let completion = poll_fn(|cx| source.poll_completion(cx)).await?; + assert_eq!(completion.owner, IoOwnerId(2)); + let vortex_io::request::IoResult::Bytes(bytes) = completion.result? else { + return Err(vortex_err!("expected bytes")); + }; + assert_eq!(bytes.to_host().await.as_slice(), b"data"); + assert_eq!(source.state.lock().stats.forgotten, 1); + assert_eq!(source.state.lock().stats.forgotten_bytes, 4); + Ok(()) + } +} diff --git a/vortex-file/src/segments/source.rs b/vortex-file/src/segments/source.rs index 1ab7af31064..2684b58bc5e 100644 --- a/vortex-file/src/segments/source.rs +++ b/vortex-file/src/segments/source.rs @@ -6,6 +6,7 @@ use std::collections::VecDeque; use std::future::Future; use std::pin::Pin; use std::sync::Arc; +use std::sync::LazyLock; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; @@ -23,6 +24,7 @@ use futures::stream::BoxStream; use futures::stream::Fuse; use futures::stream::SelectAll; use parking_lot::Mutex; +use tracing::Instrument; use vortex_array::buffer::BufferHandle; use vortex_buffer::Alignment; use vortex_buffer::ByteBuffer; @@ -33,6 +35,8 @@ use vortex_error::vortex_panic; use vortex_io::ReadAtRequest; use vortex_io::ReadAtStream; use vortex_io::VortexReadAt; +use vortex_io::request::trace::next_id as next_trace_id; +use vortex_io::request::trace::timestamp_ns; use vortex_io::runtime::Handle; use vortex_io::runtime::JoinOutcome; use vortex_layout::segments::SegmentFuture; @@ -49,6 +53,7 @@ use crate::read::IoRequest; use crate::read::IoRequestStream; use crate::read::ReadRequest; use crate::read::RequestId; +use crate::segments::FileScanIo; #[derive(Debug)] /// Events sent from segment futures to the coalescing read driver. @@ -88,7 +93,7 @@ pub enum ReadEvent { /// reader is always woken when the driver finishes — even if another reader polled the driver more /// recently and was then dropped. Its output is `()`; a driver panic is carried out of band in /// [`DriverPanic`] so it can be re-raised on the reader side. -type SharedDriver = Shared>; +pub(crate) type SharedDriver = Shared>; /// Slot holding the driver's panic payload, if it panicked while driving reads. The first reader to /// observe completion takes the payload and re-raises it; later readers report a graceful error. @@ -111,7 +116,7 @@ fn validate_read_result( }) } -type IoBatchStream = Fuse>>; +type IoBatchStream = Fuse>>; enum ReadRangeResultsState { Reading(ReadAtStream), @@ -176,6 +181,7 @@ impl Stream for ReadRangeResults { /// Drives request batches while keeping the reader's concurrency slots occupied. struct ReadDriver { + trace_source: u64, reader: Arc, batches: IoBatchStream, pending: VecDeque, @@ -183,17 +189,20 @@ struct ReadDriver { num_active: usize, batches_done: bool, concurrency: usize, + lookahead: bool, metrics: RequestMetrics, } impl ReadDriver { fn new( reader: R, - batches: BoxStream<'static, Vec>, + batches: IoRequestStream>, concurrency: usize, metrics: RequestMetrics, + trace_source: u64, ) -> Self { Self { + trace_source, reader: Arc::new(reader), batches: batches.fuse(), pending: VecDeque::new(), @@ -201,6 +210,7 @@ impl ReadDriver { num_active: 0, batches_done: false, concurrency, + lookahead: std::env::var("VORTEX_SCAN_IO_LOOKAHEAD").is_ok_and(|value| value == "1"), metrics, } } @@ -209,6 +219,12 @@ impl ReadDriver { 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::>(); + trace_read_start( + self.trace_source, + self.reader.uri().map_or("", |uri| uri.as_ref()), + self.pending.len(), + &reqs, + ); self.num_active += batch_len; self.metrics.read_ranges_calls.add(1); @@ -216,12 +232,7 @@ impl ReadDriver { 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" - ); + trace_read_batch(batch_len, self.num_active); let requests = reqs .iter() @@ -240,11 +251,21 @@ impl Stream for ReadDriver { 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. + // Choose physical ranges only for free slots. Later registrations and cancellations must + // stay visible to the coalescer until their reads can actually start. if !this.batches_done { - loop { + let available = this.concurrency - this.num_active; + while this.lookahead || this.pending.len() < available { + this.batches.get_mut().set_batch_size(if this.lookahead { + this.concurrency + } else { + available - this.pending.len() + }); match this.batches.poll_next_unpin(cx) { - Poll::Ready(Some(batch)) => this.pending.extend(batch), + Poll::Ready(Some(batch)) => { + trace_read_queued(this.trace_source, this.pending.len(), &batch); + this.pending.extend(batch); + } Poll::Ready(None) => { this.batches_done = true; break; @@ -262,6 +283,7 @@ impl Stream for ReadDriver { match this.reads.poll_next_unpin(cx) { Poll::Ready(Some((req, result))) => { + trace_read_end(this.trace_source, &req, result.is_ok()); this.num_active -= 1; let result = validate_read_result(&req, result); req.resolve(result); @@ -276,7 +298,62 @@ impl Stream for ReadDriver { } } +// The read driver's tracing lives in these helpers so its scheduling stays readable. + +/// Traces the start of each physical read in `reqs`, and the logical reads it serves. +fn trace_read_start(trace_source: u64, uri: &str, queued: usize, reqs: &[IoRequest]) { + if !tracing::enabled!(target: "vortex_scan::driver", tracing::Level::DEBUG) { + return; + } + for req in reqs { + let physical = req.requests()[0].id; + tracing::debug!(target: "vortex_scan::driver", + ts_ns = timestamp_ns(), event = "read_start", + source = trace_source, physical, offset = req.offset(), + queued, length = req.len(), uri, "scan IO"); + for member in req.requests() { + tracing::debug!(target: "vortex_scan::driver", + ts_ns = timestamp_ns(), event = "read_member", + source = trace_source, physical, read = member.id, + offset = member.offset, length = member.length, "scan IO"); + } + } +} + +fn trace_read_batch(num_ranges: usize, num_active: usize) { + tracing::trace!( + target: "vortex_file::read_ranges", + num_ranges, + num_active, + "submitting positional read batch" + ); +} + +/// Traces each physical read of `batch` joining the `queued` reads already waiting. +fn trace_read_queued(trace_source: u64, queued: usize, batch: &[IoRequest]) { + if !tracing::enabled!(target: "vortex_scan::driver", tracing::Level::DEBUG) { + return; + } + for req in batch { + tracing::debug!(target: "vortex_scan::driver", + ts_ns = timestamp_ns(), event = "read_queued", + source = trace_source, physical = req.requests()[0].id, + queued = queued + batch.len(), "scan IO"); + } +} + +fn trace_read_end(trace_source: u64, req: &IoRequest, success: bool) { + if tracing::enabled!(target: "vortex_scan::driver", tracing::Level::DEBUG) { + tracing::debug!(target: "vortex_scan::driver", + ts_ns = timestamp_ns(), event = "read_end", + source = trace_source, physical = req.requests()[0].id, + success, "scan IO"); + } +} + pub struct FileSegmentSource { + trace_source: u64, + uri: Option>, segments: Arc<[SegmentSpec]>, /// A queue for sending read request events to the I/O stream. events: mpsc::UnboundedSender, @@ -307,6 +384,23 @@ impl FileSegmentSource { .max() .unwrap_or_else(Alignment::none); let coalesce_config = reader.coalesce_config().map(|mut config| { + static MAX_SIZE: LazyLock> = LazyLock::new(|| { + std::env::var("VORTEX_SCAN_IO_COALESCE_MAX_BYTES") + .ok() + .and_then(|value| value.parse().ok()) + .filter(|&value| value > 0) + }); + static DISTANCE: LazyLock> = LazyLock::new(|| { + std::env::var("VORTEX_SCAN_IO_COALESCE_DISTANCE_BYTES") + .ok() + .and_then(|value| value.parse().ok()) + }); + if let Some(max_size) = *MAX_SIZE { + config.max_size = max_size; + } + if let Some(distance) = *DISTANCE { + config.distance = distance; + } // Aligning the coalesced start down can add up to (alignment - 1) bytes. // Increase max_size to keep the effective payload window consistent. let extra = (max_alignment.as_usize() as u64).saturating_sub(1); @@ -320,6 +414,13 @@ impl FileSegmentSource { reader.uri() ); } + static READ_CONCURRENCY: LazyLock> = LazyLock::new(|| { + std::env::var("VORTEX_SCAN_IO_READ_CONCURRENCY") + .ok() + .and_then(|value| value.parse().ok()) + .filter(|&value| value > 0) + }); + let concurrency = READ_CONCURRENCY.map_or(concurrency, |limit| concurrency.min(limit)); let stream = IoRequestStream::new( StreamExt::boxed(recv), @@ -327,14 +428,19 @@ impl FileSegmentSource { max_alignment, concurrency, metrics.clone(), - ) - .boxed(); + ); - let drive_fut = ReadDriver::new(reader, stream, concurrency, metrics).collect::<()>(); + let uri = reader.uri().cloned(); + let next_id = Arc::new(AtomicUsize::new(0)); + let trace_source = next_trace_id(); + let span = tracing::debug_span!(target: "vortex_file::read_lifecycle", "file_reads", uri = ?reader.uri()); + let drive_fut = ReadDriver::new(reader, stream, concurrency, metrics, trace_source) + .collect::<()>() + .instrument(span); // 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. - let mut task = handle.spawn(drive_fut); + let mut task = handle.spawn_io(drive_fut); let driver_panic: DriverPanic = Arc::new(Mutex::new(None)); let driver = { let driver_panic = Arc::clone(&driver_panic); @@ -352,11 +458,13 @@ impl FileSegmentSource { }; Self { + uri, + trace_source, segments, events: send, driver, driver_panic, - next_id: Arc::new(AtomicUsize::new(0)), + next_id, } } } @@ -408,6 +516,20 @@ impl SegmentSource for FileSegmentSource { } } +impl FileSegmentSource { + /// A service that serves a scan's protocol requests through this source's read driver, so + /// they coalesce with each other and with this source's own reads. + pub fn scan_io(&self) -> FileScanIo { + FileScanIo::new( + self.events.clone(), + self.driver.clone(), + Arc::clone(&self.next_id), + self.uri.clone(), + self.trace_source, + ) + } +} + /// A future that resolves a read request from a [`FileSegmentSource`]. /// /// See the documentation for [`FileSegmentSource`] for details on coalescing and pre-fetching. @@ -563,15 +685,42 @@ impl SegmentSource for BufferSegmentSource { #[cfg(test)] mod tests { use std::panic::AssertUnwindSafe; + use std::task::Waker; use futures::future::BoxFuture; use vortex_error::vortex_bail; + use vortex_io::runtime::AbortHandleRef; + use vortex_io::runtime::Executor; use vortex_io::runtime::tokio::TokioRuntime; use vortex_layout::segments::SegmentSource; use vortex_metrics::DefaultMetricsRegistry; use super::*; + #[derive(Default)] + struct IoExecutor { + io_spawns: AtomicUsize, + } + + impl Executor for IoExecutor { + fn spawn(&self, future: BoxFuture<'static, ()>) -> AbortHandleRef { + Executor::spawn(&tokio::runtime::Handle::current(), future) + } + + fn spawn_io(&self, future: BoxFuture<'static, ()>) -> AbortHandleRef { + self.io_spawns.fetch_add(1, Ordering::Relaxed); + self.spawn(future) + } + + fn spawn_cpu(&self, task: Box) -> AbortHandleRef { + Executor::spawn_cpu(&tokio::runtime::Handle::current(), task) + } + + fn spawn_blocking_io(&self, task: Box) -> AbortHandleRef { + Executor::spawn_blocking_io(&tokio::runtime::Handle::current(), task) + } + } + fn io_request(id: RequestId, offset: u64, length: usize) -> IoRequest { let (callback, _receiver) = oneshot::channel(); IoRequest::new_single(ReadRequest { @@ -763,6 +912,28 @@ mod tests { } } + #[tokio::test] + async fn file_segment_source_driver_uses_io_executor() -> VortexResult<()> { + let executor = Arc::new(IoExecutor::default()); + let runtime = Arc::clone(&executor) as Arc; + let metrics = DefaultMetricsRegistry::default(); + let source = FileSegmentSource::open( + Arc::from([SegmentSpec { + offset: 0, + length: 4, + alignment: Alignment::none(), + }]), + ReadRangesOnly { + calls: Arc::default(), + }, + Handle::new(Arc::downgrade(&runtime)), + RequestMetrics::new(&metrics, vec![]), + ); + assert_eq!(executor.io_spawns.load(Ordering::Relaxed), 1); + assert_eq!(source.request(SegmentId::from(0)).await?.len(), 4); + Ok(()) + } + #[tokio::test] async fn read_driver_batches_ready_requests() -> VortexResult<()> { let calls = Arc::new(AtomicUsize::new(0)); @@ -911,6 +1082,78 @@ mod tests { Ok(()) } + #[rstest::rstest] + #[case(false, 2, 1)] + #[case(true, 3, 0)] + fn read_driver_coalesces_neighbors_registered_while_slots_are_full( + #[case] lookahead: bool, + #[case] expected_reads: usize, + #[case] expected_coalesced: u64, + ) -> VortexResult<()> { + let (send, recv) = mpsc::unbounded(); + let mut receivers = Vec::new(); + let register = |id, offset| -> VortexResult<_> { + let (callback, receiver) = oneshot::channel(); + send.unbounded_send(ReadEvent::Request(ReadRequest { + id, + offset, + length: 4, + alignment: Alignment::none(), + callback, + })) + .map_err(|err| vortex_err!("{err}"))?; + send.unbounded_send(ReadEvent::Polled(id)) + .map_err(|err| vortex_err!("{err}"))?; + Ok(receiver) + }; + receivers.push(register(0, 0)?); + receivers.push(register(1, 64)?); + let registry = DefaultMetricsRegistry::default(); + let metrics = RequestMetrics::new(®istry, vec![]); + let batches = IoRequestStream::new( + recv.boxed(), + Some(vortex_io::CoalesceConfig::new(4, 16)), + Alignment::none(), + 1, + metrics.clone(), + ); + let batch_sizes = Arc::new(Mutex::new(Vec::new())); + let permits = Arc::new(tokio::sync::Semaphore::new(0)); + let reader = ControlledReadRanges { + active: Arc::default(), + max_active: Arc::default(), + batch_sizes: Arc::clone(&batch_sizes), + permits: Arc::clone(&permits), + }; + let mut driver = ReadDriver::new(reader, batches, 1, metrics.clone(), 0); + driver.lookahead = lookahead; + let mut cx = Context::from_waker(Waker::noop()); + assert!(Pin::new(&mut driver).poll_next(&mut cx).is_pending()); + + // The neighbor arrives after the first read occupies the only slot. The next physical + // range must be chosen after that slot is freed, so it includes both neighboring reads. + receivers.push(register(2, 68)?); + drop(send); + permits.add_permits(3); + while matches!( + Pin::new(&mut driver).poll_next(&mut cx), + Poll::Ready(Some(())) + ) {} + assert_eq!(batch_sizes.lock().len(), expected_reads); + assert_eq!(metrics.coalesced_requests.value(), expected_coalesced); + for mut receiver in receivers { + assert_eq!( + receiver + .try_recv() + .map_err(|err| vortex_err!("{err}"))? + .vortex_expect("read completed")? + .len(), + 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)); diff --git a/vortex-io/Cargo.toml b/vortex-io/Cargo.toml index 49bb29eb876..2c77deed983 100644 --- a/vortex-io/Cargo.toml +++ b/vortex-io/Cargo.toml @@ -41,6 +41,9 @@ vortex-metrics = { workspace = true } vortex-session = { workspace = true } vortex-utils = { workspace = true } +[target.'cfg(target_os = "linux")'.dependencies] +rustix = { workspace = true } + [target.'cfg(unix)'.dependencies] custom-labels = { workspace = true } diff --git a/vortex-io/src/object_store/read_at.rs b/vortex-io/src/object_store/read_at.rs index a36c4517817..aa869241da4 100644 --- a/vortex-io/src/object_store/read_at.rs +++ b/vortex-io/src/object_store/read_at.rs @@ -1,8 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +#[cfg(not(target_arch = "wasm32"))] +use std::fs::File; use std::io; +use std::ops::Range; use std::sync::Arc; +#[cfg(not(target_arch = "wasm32"))] +use std::sync::LazyLock; +#[cfg(not(target_arch = "wasm32"))] +use std::time::Instant; +#[cfg(not(target_arch = "wasm32"))] +use std::time::SystemTime; +#[cfg(not(target_arch = "wasm32"))] +use std::time::UNIX_EPOCH; use futures::FutureExt; use futures::SinkExt; @@ -10,12 +21,17 @@ 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::GetResultPayload; use object_store::ObjectStore; use object_store::ObjectStoreExt; use object_store::path::Path as ObjectPath; +#[cfg(not(target_arch = "wasm32"))] +use tokio::sync::OnceCell; +#[cfg(not(target_arch = "wasm32"))] +use tokio::sync::Semaphore; use vortex_array::buffer::BufferHandle; use vortex_array::memory::BufferAllocatorRef; use vortex_buffer::Alignment; @@ -28,6 +44,8 @@ use crate::CoalesceConfig; use crate::ReadAtRequest; use crate::ReadAtStream; use crate::VortexReadAt; +#[cfg(not(target_arch = "wasm32"))] +use crate::request::trace::timestamp_ns; use crate::runtime::Handle; #[cfg(not(target_arch = "wasm32"))] use crate::std_file::read_exact_at; @@ -35,8 +53,36 @@ use crate::std_file::read_exact_at; /// Default number of concurrent requests to allow. pub const DEFAULT_CONCURRENCY: usize = 192; +// Local object-store payloads use blocking pread rather than network GETs. An optional aggregate +// limit across files bounds allocation and blocking work; network payloads retain the GET limit. +#[cfg(not(target_arch = "wasm32"))] +fn local_read_limit() -> Option> { + static LIMIT: LazyLock>> = LazyLock::new(|| { + let concurrency = std::env::var("VORTEX_LOCAL_READ_CONCURRENCY") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(0); + (concurrency > 0).then(|| Arc::new(Semaphore::new(concurrency))) + }); + LIMIT.as_ref().map(Arc::clone) +} + /// An object store backed I/O source. +/// +/// A positive `VORTEX_LOCAL_READ_CONCURRENCY` enables a process-wide blocking-read limit for local +/// file payloads. The default, `0`, leaves local admission unlimited. Network payloads use this +/// reader's object-store concurrency. +/// +/// `VORTEX_LOCAL_FILE_HANDLE_REUSE=1` experimentally retains the first local file payload's +/// descriptor for this reader. Reads and size queries then refer to that opened inode, including +/// after the path is replaced. This removes repeated local GET/open/metadata preparation. +/// On Linux, `VORTEX_LOCAL_READ_RANDOM=1` disables kernel readahead on local file payloads; +/// with handle reuse this hint is applied once when the descriptor is first opened. pub struct ObjectStoreReadAt { + #[cfg(not(target_arch = "wasm32"))] + local_limit: Option>, + #[cfg(not(target_arch = "wasm32"))] + local_file: Option>>>>, store: Arc, path: ObjectPath, uri: Arc, @@ -65,7 +111,15 @@ impl ObjectStoreReadAt { allocator: BufferAllocatorRef, ) -> Self { let uri = Arc::from(path.to_string()); + #[cfg(not(target_arch = "wasm32"))] + static REUSE_LOCAL_FILE: LazyLock = LazyLock::new(|| { + std::env::var("VORTEX_LOCAL_FILE_HANDLE_REUSE").is_ok_and(|value| value == "1") + }); Self { + #[cfg(not(target_arch = "wasm32"))] + local_limit: local_read_limit(), + #[cfg(not(target_arch = "wasm32"))] + local_file: REUSE_LOCAL_FILE.then(Arc::default), store, path, uri, @@ -89,47 +143,221 @@ impl ObjectStoreReadAt { } } +#[cfg(not(target_arch = "wasm32"))] +enum ReadFile { + Owned(File), + Shared(Arc), +} + +#[cfg(not(target_arch = "wasm32"))] +impl AsRef for ReadFile { + fn as_ref(&self) -> &File { + match self { + Self::Owned(file) => file, + Self::Shared(file) => file, + } + } +} + +enum ReadPayload { + #[cfg(not(target_arch = "wasm32"))] + File(ReadFile), + Stream(BoxStream<'static, object_store::Result>), +} + +#[cfg(not(target_arch = "wasm32"))] +fn configure_local_file(file: File) -> File { + #[cfg(target_os = "linux")] + { + static RANDOM_ACCESS: LazyLock = LazyLock::new(|| { + std::env::var("VORTEX_LOCAL_READ_RANDOM").is_ok_and(|value| value == "1") + }); + if *RANDOM_ACCESS { + // Coalescing already selects the ranges. Shared descriptors otherwise accumulate + // readahead history across concurrently read, unrelated column segments. + if let Err(error) = rustix::fs::fadvise(&file, 0, None, rustix::fs::Advice::Random) { + tracing::debug!(%error, "local random-access hint unavailable"); + } + } + } + file +} + +async fn read_payload( + store: &dyn ObjectStore, + path: &ObjectPath, + range: Range, + #[cfg(not(target_arch = "wasm32"))] local_file: Option<&OnceCell>>>, +) -> VortexResult<(ReadPayload, bool)> { + #[cfg(not(target_arch = "wasm32"))] + let mut initial_payload = None; + #[cfg(not(target_arch = "wasm32"))] + if let Some(local_file) = local_file { + let mut opened_here = false; + let file = local_file + .get_or_try_init(|| async { + opened_here = true; + let response = store + .get_opts( + path, + GetOptions { + range: Some(GetRange::Bounded(range.clone())), + ..Default::default() + }, + ) + .await?; + match response.payload { + GetResultPayload::File(file, _) => { + Ok::<_, VortexError>(Some(Arc::new(configure_local_file(file)))) + } + stream => { + // The initializer keeps its stream; other requests do their own GETs. + initial_payload = Some(stream); + Ok(None) + } + } + }) + .await?; + if let Some(file) = file { + return Ok(( + ReadPayload::File(ReadFile::Shared(Arc::clone(file))), + !opened_here, + )); + } + } + #[cfg(target_arch = "wasm32")] + let initial_payload = None; + let payload = match initial_payload { + Some(payload) => payload, + None => { + store + .get_opts( + path, + GetOptions { + range: Some(GetRange::Bounded(range)), + ..Default::default() + }, + ) + .await? + .payload + } + }; + let payload = match payload { + #[cfg(not(target_arch = "wasm32"))] + GetResultPayload::File(file, _) => { + ReadPayload::File(ReadFile::Owned(configure_local_file(file))) + } + #[cfg(target_arch = "wasm32")] + GetResultPayload::File(..) => unreachable!("File payload not supported on wasm32"), + GetResultPayload::Stream(stream) => ReadPayload::Stream(stream), + }; + Ok((payload, false)) +} + async fn read_object_store_range( store: Arc, path: ObjectPath, io_handle: Handle, allocator: BufferAllocatorRef, request: ReadAtRequest, + #[cfg(not(target_arch = "wasm32"))] local_limit: Option>, + #[cfg(not(target_arch = "wasm32"))] local_file: Option>>>>, ) -> VortexResult { + #[cfg(not(target_arch = "wasm32"))] + let timing = tracing::enabled!(target: "vortex_io::read_timing", tracing::Level::DEBUG) + .then(Instant::now); + #[cfg(not(target_arch = "wasm32"))] + let start_ns = timing.map(|_| timestamp_ns()); let ReadAtRequest { offset, length, alignment, } = request; let range = offset..(offset + length as u64); + let (payload, file_handle_reused) = read_payload( + store.as_ref(), + &path, + range.clone(), + #[cfg(not(target_arch = "wasm32"))] + local_file.as_deref(), + ) + .await?; + #[cfg(target_arch = "wasm32")] + let _ = file_handle_reused; + #[cfg(not(target_arch = "wasm32"))] + let received = timing.map(|_| Instant::now()); + #[cfg(not(target_arch = "wasm32"))] + let local_permit = if matches!(&payload, ReadPayload::File(..)) { + match local_limit { + Some(limit) => Some(limit.acquire_owned().await.map_err(io::Error::other)?), + None => None, + } + } else { + None + }; + #[cfg(not(target_arch = "wasm32"))] + let admitted = timing.map(|_| Instant::now()); let mut buffer = allocator.with_capacity_aligned::(length, alignment); // SAFETY: each return path checks that every byte was initialized. unsafe { buffer.set_len(length) }; + #[cfg(not(target_arch = "wasm32"))] + let allocated = timing.map(|_| Instant::now()); - let response = store - .get_opts( - &path, - GetOptions { - range: Some(GetRange::Bounded(range.clone())), - ..Default::default() - }, - ) - .await?; - - let buffer = match response.payload { + let buffer = match 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") + ReadPayload::File(file) => { + let submitted = timing.map(|_| Instant::now()); + let queued_at = timing.map(|_| SystemTime::now()); + let (buffer, phases) = io_handle + .spawn_blocking(move || { + // A cancelled async read must not free the slot while pread is still running. + let _permit = local_permit; + let started = submitted.map(|_| (Instant::now(), SystemTime::now())); + read_exact_at(file.as_ref(), buffer.as_mut_slice(), range.start)?; + let finished = started.map(|_| (Instant::now(), SystemTime::now())); + Ok::<_, io::Error>((buffer, submitted.zip(started).zip(finished))) + }) + .await + .map_err(io::Error::other)?; + if let Some(( + ( + (((start, received), admitted), allocated), + ((submitted, (started, reading_at)), (finished, completed_at)), + ), + queued_at, + )) = timing + .zip(received) + .zip(admitted) + .zip(allocated) + .zip(phases) + .zip(queued_at) + { + let resumed = Instant::now(); + tracing::debug!( + target: "vortex_io::read_timing", + path = %path, + ts_ns = timestamp_ns(), + start_ns = start_ns.unwrap_or_default(), + offset, + length, + queued_unix_ns = queued_at.duration_since(UNIX_EPOCH).unwrap_or_default().as_nanos(), + started_unix_ns = reading_at.duration_since(UNIX_EPOCH).unwrap_or_default().as_nanos(), + reading_unix_ns = reading_at.duration_since(UNIX_EPOCH).unwrap_or_default().as_nanos(), + completed_unix_ns = completed_at.duration_since(UNIX_EPOCH).unwrap_or_default().as_nanos(), + file_handle_reused, + prepare_ns = u64::try_from(submitted.duration_since(start).as_nanos()).unwrap_or(u64::MAX), + allocation_ns = u64::try_from(allocated.duration_since(admitted).as_nanos()).unwrap_or(u64::MAX), + get_ns = u64::try_from(received.duration_since(start).as_nanos()).unwrap_or(u64::MAX), + admission_ns = u64::try_from(admitted.duration_since(received).as_nanos()).unwrap_or(u64::MAX), + queue_ns = u64::try_from(started.duration_since(submitted).as_nanos()).unwrap_or(u64::MAX), + read_ns = u64::try_from(finished.duration_since(started).as_nanos()).unwrap_or(u64::MAX), + resume_ns = u64::try_from(resumed.duration_since(finished).as_nanos()).unwrap_or(u64::MAX), + "local object-store read" + ); + } + buffer } - GetResultPayload::Stream(mut byte_stream) => { + ReadPayload::Stream(mut byte_stream) => { let mut written = 0usize; while let Some(bytes) = byte_stream.next().await { let bytes = bytes?; @@ -173,6 +401,19 @@ impl VortexReadAt for ObjectStoreReadAt { } fn size(&self) -> BoxFuture<'static, VortexResult> { + #[cfg(not(target_arch = "wasm32"))] + if let Some(file) = self + .local_file + .as_ref() + .and_then(|cache| cache.get()) + .and_then(Option::as_ref) + { + let file = Arc::clone(file); + return self + .handle + .spawn_blocking(move || Ok(file.metadata()?.len())) + .boxed(); + } let store = Arc::clone(&self.store); let path = self.path.clone(); async move { @@ -196,6 +437,10 @@ impl VortexReadAt for ObjectStoreReadAt { let handle = self.handle.clone(); let allocator = self.allocator.clone(); let io_handle = handle.clone(); + #[cfg(not(target_arch = "wasm32"))] + let local_limit = self.local_limit.clone(); + #[cfg(not(target_arch = "wasm32"))] + let local_file = self.local_file.clone(); handle .spawn_io(read_object_store_range( store, @@ -203,6 +448,10 @@ impl VortexReadAt for ObjectStoreReadAt { io_handle, allocator, ReadAtRequest::new(offset, length, alignment), + #[cfg(not(target_arch = "wasm32"))] + local_limit, + #[cfg(not(target_arch = "wasm32"))] + local_file, )) .boxed() } @@ -219,6 +468,10 @@ impl VortexReadAt for ObjectStoreReadAt { let concurrency = self.concurrency.max(1); let (mut send, recv) = mpsc::channel(concurrency); let io_handle = handle.clone(); + #[cfg(not(target_arch = "wasm32"))] + let local_limit = self.local_limit.clone(); + #[cfg(not(target_arch = "wasm32"))] + let local_file = self.local_file.clone(); // 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 @@ -229,9 +482,23 @@ impl VortexReadAt for ObjectStoreReadAt { let path = path.clone(); let io_handle = io_handle.clone(); let allocator = allocator.clone(); + #[cfg(not(target_arch = "wasm32"))] + let local_limit = local_limit.clone(); + #[cfg(not(target_arch = "wasm32"))] + let local_file = local_file.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, + #[cfg(not(target_arch = "wasm32"))] + local_limit, + #[cfg(not(target_arch = "wasm32"))] + local_file, + ) + .await; (request, result) } }); @@ -260,9 +527,15 @@ mod tests { use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; + #[cfg(not(target_arch = "wasm32"))] + use std::time::Duration; use object_store::PutPayload; + #[cfg(not(target_arch = "wasm32"))] + use object_store::local::LocalFileSystem; use object_store::memory::InMemory; + #[cfg(not(target_arch = "wasm32"))] + use parking_lot::Mutex; use super::*; use crate::runtime::AbortHandle; @@ -299,6 +572,80 @@ mod tests { struct TokioAbortHandle(tokio::task::AbortHandle); + /// Holds submitted blocking work until the test explicitly executes it. Aborting its caller + /// does not abort a blocking operation that has already been submitted. + #[cfg(not(target_arch = "wasm32"))] + #[derive(Default)] + struct DeferredExecutor { + executor: CountingExecutor, + blocking: Mutex>>, + } + + #[cfg(not(target_arch = "wasm32"))] + impl Executor for DeferredExecutor { + fn spawn(&self, fut: BoxFuture<'static, ()>) -> AbortHandleRef { + self.executor.spawn(fut) + } + + fn spawn_io(&self, fut: BoxFuture<'static, ()>) -> AbortHandleRef { + self.executor.spawn_io(fut) + } + + fn spawn_cpu(&self, task: Box) -> AbortHandleRef { + self.executor.spawn_cpu(task) + } + + fn spawn_blocking_io(&self, task: Box) -> AbortHandleRef { + self.blocking.lock().push(task); + TokioAbortHandle::new_handle(tokio::spawn(async {}).abort_handle()) + } + } + + #[cfg(not(target_arch = "wasm32"))] + impl DeferredExecutor { + async fn take_blocking(&self) -> VortexResult> { + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if let Some(task) = self.blocking.lock().pop() { + return task; + } + tokio::task::yield_now().await; + } + }) + .await + .map_err(|err| vortex_error::vortex_err!("{err}")) + } + } + + #[cfg(not(target_arch = "wasm32"))] + #[tokio::test] + async fn local_reads_keep_admission_until_cancelled_blocking_work_finishes() -> VortexResult<()> + { + let directory = tempfile::tempdir()?; + let store = + Arc::new(LocalFileSystem::new_with_prefix(directory.path())?) as Arc; + let path = ObjectPath::from("test.bin"); + store.put(&path, PutPayload::from_static(TEST_DATA)).await?; + let executor = Arc::new(DeferredExecutor::default()); + let runtime = Arc::clone(&executor) as Arc; + let handle = Handle::new(Arc::downgrade(&runtime)); + let limit = Arc::new(Semaphore::new(1)); + let mut reader = ObjectStoreReadAt::new(store, path, handle); + reader.local_limit = Some(Arc::clone(&limit)); + + let first = reader.read_at(0, 6, Alignment::none()); + let blocking = executor.take_blocking().await?; + drop(first); + assert_eq!(limit.available_permits(), 0); + let second = reader.read_at(7, 5, Alignment::none()); + assert!(executor.blocking.lock().is_empty()); + blocking(); + executor.take_blocking().await?(); + assert_eq!(second.await?.to_host().await.as_slice(), b"store"); + assert_eq!(limit.available_permits(), 1); + Ok(()) + } + impl TokioAbortHandle { fn new_handle(handle: tokio::task::AbortHandle) -> AbortHandleRef { Box::new(Self(handle)) @@ -331,6 +678,105 @@ mod tests { Ok(()) } + #[cfg(not(target_arch = "wasm32"))] + #[rstest::rstest] + #[case(false)] + #[case(true)] + #[tokio::test] + async fn local_handle_reuse_pins_the_opened_file(#[case] reuse: bool) -> VortexResult<()> { + let directory = tempfile::tempdir()?; + let store = + Arc::new(LocalFileSystem::new_with_prefix(directory.path())?) as Arc; + let path = ObjectPath::from("test.bin"); + store.put(&path, PutPayload::from_static(TEST_DATA)).await?; + let executor = Arc::new(CountingExecutor::default()) as Arc; + let handle = Handle::new(Arc::downgrade(&executor)); + let mut reader = ObjectStoreReadAt::new(Arc::clone(&store), path.clone(), handle); + reader.local_file = reuse.then(Arc::default); + assert!(reader.read_at(1024, 1, Alignment::none()).await.is_err()); + if let Some(cache) = &reader.local_file { + assert!(cache.get().is_none()); + } + let initial = reader + .read_ranges(Arc::from([ + ReadAtRequest::new(1024, 1, Alignment::none()), + ReadAtRequest::new(0, 6, Alignment::new(64)), + ReadAtRequest::new(7, 5, Alignment::new(64)), + ])) + .collect::>() + .await; + assert_eq!(initial.len(), 3); + for (request, result) in initial { + if request.offset == 1024 { + assert!(result.is_err()); + continue; + } + let offset = usize::try_from(request.offset)?; + assert_eq!( + result?.to_host().await.as_slice(), + &TEST_DATA[offset..offset + request.length] + ); + } + + const REPLACEMENT: &[u8] = b"new file"; + // Object-store replacement uses another inode. Reuse deliberately keeps the opened one. + store + .put(&path, PutPayload::from_static(REPLACEMENT)) + .await?; + let expected = if reuse { TEST_DATA } else { REPLACEMENT }; + assert_eq!(reader.size().await?, expected.len() as u64); + let requests = Arc::from([ + ReadAtRequest::new(0, 4, Alignment::new(64)), + ReadAtRequest::new(0, 64, Alignment::none()), + ReadAtRequest::new(4, 4, Alignment::new(64)), + ]); + let results = reader.read_ranges(requests).collect::>().await; + assert_eq!(results.len(), 3); + for (request, result) in results { + if request.length == 64 { + assert!(result.is_err()); + continue; + } + let buffer = result?.to_host().await; + assert_eq!(buffer.alignment(), Alignment::new(64)); + let offset = usize::try_from(request.offset)?; + assert_eq!( + buffer.as_slice(), + &expected[offset..offset + request.length] + ); + } + Ok(()) + } + + #[cfg(not(target_arch = "wasm32"))] + #[tokio::test] + async fn local_handle_reuse_keeps_stream_payloads_dynamic() -> VortexResult<()> { + let store = Arc::new(InMemory::new()) as Arc; + let path = ObjectPath::from("test.bin"); + store.put(&path, PutPayload::from_static(TEST_DATA)).await?; + let executor = Arc::new(CountingExecutor::default()) as Arc; + let handle = Handle::new(Arc::downgrade(&executor)); + let mut reader = ObjectStoreReadAt::new(Arc::clone(&store), path.clone(), handle); + let local_file = Arc::new(OnceCell::new()); + reader.local_file = Some(Arc::clone(&local_file)); + reader.read_at(7, 5, Alignment::none()).await?; + store + .put(&path, PutPayload::from_static(b"new file")) + .await?; + assert_eq!( + reader + .read_at(0, 8, Alignment::none()) + .await? + .to_host() + .await + .as_slice(), + b"new file" + ); + assert_eq!(reader.size().await?, 8); + assert!(local_file.get().is_some_and(Option::is_none)); + Ok(()) + } + #[tokio::test] async fn read_ranges_uses_one_io_task() -> anyhow::Result<()> { let executor = Arc::new(CountingExecutor::default()); diff --git a/vortex-io/src/request/mod.rs b/vortex-io/src/request/mod.rs index efce222ba0f..0adb14b692c 100644 --- a/vortex-io/src/request/mod.rs +++ b/vortex-io/src/request/mod.rs @@ -18,6 +18,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_panic; mod source; +pub mod trace; pub use source::ReadAtIoSource; @@ -62,12 +63,15 @@ pub enum IoIntent { Prefetch, /// Coalescing interest only; never starts an independent read. Announce, + /// Optional withdrawal of a range that this registration scope no longer needs. + /// Required fetches and other scopes' interests must remain live. Services may decline it. + Forget, } /// One request: a consumer-chosen id and the target it names. #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)] pub struct IoRequest { - /// Required delivery, background prefetch, or coalescing-only interest. + /// Required delivery, an optional hint, or withdrawal of optional interest. pub intent: IoIntent, /// Identifies the request within the issuing object. pub request: IoRequestId, @@ -122,8 +126,14 @@ pub struct Completion { /// IO registration and dispatch for one registration scope, such as the work descended from one /// root planner. Services may decline optional hints. Accepted optional work belongs to the -/// source, not to the publishing stage, and lives until the source is cleared or dropped. +/// source, not to the publishing stage, and lives until forgotten, cleared or dropped. pub trait IoSource: Send + Sync { + /// A process-unique session identity for optional scan/IO tracing. + /// Sources without correlated physical-read diagnostics may leave this unset. + fn trace_id(&self) -> Option { + None + } + /// Registers a whole batch before any of its reads become eligible for dispatch. /// Registration does not perform IO and never blocks. fn submit(&self, owner: IoOwnerId, batch: IoBatch) -> VortexResult<()>; diff --git a/vortex-io/src/request/source/tests.rs b/vortex-io/src/request/source/tests.rs index fd6bcfd12a9..ffd0b2b7e86 100644 --- a/vortex-io/src/request/source/tests.rs +++ b/vortex-io/src/request/source/tests.rs @@ -25,6 +25,7 @@ fn request(intent: IoIntent) -> IoRequest { #[rstest] #[case(IoIntent::Announce)] #[case(IoIntent::Prefetch)] +#[case(IoIntent::Forget)] fn optional_hints_can_be_declined(#[case] intent: IoIntent) -> VortexResult<()> { let source = ReadAtIoSource::new(Arc::new(PanickingReadAt), Arc::new(RUNTIME.clone())); source.submit(IoOwnerId(1), vec![request(intent)])?; diff --git a/vortex-io/src/request/trace.rs b/vortex-io/src/request/trace.rs new file mode 100644 index 00000000000..3d38e8c063c --- /dev/null +++ b/vortex-io/src/request/trace.rs @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A shared monotonic clock for scan-driver and physical-read diagnostics. + +use std::sync::LazyLock; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::time::Instant; + +/// Nanoseconds since the first diagnostic event in this process. +/// +/// Only call this when the diagnostic tracing target is enabled. All scan and IO workers share +/// the clock, so their intervals can be joined without relying on wall-clock timestamps. +pub fn timestamp_ns() -> u64 { + static START: LazyLock = LazyLock::new(Instant::now); + u64::try_from(START.elapsed().as_nanos()).unwrap_or(u64::MAX) +} + +/// A process-unique diagnostic identity, never reused when a scan or source is dropped. +pub fn next_id() -> u64 { + static NEXT_ID: AtomicU64 = AtomicU64::new(0); + NEXT_ID.fetch_add(1, Ordering::Relaxed) +} diff --git a/vortex-io/src/std_file/read_at.rs b/vortex-io/src/std_file/read_at.rs index 21aac88923f..cf3cb79473a 100644 --- a/vortex-io/src/std_file/read_at.rs +++ b/vortex-io/src/std_file/read_at.rs @@ -13,6 +13,9 @@ use std::os::unix::fs::FileExt; use std::os::windows::fs::FileExt; use std::path::Path; use std::sync::Arc; +use std::time::Instant; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; use futures::FutureExt; use futures::future::BoxFuture; @@ -123,13 +126,37 @@ impl VortexReadAt for FileReadAt { let file = Arc::clone(&self.file); let handle = self.handle.clone(); let allocator = self.allocator.clone(); + let uri = tracing::enabled!(target: "vortex_io::read_timing", tracing::Level::DEBUG) + .then(|| Arc::clone(&self.uri)); async move { + let submitted = uri.as_ref().map(|_| Instant::now()); handle .spawn_blocking(move || { + let started = submitted.map(|_| Instant::now()); let mut buffer = allocator.with_capacity_aligned::(length, alignment); // SAFETY: read_exact_at initializes every byte before the buffer is frozen. unsafe { buffer.set_len(length) }; - read_exact_at(&file, buffer.as_mut_slice(), offset)?; + let read_started = started.map(|_| Instant::now()); + let result = read_exact_at(&file, buffer.as_mut_slice(), offset); + if let Some(((submitted, started), (read_started, uri))) = + submitted.zip(started).zip(read_started.zip(uri)) + { + let finished = Instant::now(); + tracing::debug!( + target: "vortex_io::read_timing", + path = %uri, + offset, + length, + completed_unix_ns = u64::try_from( + SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_nanos() + ).unwrap_or(u64::MAX), + queue_ns = u64::try_from(started.duration_since(submitted).as_nanos()).unwrap_or(u64::MAX), + allocation_ns = u64::try_from(read_started.duration_since(started).as_nanos()).unwrap_or(u64::MAX), + read_ns = u64::try_from(finished.duration_since(read_started).as_nanos()).unwrap_or(u64::MAX), + "local positional read" + ); + } + result?; Ok(BufferHandle::new_host(buffer.freeze())) }) .await diff --git a/vortex-layout/src/segments/cache.rs b/vortex-layout/src/segments/cache.rs index f908708a801..d879c087a42 100644 --- a/vortex-layout/src/segments/cache.rs +++ b/vortex-layout/src/segments/cache.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use async_trait::async_trait; use futures::FutureExt; +use futures::future; use moka::future::Cache; use moka::future::CacheBuilder; use moka::policy::EvictionPolicy; @@ -48,6 +49,13 @@ use crate::segments::SegmentSource; /// it through `__getbuffer__`, as `ReadBuffer` does, so Python can keep it without a copy. #[async_trait] pub trait SegmentCache: Send + Sync { + /// Return immediately available cached bytes, or `None` to use the asynchronous lookup. + /// This must not perform IO. A miss still goes through [`Self::get`], which may resolve + /// while the source's announced request waits to be polled. + fn get_if_ready(&self, _id: SegmentId) -> Option { + None + } + /// Return a cached segment, or `None` on cache miss. async fn get(&self, id: SegmentId) -> VortexResult>; /// Store a segment in the cache. @@ -183,6 +191,14 @@ pub struct FileSegmentCache { #[async_trait] impl SegmentCache for FileSegmentCache { + fn get_if_ready(&self, id: SegmentId) -> Option { + self.cache + .0 + .get(&(self.source_id.clone(), id)) + .now_or_never() + .flatten() + } + async fn get(&self, id: SegmentId) -> VortexResult> { Ok(self.cache.0.get(&(self.source_id.clone(), id)).await) } @@ -229,6 +245,12 @@ impl InstrumentedSegmentCache { #[async_trait] impl SegmentCache for InstrumentedSegmentCache { + fn get_if_ready(&self, id: SegmentId) -> Option { + let segment = self.segment_cache.get_if_ready(id)?; + self.hits.add(1); + Some(segment) + } + async fn get(&self, id: SegmentId) -> VortexResult> { let result = self.segment_cache.get(id).await?; if result.is_some() { @@ -261,6 +283,10 @@ impl SegmentCacheSourceAdapter { impl SegmentSource for SegmentCacheSourceAdapter { fn request(&self, id: SegmentId) -> SegmentFuture { + if let Some(segment) = self.cache.get_if_ready(id) { + tracing::debug!("Resolved segment {} from cache", id); + return future::ready(Ok(BufferHandle::new_host(segment))).boxed(); + } let cache = Arc::clone(&self.cache); let delegate = self.source.request(id); @@ -284,6 +310,9 @@ impl SegmentSource for SegmentCacheSourceAdapter { #[cfg(test)] mod tests { + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use rstest::rstest; use vortex_buffer::ByteBuffer; @@ -302,9 +331,15 @@ mod tests { a.put(id, ByteBuffer::copy_from(b"from a")).await?; assert_eq!(a.get(id).await?.as_deref(), Some(b"from a".as_slice())); assert!(b.get(id).await?.is_none()); + assert_eq!(a.get_if_ready(id).as_deref(), Some(b"from a".as_slice())); + assert!(b.get_if_ready(id).is_none()); // A later view with the same source ID sees what the earlier one stored. let a_again = shared.for_file("a"); + assert_eq!( + a_again.get_if_ready(id).as_deref(), + Some(b"from a".as_slice()) + ); assert_eq!( a_again.get(id).await?.as_deref(), Some(b"from a".as_slice()) @@ -336,4 +371,75 @@ mod tests { } Ok(()) } + + struct TestCache { + bytes: Option, + immediate: bool, + } + + #[async_trait] + impl SegmentCache for TestCache { + fn get_if_ready(&self, _id: SegmentId) -> Option { + self.immediate.then(|| self.bytes.clone()).flatten() + } + + async fn get(&self, _id: SegmentId) -> VortexResult> { + Ok(self.bytes.clone()) + } + + async fn put(&self, _id: SegmentId, _buffer: ByteBuffer) -> VortexResult<()> { + Ok(()) + } + } + + struct RecordingSource { + bytes: ByteBuffer, + requests: AtomicUsize, + polls: Arc, + } + + impl SegmentSource for RecordingSource { + fn request(&self, _id: SegmentId) -> SegmentFuture { + self.requests.fetch_add(1, Ordering::Relaxed); + let polls = Arc::clone(&self.polls); + let bytes = self.bytes.clone(); + async move { + polls.fetch_add(1, Ordering::Relaxed); + Ok(BufferHandle::new_host(bytes)) + } + .boxed() + } + } + + #[rstest] + #[case::immediate_hit(true, true, 0, 0)] + #[case::asynchronous_hit(true, false, 1, 0)] + #[case::miss(false, true, 1, 1)] + #[tokio::test] + async fn cache_hits_skip_reads_and_misses_keep_early_announcements( + #[case] cached: bool, + #[case] immediate: bool, + #[case] requests: usize, + #[case] polls: usize, + ) -> VortexResult<()> { + let bytes = ByteBuffer::from(vec![1, 2, 3, 4]); + let source = Arc::new(RecordingSource { + bytes: bytes.clone(), + requests: AtomicUsize::new(0), + polls: Arc::default(), + }); + let adapter = SegmentCacheSourceAdapter::new( + Arc::new(TestCache { + bytes: cached.then(|| bytes.clone()), + immediate, + }), + Arc::::clone(&source), + ); + let request = adapter.request(SegmentId::from(0)); + assert_eq!(source.requests.load(Ordering::Relaxed), requests); + assert_eq!(source.polls.load(Ordering::Relaxed), 0); + assert_eq!(request.await?.try_into_host()?.await?, bytes); + assert_eq!(source.polls.load(Ordering::Relaxed), polls); + Ok(()) + } }