diff --git a/vortex-io/src/lib.rs b/vortex-io/src/lib.rs index fabeb220934..f7b3cde6161 100644 --- a/vortex-io/src/lib.rs +++ b/vortex-io/src/lib.rs @@ -25,6 +25,7 @@ pub mod object_store; mod read_at; pub mod request; pub mod runtime; +pub mod segments; pub mod session; #[cfg(not(target_arch = "wasm32"))] pub mod std_file; diff --git a/vortex-io/src/segments.rs b/vortex-io/src/segments.rs new file mode 100644 index 00000000000..390419528cb --- /dev/null +++ b/vortex-io/src/segments.rs @@ -0,0 +1,382 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Reads for scans that drive their own IO. +//! +//! A request names what to read by [`IoKey`], a registered segment source and a +//! [`SourceRequest`] of it (one of its segments, or its size), and says how much it needs it with +//! an [`IoIntent`]. One service serves many sources, so one session can read from all of them. +//! A [`SegmentIoSession`] registers a whole batch of them before any can be read, and returns one +//! [`IoHandle`] per request. The handle is the request: its result arrives in it. The session holds fetches and prefetches until they finish; an announcement is +//! withdrawn by dropping every handle of it. +//! +//! A [`SegmentIoService`] serves every registered source. Its sessions share and coalesce reads, +//! so one session's announcement can widen another session's read. +//! +//! This protocol is to replace the byte-range protocol in [`request`](crate::request), which the +//! planner-driven scan still uses. + +use std::any::Any; +use std::pin::Pin; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; + +use futures::FutureExt; +use futures::StreamExt; +use futures::future::BoxFuture; +use futures::future::Shared; +use futures::stream::FuturesUnordered; +use vortex_array::buffer::BufferHandle; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; + +/// A segment source registered with a [`SegmentIoService`], such as one opened file. +#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct SegmentSourceId(pub u32); + +/// What to read from a segment source. +#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum SourceRequest { + /// The source's total length in bytes, answered with [`SourceResult::Size`]. + Size, + /// One segment, by its id within the source (the value of a + /// `vortex_layout::segments::SegmentId`), answered with [`SourceResult::Bytes`]. + Segment(u32), +} + +/// What one request reads: a source, and what to read from it. +#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct IoKey { + /// The source to read. + pub source: SegmentSourceId, + /// What to read from it. + pub request: SourceRequest, +} + +impl IoKey { + /// The key of `request` of `source`. + pub fn new(source: SegmentSourceId, request: SourceRequest) -> Self { + Self { source, request } + } + + /// The key of `segment` of `source`. + pub fn segment(source: SegmentSourceId, segment: u32) -> Self { + Self::new(source, SourceRequest::Segment(segment)) + } + + /// The key of the size of `source`. + pub fn size(source: SegmentSourceId) -> Self { + Self::new(source, SourceRequest::Size) + } +} + +/// The answer to one request. +#[derive(Clone, Debug)] +pub enum SourceResult { + /// The source's total length in bytes. + Size(u64), + /// A segment's bytes. + Bytes(BufferHandle), +} + +impl SourceResult { + /// The bytes of a segment read, failing for any other answer. + pub fn into_bytes(self) -> VortexResult { + match self { + Self::Bytes(bytes) => Ok(bytes), + Self::Size(_) => vortex_bail!("expected segment bytes, got a size"), + } + } + + /// The size a size request was answered with, failing for any other answer. + pub fn size(&self) -> VortexResult { + match self { + Self::Size(size) => Ok(*size), + Self::Bytes(_) => vortex_bail!("expected a size, got segment bytes"), + } + } +} + +/// How much a request needs what it reads. +#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)] +pub enum IoIntent { + /// Needed now. Its owner is reported by [`SegmentIoSession::poll_ready`] when it finishes. + /// A service may read fetches ahead of prefetches. + Fetch, + /// Needed soon: read, but its owner is not reported when it finishes. + Prefetch, + /// May be needed: only widens nearby reads, and is never read on its own. Dropping its + /// handles withdraws it. + Announce, +} + +/// One read a planner or morsel asks for. +#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)] +pub struct IoRequest { + /// What to read. + pub key: IoKey, + /// How much it is needed. + pub intent: IoIntent, +} + +impl IoRequest { + /// A request for `key` with `intent`. + pub fn new(key: IoKey, intent: IoIntent) -> Self { + Self { key, intent } + } +} + +/// Identifies the object that submitted a request, chosen by the caller. +#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct IoOwnerId(pub u64); + +/// The result of one read, shared by every handle of it. +pub type SharedResult = Shared>>>; + +/// One registered read. +/// +/// Clones share the read. The read stays registered while any handle of it, from any session, is +/// alive. A session also holds its fetches and prefetches until they finish, so dropping every +/// handle withdraws a read only when it was announced alone. +#[derive(Clone)] +pub struct IoHandle { + result: SharedResult, + /// Keeps the read registered with its service while the handle lives. + _registration: Option>, +} + +impl IoHandle { + /// A handle whose result arrives through `result`, keeping `registration` alive while it + /// lives. + pub fn new(result: SharedResult, registration: Arc) -> Self { + Self { + result, + _registration: Some(registration), + } + } + + /// A handle that already holds its result, such as a cache hit or a source's size. + pub fn ready(result: VortexResult) -> Self { + let result = futures::future::ready(result.map_err(Arc::new)) + .boxed() + .shared(); + // `Shared::peek` only sees a result once a poll has completed the future, so poll it once + // here for `try_get` to see it without a waker. + drop(result.clone().now_or_never()); + Self { + result, + _registration: None, + } + } + + /// The read's result once it has arrived. Never blocks and never reads. + pub fn try_get(&self) -> Option> { + self.result + .peek() + .map(|result| result.clone().map_err(VortexError::from)) + } + + /// The shared future the result arrives through, for sessions that wake on it. + pub fn result(&self) -> SharedResult { + self.result.clone() + } +} + +/// One scope's reads, such as the work descended from one root planner. +/// +/// Registration does no IO and never blocks. Every read is driven by the service, so nothing has +/// to be polled for reads to progress; polling only learns which reads finished. +pub trait SegmentIoSession: Send + Sync { + /// Registers the whole batch before any of it can be read, and returns one handle per + /// request, in order. Every fetch is reported by [`poll_ready`](Self::poll_ready) once it + /// finishes, including one already finished at submit, such as a cache hit or an error. + fn submit(&self, owner: IoOwnerId, batch: &[IoRequest]) -> VortexResult>; + + /// Moves the owners of fetches that finished since the last call into `ready`, once per + /// fetch, and drives the session's prefetches. Returns `Pending`, registering `cx` to be + /// woken by the next one, if none have, and fails if no fetch is in flight, since nothing + /// would wake the caller. An owner is reported even if it dropped the fetch's handles. + fn poll_ready( + &self, + cx: &mut Context<'_>, + ready: &mut Vec, + ) -> Poll>; +} + +/// Serves the reads of every registered segment source. +pub trait SegmentIoService: Send + Sync { + /// A new session for one scope. + fn session(&self) -> Arc; +} + +/// A source registered with a [`SegmentIoService`]: the service, the source's id in it, and the +/// registration. The service can read the source while any clone of this is alive. +#[derive(Clone)] +pub struct RegisteredSource { + service: Arc, + id: SegmentSourceId, + _registration: Arc, +} + +impl RegisteredSource { + /// The registration of `id` with `service`, kept alive by `registration`. + pub fn new( + service: Arc, + id: SegmentSourceId, + registration: Arc, + ) -> Self { + Self { + service, + id, + _registration: registration, + } + } + + /// The service that reads the source. + pub fn service(&self) -> &Arc { + &self.service + } + + /// The source's id in its service. + pub fn id(&self) -> SegmentSourceId { + self.id + } +} + +/// The fetches and prefetches of one session, waking the session when a fetch finishes. +/// +/// Sessions keep one of these to implement [`SegmentIoSession::poll_ready`]. Prefetches are +/// polled until they finish, so their reads start and their results land in their handles even +/// though nobody awaits them. +#[derive(Default)] +pub struct FetchTracker { + fetches: FuturesUnordered + Send>>>, + prefetches: FuturesUnordered, +} + +impl FetchTracker { + /// Tracks a prefetch until it finishes, without reporting it. + pub fn track_prefetch(&mut self, handle: &IoHandle) { + self.prefetches.push(handle.result()); + } + + /// Tracks a fetch submitted by `owner`. + pub fn track(&mut self, owner: IoOwnerId, handle: &IoHandle) { + let bytes = handle.result(); + self.fetches.push(Box::pin(async move { + // Errors are taken from the handle by its owner. + drop(bytes.await); + owner + })); + } + + /// Moves the owners of finished fetches into `ready`, as [`SegmentIoSession::poll_ready`]. + pub fn poll_ready( + &mut self, + cx: &mut Context<'_>, + ready: &mut Vec, + ) -> Poll> { + // Finished prefetches are dropped; their handles hold the results. + while let Poll::Ready(Some(_)) = self.prefetches.poll_next_unpin(cx) {} + let before = ready.len(); + loop { + match self.fetches.poll_next_unpin(cx) { + Poll::Ready(Some(owner)) => ready.push(owner), + Poll::Ready(None) if ready.len() > before => return Poll::Ready(Ok(())), + Poll::Ready(None) => { + return Poll::Ready(Err(vortex_err!( + "waiting for segment reads with no fetch in flight" + ))); + } + Poll::Pending if ready.len() > before => return Poll::Ready(Ok(())), + Poll::Pending => return Poll::Pending, + } + } + } + + /// Drops every tracked fetch and prefetch. + pub fn clear(&mut self) { + self.fetches.clear(); + self.prefetches.clear(); + } +} + +#[cfg(test)] +mod tests { + use std::task::Waker; + + use futures::channel::oneshot; + use vortex_buffer::ByteBuffer; + + use super::*; + + fn pending() -> (oneshot::Sender, IoHandle) { + let (send, recv) = oneshot::channel::(); + let result = recv + .map(|result| { + result + .map(SourceResult::Bytes) + .map_err(|_| Arc::new(vortex_err!("cancelled"))) + }) + .boxed() + .shared(); + (send, IoHandle::new(result, Arc::new(()))) + } + + #[test] + fn ready_handle_needs_no_poll() -> VortexResult<()> { + let handle = IoHandle::ready(Ok(SourceResult::Bytes(BufferHandle::new_host( + ByteBuffer::copy_from(b"abc"), + )))); + let result = handle.try_get().ok_or_else(|| vortex_err!("not ready"))??; + assert_eq!(result.into_bytes()?.len(), 3); + Ok(()) + } + + #[test] + fn results_are_checked_against_what_was_asked() -> VortexResult<()> { + assert_eq!(SourceResult::Size(7).size()?, 7); + assert!(SourceResult::Size(7).into_bytes().is_err()); + let bytes = SourceResult::Bytes(BufferHandle::new_host(ByteBuffer::copy_from(b"a"))); + assert!(bytes.size().is_err()); + Ok(()) + } + + #[test] + fn tracker_reports_owners_as_fetches_finish() -> VortexResult<()> { + let mut tracker = FetchTracker::default(); + let (send_a, a) = pending(); + let (_send_b, b) = pending(); + tracker.track(IoOwnerId(1), &a); + tracker.track(IoOwnerId(2), &b); + let mut cx = Context::from_waker(Waker::noop()); + let mut ready = Vec::new(); + assert!(tracker.poll_ready(&mut cx, &mut ready).is_pending()); + assert!(a.try_get().is_none()); + + send_a + .send(BufferHandle::new_host(ByteBuffer::copy_from(b"a"))) + .map_err(|_| vortex_err!("receiver dropped"))?; + assert!(matches!( + tracker.poll_ready(&mut cx, &mut ready), + Poll::Ready(Ok(())) + )); + assert_eq!(ready, vec![IoOwnerId(1)]); + assert!(a.try_get().is_some()); + assert!(b.try_get().is_none()); + Ok(()) + } + + #[test] + fn tracker_fails_with_nothing_in_flight() { + let mut tracker = FetchTracker::default(); + let mut cx = Context::from_waker(Waker::noop()); + assert!(matches!( + tracker.poll_ready(&mut cx, &mut Vec::new()), + Poll::Ready(Err(_)) + )); + } +}