lance_io/uring/
reader.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! UringReader implementation.
5
6use super::future::UringReadFuture;
7use super::requests::IoRequest;
8use super::thread::{QueuedRequest, THREAD_SELECTOR, URING_THREADS};
9use super::{DEFAULT_URING_BLOCK_SIZE, DEFAULT_URING_IO_PARALLELISM, URING_BLOCK_SIZE};
10use crate::local::to_local_path;
11use crate::traits::Reader;
12use crate::uring::requests::RequestState;
13use crate::utils::tracking_store::IOTracker;
14use bytes::{Bytes, BytesMut};
15use futures::FutureExt;
16use futures::future::BoxFuture;
17use lance_core::deepsize::DeepSizeOf;
18use lance_core::{Error, Result};
19use object_store::path::Path;
20use std::fs::File;
21use std::future::Future;
22use std::io::{self, ErrorKind};
23use std::ops::Range;
24use std::os::unix::io::{AsRawFd, RawFd};
25use std::pin::Pin;
26use std::sync::atomic::Ordering;
27use std::sync::{Arc, LazyLock, Mutex};
28use std::time::Duration;
29use tracing::instrument;
30
31/// Cache key for UringReader instances.
32/// We cache by (path, block_size) because block_size affects reader behavior.
33#[derive(Clone, Debug, Hash, Eq, PartialEq)]
34pub(super) struct CacheKey {
35    path: String,
36    block_size: usize,
37}
38
39impl CacheKey {
40    pub(super) fn new(path: &Path, block_size: usize) -> Self {
41        Self {
42            path: path.to_string(),
43            block_size,
44        }
45    }
46}
47
48/// Data stored in the cache for each opened file.
49#[derive(Clone)]
50pub(super) struct CachedReaderData {
51    pub(super) handle: Arc<UringFileHandle>,
52    pub(super) size: usize,
53}
54
55/// Global cache of open file handles.
56/// Entries expire after 60 seconds to ensure files are eventually closed.
57pub(super) static HANDLE_CACHE: LazyLock<moka::future::Cache<CacheKey, CachedReaderData>> =
58    LazyLock::new(|| {
59        moka::future::Cache::builder()
60            .time_to_live(Duration::from_secs(60))
61            .max_capacity(10_000)
62            .build()
63    });
64
65/// File handle for io_uring operations.
66///
67/// Keeps the file alive and provides the raw file descriptor.
68#[derive(Debug)]
69pub(super) struct UringFileHandle {
70    /// The file (kept alive via Arc)
71    #[allow(unused)]
72    file: Arc<File>,
73
74    /// Raw file descriptor for io_uring
75    pub(super) fd: RawFd,
76
77    /// Object store path
78    pub(super) path: Path,
79}
80
81impl UringFileHandle {
82    pub(super) fn new(file: File, path: Path) -> Self {
83        let fd = file.as_raw_fd();
84        Self {
85            file: Arc::new(file),
86            fd,
87            path,
88        }
89    }
90}
91
92/// io_uring-based reader for local files.
93///
94/// This reader uses a dedicated process-wide thread running an io_uring event loop
95/// for high-performance asynchronous I/O.
96#[derive(Debug)]
97pub struct UringReader {
98    /// File handle
99    handle: Arc<UringFileHandle>,
100
101    /// Block size for I/O operations
102    block_size: usize,
103
104    /// File size (determined at open time)
105    size: usize,
106
107    /// I/O tracker for monitoring operations
108    io_tracker: Arc<IOTracker>,
109}
110
111impl DeepSizeOf for UringReader {
112    fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize {
113        // Skip file handle (just a system resource)
114        // Only count the path's deep size
115        self.handle.path.as_ref().deep_size_of_children(context)
116    }
117}
118
119impl UringReader {
120    /// Open a file with io_uring.
121    ///
122    /// This is the internal constructor used by ObjectStore.
123    #[instrument(level = "debug")]
124    pub(crate) async fn open(
125        path: &Path,
126        block_size: usize,
127        known_size: Option<usize>,
128        io_tracker: Arc<IOTracker>,
129    ) -> Result<Box<dyn Reader>> {
130        // Determine block size with environment variable override
131        let block_size = URING_BLOCK_SIZE.unwrap_or(block_size.max(DEFAULT_URING_BLOCK_SIZE));
132
133        let cache_key = CacheKey::new(path, block_size);
134
135        // Try to get from cache first
136        if let Some(data) = HANDLE_CACHE.get(&cache_key).await {
137            // Use known_size if provided, otherwise use cached size
138            let size = known_size.unwrap_or(data.size);
139            return Ok(Box::new(Self {
140                handle: data.handle,
141                block_size,
142                size,
143                io_tracker,
144            }) as Box<dyn Reader>);
145        }
146
147        // Cache miss - open file and get size
148        let path_clone = path.clone();
149        let local_path = to_local_path(path);
150
151        let data = tokio::task::spawn_blocking(move || {
152            let file = File::open(&local_path).map_err(|e| match e.kind() {
153                ErrorKind::NotFound => Error::not_found(path_clone.to_string()),
154                _ => e.into(),
155            })?;
156
157            // Get size from known_size or file metadata
158            let size = match known_size {
159                Some(s) => s,
160                None => file.metadata()?.len() as usize,
161            };
162
163            Ok::<_, Error>(CachedReaderData {
164                handle: Arc::new(UringFileHandle::new(file, path_clone)),
165                size,
166            })
167        })
168        .await??;
169
170        // Insert into cache
171        HANDLE_CACHE.insert(cache_key, data.clone()).await;
172
173        // Return new reader instance
174        Ok(Box::new(Self {
175            handle: data.handle.clone(),
176            block_size,
177            size: data.size,
178            io_tracker,
179        }) as Box<dyn Reader>)
180    }
181
182    /// Submit a read request to the io_uring thread via channel and return a future.
183    fn submit_read(
184        &self,
185        offset: u64,
186        length: usize,
187    ) -> Pin<Box<dyn Future<Output = object_store::Result<Bytes>> + Send>> {
188        let mut buffer = BytesMut::with_capacity(length);
189        unsafe {
190            buffer.set_len(length);
191        }
192
193        // Create IoRequest with all data
194        let request = Arc::new(IoRequest {
195            fd: self.handle.fd,
196            offset,
197            length,
198            thread_id: std::thread::current().id(),
199            state: Mutex::new(RequestState {
200                completed: false,
201                waker: None,
202                err: None,
203                buffer,
204                bytes_read: 0,
205            }),
206        });
207
208        if URING_THREADS.threads.is_empty() {
209            let initialization_errors = if URING_THREADS.initialization_errors.is_empty() {
210                "LANCE_URING_THREAD_COUNT is 0".to_owned()
211            } else {
212                URING_THREADS.initialization_errors.join("; ")
213            };
214            return Box::pin(async move {
215                Err(object_store::Error::Generic {
216                    store: "UringReader",
217                    source: Box::new(io::Error::other(format!(
218                        "no io_uring worker threads are available: {initialization_errors}"
219                    ))),
220                })
221            });
222        }
223
224        // Select thread in round-robin fashion
225        let thread_idx = (THREAD_SELECTOR.fetch_add(1, Ordering::Relaxed) as usize)
226            % URING_THREADS.threads.len();
227        let thread = &URING_THREADS.threads[thread_idx];
228
229        if !thread.is_alive.load(Ordering::Acquire) {
230            return Box::pin(async move {
231                Err(object_store::Error::Generic {
232                    store: "UringReader",
233                    source: Box::new(io::Error::new(
234                        io::ErrorKind::BrokenPipe,
235                        "io_uring thread died",
236                    )),
237                })
238            });
239        }
240
241        // Send to selected thread via channel
242        match thread
243            .request_tx
244            .send(QueuedRequest::new(Arc::clone(&request)))
245        {
246            Ok(()) => Box::pin(UringReadFuture { request }),
247            Err(_) => Box::pin(async move {
248                Err(object_store::Error::Generic {
249                    store: "UringReader",
250                    source: Box::new(io::Error::new(
251                        io::ErrorKind::BrokenPipe,
252                        "io_uring thread died",
253                    )),
254                })
255            }),
256        }
257    }
258}
259
260impl Reader for UringReader {
261    fn path(&self) -> &Path {
262        &self.handle.path
263    }
264
265    fn block_size(&self) -> usize {
266        self.block_size
267    }
268
269    fn io_parallelism(&self) -> usize {
270        std::env::var("LANCE_URING_IO_PARALLELISM")
271            .ok()
272            .and_then(|s| s.parse().ok())
273            .unwrap_or(DEFAULT_URING_IO_PARALLELISM)
274    }
275
276    /// Returns the file size.
277    fn size(&self) -> BoxFuture<'_, object_store::Result<usize>> {
278        Box::pin(async move { Ok(self.size) })
279    }
280
281    /// Read a range of bytes using io_uring.
282    #[instrument(level = "debug", skip(self))]
283    fn get_range(&self, range: Range<usize>) -> BoxFuture<'static, object_store::Result<Bytes>> {
284        let io_tracker = self.io_tracker.clone();
285        let path = self.handle.path.clone();
286        let num_bytes = range.len() as u64;
287        let range_u64 = (range.start as u64)..(range.end as u64);
288
289        let metrics = self.io_tracker.begin_io("get");
290        self.submit_read(range.start as u64, range.len())
291            .map(move |result| {
292                metrics.record(&result, num_bytes);
293                if result.is_ok() {
294                    io_tracker.record_read("get_range", path, num_bytes, Some(range_u64));
295                }
296                result
297            })
298            .boxed()
299    }
300
301    /// Read the entire file using io_uring.
302    #[instrument(level = "debug", skip(self))]
303    fn get_all(&self) -> BoxFuture<'static, object_store::Result<Bytes>> {
304        let size = self.size;
305        let io_tracker = self.io_tracker.clone();
306        let path = self.handle.path.clone();
307
308        let metrics = self.io_tracker.begin_io("get");
309        self.submit_read(0, size)
310            .map(move |result| {
311                let num_bytes = result.as_ref().map_or(0, |bytes| bytes.len() as u64);
312                metrics.record(&result, num_bytes);
313                if result.is_ok() {
314                    io_tracker.record_read("get_all", path, num_bytes, None);
315                }
316                result
317            })
318            .boxed()
319    }
320}
lance_io/uring/reader.rs

lance_io/uring/
reader.rs
