agave_fs/
io_setup.rs

1use std::io;
2#[cfg(target_os = "linux")]
3use {
4    crate::io_uring::sqpoll::SharedSqPoll,
5    std::os::fd::{AsFd as _, BorrowedFd},
6};
7
8/// State used by IO utilities for managing shared resources and configuration during setup.
9///
10/// This may include io_uring file descriptors, flag whether to register memory buffers in kernel,
11/// and other resources that need to be accessed by multiple functions performing IO operations
12/// such that they can efficiently cooperate with each other.
13///
14/// This is achieved by creating `IoSetupState` at the beginning of the processing setup and
15/// passing its reference to IO utilities that need sharing / customized options.
16///
17/// Note: the state needs to live only during creation of the IO utilities, not during their usage,
18/// so it's generally advisable to drop it after setup is done such that e.g. init-only squeue is
19/// released.
20#[derive(Default)]
21pub struct IoSetupState {
22    #[cfg(target_os = "linux")]
23    shared_sqpoll: Option<SharedSqPoll>,
24    pub use_direct_io: bool,
25    pub use_registered_io_uring_buffers: bool,
26}
27
28impl IoSetupState {
29    /// Enables shared io-uring worker pool and sqpoll based kernel thread.
30    ///
31    /// The sqpoll thread will drain submission queues from all io-uring instances created
32    /// through builder obtained from `create_io_uring_builder()` with FD obtained from
33    /// `shared_sqpoll_fd()` after this call.
34    pub fn with_shared_sqpoll(self) -> io::Result<Self> {
35        Ok(Self {
36            #[cfg(target_os = "linux")]
37            shared_sqpoll: Some(SharedSqPoll::new()?),
38            ..self
39        })
40    }
41
42    /// Enables registering of buffers in io-uring (as `fixed`).
43    ///
44    /// Speeds up kernel operations on the memory, but requires appropriate memlock ulimit.
45    pub fn with_buffers_registered(mut self, fixed: bool) -> Self {
46        self.use_registered_io_uring_buffers = fixed;
47        self
48    }
49
50    /// Enables direct I/O for operations that bypass the operating system's caching layer.
51    ///
52    /// File system is required to support opening files with `O_DIRECT` flag.
53    ///
54    /// This can improve performance when allocation and checking of caches by the kernel is slower
55    /// than the overall savings from re-using cached file data (e.g. for read / write once data).
56    pub fn with_direct_io(mut self, use_direct_io: bool) -> Self {
57        self.use_direct_io = use_direct_io;
58        self
59    }
60
61    #[cfg(target_os = "linux")]
62    pub fn shared_sqpoll_fd(&self) -> Option<BorrowedFd<'_>> {
63        self.shared_sqpoll.as_ref().map(|s| s.as_fd())
64    }
65}
66
67#[cfg(all(test, target_os = "linux"))]
68mod tests {
69    use {
70        super::*,
71        crate::{
72            file_io::FileCreator,
73            io_uring::{
74                file_creator::IoUringFileCreatorBuilder,
75                sequential_file_reader::SequentialFileReaderBuilder,
76            },
77        },
78        rand::RngCore,
79        std::{
80            fs::File,
81            io::{Cursor, Read},
82            sync::{Arc, RwLock},
83        },
84    };
85
86    #[test]
87    fn test_shared_sqpoll_read_and_create() {
88        let io_setup = &IoSetupState::default().with_shared_sqpoll().unwrap();
89
90        let read_bytes = RwLock::new(vec![]);
91        let read_bytes_ref = &read_bytes;
92        let mut file_creator = IoUringFileCreatorBuilder::new()
93            .shared_sqpoll(io_setup.shared_sqpoll_fd())
94            .build(1 << 20, move |file_info| {
95                let mut reader = SequentialFileReaderBuilder::new()
96                    .shared_sqpoll(io_setup.shared_sqpoll_fd())
97                    .build(1 << 20)
98                    .unwrap();
99                reader.set_path(file_info.path).unwrap();
100                reader
101                    .read_to_end(read_bytes_ref.write().unwrap().as_mut())
102                    .unwrap();
103                None
104            })
105            .unwrap();
106
107        let temp_dir = tempfile::tempdir().unwrap();
108        let dir_handle = Arc::new(File::open(temp_dir.path()).unwrap());
109        let mut write_bytes = vec![0; 2 << 20];
110        rand::rng().fill_bytes(&mut write_bytes);
111
112        let file_path1 = temp_dir.path().join("test-1.txt");
113        let file_path2 = temp_dir.path().join("test-2.txt");
114        for path in [file_path1, file_path2] {
115            let dir_handle = dir_handle.clone();
116            file_creator
117                .schedule_create_at_dir(path, 0o644, dir_handle, &mut Cursor::new(&write_bytes))
118                .unwrap();
119        }
120        file_creator.drain().unwrap();
121        drop(file_creator);
122
123        // After drain all the callbacks that read data into `read_bytes` should be done.
124        let read_bytes = read_bytes.into_inner().unwrap();
125        // Expect atomically appended two copies, each from different file.
126        assert_eq!(&read_bytes[..write_bytes.len()], &write_bytes);
127        assert_eq!(&read_bytes[write_bytes.len()..], &write_bytes);
128    }
129}
agave_fs/io_setup.rs

agave_fs/
io_setup.rs
