vortex_io/request/
source.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! A blocking byte source for exercising the protocol with real footer reads.
5
6use std::collections::VecDeque;
7use std::sync::Arc;
8use std::task::Context;
9use std::task::Poll;
10
11use parking_lot::Mutex;
12use vortex_error::VortexResult;
13use vortex_error::vortex_bail;
14use vortex_error::vortex_err;
15
16use super::Completion;
17use super::IoBatch;
18use super::IoIntent;
19use super::IoOwnerId;
20use super::IoRequest;
21use super::IoResult;
22use super::IoSource;
23use super::IoTarget;
24use crate::VortexReadAt;
25use crate::runtime::BlockingRuntime;
26
27/// A simple source that registers batches and performs one fetch per `wait()`.
28/// Optional hints are declined. Shared reads, prefetching, and priority scheduling
29/// belong to a later IO service implementation.
30pub struct ReadAtIoSource<R: BlockingRuntime> {
31    read: Arc<dyn VortexReadAt>,
32    runtime: Arc<R>,
33    queued: Mutex<VecDeque<(IoOwnerId, IoRequest)>>,
34}
35
36impl<R: BlockingRuntime> ReadAtIoSource<R> {
37    /// Creates a source that drives its reads on `runtime`.
38    pub fn new(read: Arc<dyn VortexReadAt>, runtime: Arc<R>) -> Self {
39        Self {
40            read,
41            runtime,
42            queued: Mutex::default(),
43        }
44    }
45}
46
47impl<R: BlockingRuntime + Send + Sync> IoSource for ReadAtIoSource<R> {
48    fn submit(&self, work: IoOwnerId, batch: IoBatch) -> VortexResult<()> {
49        for request in &batch {
50            if let IoTarget::Range { offset, len, .. } = &request.target
51                && offset.checked_add(*len as u64).is_none()
52            {
53                vortex_bail!("ReadAtIoSource: range overflow for {:?}", request.request);
54            }
55        }
56        self.queued.lock().extend(
57            batch
58                .into_iter()
59                .filter(|request| request.intent == IoIntent::Fetch)
60                .map(|request| (work, request)),
61        );
62        Ok(())
63    }
64
65    fn poll(&self) -> VortexResult<Option<Completion>> {
66        Ok(None)
67    }
68
69    /// Performs the next queued read on the polling thread: this source has no background IO.
70    fn poll_completion(&self, _cx: &mut Context<'_>) -> Poll<VortexResult<Completion>> {
71        Poll::Ready(self.wait())
72    }
73
74    fn wait(&self) -> VortexResult<Completion> {
75        let (work, request) = self
76            .queued
77            .lock()
78            .pop_front()
79            .ok_or_else(|| vortex_err!("ReadAtIoSource: wait with nothing outstanding"))?;
80        let result = match request.target {
81            IoTarget::Size => self.runtime.block_on(self.read.size()).map(IoResult::Size),
82            IoTarget::Range {
83                offset,
84                len,
85                alignment,
86            } => self
87                .runtime
88                .block_on(self.read.read_at(offset, len, alignment))
89                .map(IoResult::Bytes),
90        };
91        Ok(Completion {
92            owner: work,
93            request: request.request,
94            result,
95        })
96    }
97
98    fn release(&self, work: IoOwnerId) {
99        self.queued.lock().retain(|(owner, _)| *owner != work);
100    }
101
102    fn clear(&self) {
103        self.queued.lock().clear();
104    }
105}
106
107#[cfg(test)]
108mod tests;
vortex_io/request/source.rs

vortex_io/request/
source.rs
