io_engine/
tasks.rs

1use std::fmt;
2use std::os::fd::RawFd;
3
4use crossfire::waitgroup::WaitGroupGuard;
5use embed_collections::SegList;
6use io_buffer::{Buffer, safe_copy};
7use rustix::io::Errno;
8
9#[derive(Copy, Clone, PartialEq, Debug)]
10#[repr(u8)]
11pub enum IOAction {
12    Read = 0,  // The same with IOCB_CMD_PREAD
13    Write = 1, // the same with IOCB_CMD_PWRITE
14    Alloc = 2,
15    Fsync = 3,
16}
17
18impl IOAction {
19    #[inline(always)]
20    pub fn is_read_write(&self) -> bool {
21        (*self as u8) < (IOAction::Alloc as u8)
22    }
23}
24
25pub trait CbArgs: Sized + 'static + Send + Unpin {}
26
27impl CbArgs for () {}
28
29impl<T: Send + Sync + 'static> CbArgs for WaitGroupGuard<T> {}
30
31// Carries the information of read/write event
32pub struct IOEvent<C: CbArgs> {
33    pub action: IOAction,
34    /// Result of the IO operation.
35    /// Initialized to i32::MIN.
36    /// - `>= 0`: Accumulated bytes transferred (used for partial IO retries).
37    /// - `<0`: Error code (negative errno).
38    pub(crate) res: i32,
39    /// make sure SListNode always in the front.
40    /// This is for putting sub_tasks in the link list, without additional allocation.
41    pub(crate) buf_or_len: BufOrLen,
42    pub offset: i64,
43    pub fd: RawFd,
44    pub(crate) args: Option<TaskArgs<C>>,
45}
46
47pub(crate) enum TaskArgs<C: CbArgs> {
48    Callback(C),
49    Merged(SegList<IOEventMerged<C>>),
50}
51
52pub(crate) enum BufOrLen {
53    Buffer(Buffer),
54    /// for fallocate
55    Len(u64),
56}
57
58pub(crate) struct IOEventMerged<C: CbArgs> {
59    pub buf: Buffer,
60    pub args: Option<C>,
61}
62
63impl<C: CbArgs> fmt::Debug for IOEvent<C> {
64    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
65        if let Some(TaskArgs::Merged(sub_tasks)) = self.args.as_ref() {
66            write!(f, "offset={} {:?} merged {}", self.offset, self.action, sub_tasks.len())
67        } else {
68            write!(f, "offset={} {:?}", self.offset, self.action)
69        }
70    }
71}
72
73impl<C: CbArgs> IOEvent<C> {
74    /// For IOAction::Read / IOAction::Write
75    #[inline]
76    pub fn new(fd: RawFd, buf: Buffer, action: IOAction, offset: i64) -> Self {
77        log_assert!(!buf.is_empty(), "{:?} offset={}, buffer size == 0", action, offset);
78        Self { buf_or_len: BufOrLen::Buffer(buf), fd, action, offset, res: i32::MIN, args: None }
79    }
80
81    /// For IOAction::Alloc / IOAction::Fsync
82    #[inline]
83    pub fn new_no_buf(fd: RawFd, action: IOAction, offset: i64, len: u64) -> Self {
84        Self {
85            buf_or_len: BufOrLen::Len(len), // No buffer for this action
86            fd,
87            action,
88            offset,
89            res: i32::MIN,
90            args: None,
91        }
92    }
93
94    #[inline(always)]
95    pub fn set_fd(&mut self, fd: RawFd) {
96        self.fd = fd;
97    }
98
99    /// Set callback for IOEvent, might be closure or a custom struct
100    #[inline(always)]
101    pub fn set_args(&mut self, args: C) {
102        self.args.replace(TaskArgs::Callback(args));
103    }
104
105    #[inline(always)]
106    pub fn get_size(&self) -> u64 {
107        match &self.buf_or_len {
108            BufOrLen::Buffer(buf) => buf.len() as u64,
109            BufOrLen::Len(l) => *l,
110        }
111    }
112
113    /// Set merged buffer and subtasks for the master event after merging.
114    #[inline(always)]
115    pub(crate) fn set_merged_tasks(
116        &mut self, merged_buf: Buffer, sub_tasks: SegList<IOEventMerged<C>>,
117    ) {
118        self.buf_or_len = BufOrLen::Buffer(merged_buf);
119        self.args.replace(TaskArgs::Merged(sub_tasks));
120    }
121
122    /// Convert this IOEvent into an IOEventMerged for storing in merge buffer.
123    /// Extracts the buffer and callback from the event.
124    #[inline(always)]
125    pub(crate) fn into_merged(mut self) -> IOEventMerged<C> {
126        let buf = match std::mem::replace(&mut self.buf_or_len, BufOrLen::Len(0)) {
127            BufOrLen::Buffer(buf) => buf,
128            BufOrLen::Len(_) => panic!("into_merged called on IOEvent with no buffer"),
129        };
130        let args = match self.args.take() {
131            Some(TaskArgs::Callback(args)) => Some(args),
132            _ => None,
133        };
134        IOEventMerged { buf, args }
135    }
136
137    /// Extract buffer and callback to create IOEventMerged, leaving this event with empty buffer.
138    /// Used when moving first event to merged_events list.
139    #[inline(always)]
140    pub(crate) fn extract_merged(&mut self) -> IOEventMerged<C> {
141        let buf = match std::mem::replace(&mut self.buf_or_len, BufOrLen::Len(0)) {
142            BufOrLen::Buffer(buf) => buf,
143            BufOrLen::Len(_) => panic!("extract_merged called on IOEvent with no buffer"),
144        };
145        let args = match self.args.take() {
146            Some(TaskArgs::Callback(args)) => Some(args),
147            _ => None,
148        };
149        IOEventMerged { buf, args }
150    }
151
152    /// return (offset, ptr, len)
153    #[inline(always)]
154    pub(crate) fn get_param_for_io(&mut self) -> (u64, *mut u8, u32) {
155        if let BufOrLen::Buffer(buf) = &mut self.buf_or_len {
156            let mut offset = self.offset as u64;
157            let mut p = buf.get_raw_mut();
158            let mut l = buf.len() as u32;
159            if self.res <= 0 {
160                (offset, p, l)
161            } else {
162                // resubmited I/O
163                offset += self.res as u64;
164                p = unsafe { p.add(self.res as usize) };
165                l += self.res as u32;
166                (offset, p, l)
167            }
168        } else {
169            panic!("get_buf_raw called on IOEvent with no buffer");
170        }
171    }
172
173    #[inline(always)]
174    pub fn get_write_result(self) -> Result<(), Errno> {
175        let res = self.res;
176        if res >= 0 {
177            return Ok(());
178        } else if res == i32::MIN {
179            panic!("IOEvent get_result before it's done");
180        } else {
181            return Err(Errno::from_raw_os_error(-res));
182        }
183    }
184
185    /// Get the result of the IO operation (bytes read/written or error).
186    /// Returns the number of bytes successfully transferred.
187    #[inline(always)]
188    pub fn get_result(&self) -> Result<usize, Errno> {
189        let res = self.res;
190        if res >= 0 {
191            return Ok(res as usize);
192        } else if res == i32::MIN {
193            panic!("IOEvent get_result before it's done");
194        } else {
195            return Err(Errno::from_raw_os_error(-res));
196        }
197    }
198
199    /// Get the buffer from a read operation.
200    /// Note: The buffer length is NOT modified. Use `get_result()` to get actual bytes read.
201    #[inline(always)]
202    pub fn get_read_result(mut self) -> Result<Buffer, Errno> {
203        let res = self.res;
204        if res >= 0 {
205            // XXX?
206            let buf_or_len = std::mem::replace(&mut self.buf_or_len, BufOrLen::Len(0));
207            if let BufOrLen::Buffer(buf) = buf_or_len {
208                // Do NOT modify buffer length - caller should use get_result() to know actual bytes read
209                return Ok(buf);
210            } else {
211                panic!("get_read_result called on IOEvent with no buffer");
212            }
213        } else if res == i32::MIN {
214            panic!("IOEvent get_result before it's done");
215        } else {
216            return Err(Errno::from_raw_os_error(-res));
217        }
218    }
219
220    #[inline(always)]
221    pub(crate) fn set_error(&mut self, mut errno: i32) {
222        if errno == 0 {
223            // TODO when errno == 0?
224            // XXX: EOF does not have code to represent,
225            // also when offset is not align to 4096, may return result 0,
226            errno = Errno::INVAL.raw_os_error();
227        }
228        if errno > 0 {
229            errno = -errno;
230        }
231        self.res = errno;
232    }
233
234    #[inline(always)]
235    pub(crate) fn set_copied(&mut self, len: usize) {
236        if self.res == i32::MIN {
237            // the initial state
238            self.res = len as i32;
239        } else {
240            // resubmit for short I/O
241            self.res += len as i32;
242        }
243    }
244
245    /// For writing custom callback workers
246    ///
247    /// Callback worker should always call this function on receiving IOEvent from Driver
248    ///
249    /// parameter: `check_short_read(offset: u64)` should be checking the offset exceed file end.
250    /// If `check_short_read()` return true, the callback function will return Err(IOEvent) for I/O resubmit.
251    ///
252    /// NOTE: you should always use a weak reference in `check_short_read` closure and
253    /// re-submission.
254    #[inline(always)]
255    pub fn callback<F, B>(mut self: Box<Self>, check_short_read: F, cb: B) -> Result<(), Box<Self>>
256    where
257        F: FnOnce(u64) -> bool,
258        B: Fn(C, i64, Result<Option<Buffer>, Errno>),
259    {
260        if self.res >= 0 {
261            if let BufOrLen::Buffer(buf) = &mut self.buf_or_len {
262                if buf.len() == self.res as usize {
263                    // most frequent case in the front, for cpu branch prediction
264                    self._callback_unchecked::<B>(false, cb);
265                } else if self.action == IOAction::Read {
266                    if check_short_read(self.offset as u64 + self.res as u64) {
267                        return Err(self);
268                    } else {
269                        // reach file ending
270                        buf.set_len(self.res as usize);
271                        self._callback_unchecked::<B>(false, cb);
272                    }
273                } else {
274                    // short write always need to resubmit
275                    return Err(self);
276                }
277            } else {
278                self._callback_unchecked::<B>(false, cb);
279            }
280        }
281        Ok(())
282    }
283
284    /// Perform callback on the IOEvent when cannot re-submit for short i/o
285    #[inline(always)]
286    pub fn callback_unchecked<B>(self, cb: B)
287    where
288        B: Fn(C, i64, Result<Option<Buffer>, Errno>),
289    {
290        self._callback_unchecked::<B>(true, cb);
291    }
292
293    /// Perform callback on the IOEvent when cannot re-submit for short i/o
294    ///
295    /// # Arguments
296    ///
297    /// - to_fix_short_io: should always be true, fix the buffer len of short I/O
298    ///
299    /// # Safety
300    ///
301    /// Only for callback worker does not re-submit when short I/O.
302    /// Buffer::len() will changed to actual I/O copied size during callback.
303    #[inline(always)]
304    pub(crate) fn _callback_unchecked<B>(mut self, to_fix_short_io: bool, cb: B)
305    where
306        B: Fn(C, i64, Result<Option<Buffer>, Errno>),
307    {
308        match self.args.take() {
309            Some(TaskArgs::Callback(args)) => {
310                let res: Result<Option<Buffer>, Errno> = if self.res >= 0 {
311                    match self.buf_or_len {
312                        BufOrLen::Buffer(mut buf) => {
313                            if to_fix_short_io && buf.len() > self.res as usize {
314                                buf.set_len(self.res as usize);
315                            }
316                            Ok(Some(buf))
317                        }
318                        BufOrLen::Len(_) => Ok(None),
319                    }
320                } else {
321                    Err(Errno::from_raw_os_error(-self.res))
322                };
323                cb(args, self.offset, res);
324            }
325            Some(TaskArgs::Merged(sub_tasks)) => {
326                if self.res >= 0 {
327                    let mut offset = self.offset;
328                    if self.action == IOAction::Read {
329                        if let BufOrLen::Buffer(parent_buf) = &self.buf_or_len {
330                            let mut b: &[u8] = &parent_buf[0..self.res as usize];
331                            for IOEventMerged { mut buf, args } in sub_tasks {
332                                if let Some(_args) = args {
333                                    let copied = safe_copy(&mut buf, b);
334                                    if copied < buf.len() {
335                                        buf.set_len(copied); // short I/O
336                                    }
337                                    cb(_args, offset, Ok(Some(buf)));
338                                    b = &b[copied..];
339                                    offset += copied as i64
340                                }
341                            }
342                        }
343                    } else if self.action == IOAction::Write {
344                        let mut l = self.res as usize;
345                        for IOEventMerged { mut buf, args } in sub_tasks {
346                            let mut copied = buf.len();
347                            if copied > l {
348                                // short write
349                                copied = l;
350                                buf.set_len(l);
351                            }
352                            if let Some(_args) = args {
353                                cb(_args, offset, Ok(Some(buf)));
354                            }
355                            l -= copied;
356                            offset += copied as i64;
357                        }
358                    }
359                } else {
360                    let mut offset = self.offset;
361                    for IOEventMerged { buf, args } in sub_tasks {
362                        let _l = buf.len() as i64;
363                        if let Some(_args) = args {
364                            cb(_args, offset, Err(Errno::from_raw_os_error(-self.res)));
365                        }
366                        offset += _l;
367                    }
368                }
369            }
370            None => {}
371        }
372    }
373}
374
375#[cfg(test)]
376mod tests {
377
378    use super::*;
379    use io_buffer::Buffer;
380    use rustix::io::Errno;
381    use std::mem::size_of;
382    use std::sync::Arc;
383    use std::sync::atomic::{AtomicI64, Ordering};
384
385    #[test]
386    fn test_ioevent_size() {
387        println!("IOEvent size {}", size_of::<IOEvent<()>>());
388        println!("BufOrLen size {}", size_of::<crate::tasks::BufOrLen>());
389        println!("IOEventMerged size {}", size_of::<IOEventMerged<()>>());
390    }
391
392    /// Test normal callback (non-merged case)
393    #[test]
394    fn test_callback_normal() {
395        let buffer = Buffer::alloc(4096).unwrap();
396        let mut event = IOEvent::<()>::new(0, buffer, IOAction::Write, 1024);
397
398        let result = Arc::new(std::sync::Mutex::new(None));
399        let result_clone = result.clone();
400
401        event.set_args(());
402        event.set_copied(4096);
403        event.callback_unchecked(move |_args, offset, res| {
404            *result_clone.lock().unwrap() = Some((offset, res));
405        });
406
407        let (offset, res) = result.lock().unwrap().take().unwrap();
408        assert_eq!(offset, 1024);
409        assert!(res.is_ok());
410        assert!(res.unwrap().is_some());
411    }
412
413    /// Test merged read callback - verifies offset correctness
414    #[test]
415    fn test_callback_merged_read() {
416        let offsets = Arc::new([AtomicI64::new(0), AtomicI64::new(0), AtomicI64::new(0)]);
417        let offsets_clone = offsets.clone();
418
419        // Create sub-tasks with their own buffers first
420        let mut sub_tasks = SegList::new();
421
422        sub_tasks.push(IOEventMerged { buf: Buffer::alloc(16).unwrap(), args: Some(()) });
423
424        sub_tasks.push(IOEventMerged { buf: Buffer::alloc(16).unwrap(), args: Some(()) });
425
426        sub_tasks.push(IOEventMerged { buf: Buffer::alloc(16).unwrap(), args: Some(()) });
427
428        // Create parent buffer and event
429        let parent_buf = Buffer::alloc(48).unwrap();
430        let mut event = IOEvent::<()>::new(0, parent_buf, IOAction::Read, 1000);
431        event.set_copied(48); // 48 bytes read
432
433        // Get the parent buffer back and fill with data
434        let parent_buf = match std::mem::replace(&mut event.buf_or_len, BufOrLen::Len(0)) {
435            BufOrLen::Buffer(buf) => buf,
436            BufOrLen::Len(_) => panic!("expected buffer"),
437        };
438        let mut parent_buf = parent_buf;
439        parent_buf.copy_from(0, b"ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@#$%^&*()");
440
441        event.set_merged_tasks(parent_buf, sub_tasks);
442        event.callback_unchecked(move |(), offset, res| {
443            let idx = (offset - 1000) / 16;
444            offsets_clone[idx as usize].store(offset, Ordering::SeqCst);
445            assert!(res.is_ok());
446            assert!(res.unwrap().is_some());
447        });
448
449        // Verify offsets
450        assert_eq!(offsets[0].load(Ordering::SeqCst), 1000);
451        assert_eq!(offsets[1].load(Ordering::SeqCst), 1016);
452        assert_eq!(offsets[2].load(Ordering::SeqCst), 1032);
453    }
454
455    /// Test merged write callback - verifies offset correctness
456    #[test]
457    fn test_callback_merged_write() {
458        let parent_buf = Buffer::alloc(4096).unwrap();
459
460        let mut event = IOEvent::<()>::new(0, parent_buf, IOAction::Write, 2000);
461        event.set_copied(48); // All 48 bytes written
462
463        let mut sub_tasks = SegList::new();
464
465        sub_tasks.push(IOEventMerged { buf: Buffer::alloc(16).unwrap(), args: Some(()) });
466
467        sub_tasks.push(IOEventMerged { buf: Buffer::alloc(16).unwrap(), args: Some(()) });
468
469        sub_tasks.push(IOEventMerged { buf: Buffer::alloc(16).unwrap(), args: Some(()) });
470
471        event.set_merged_tasks(Buffer::alloc(4096).unwrap(), sub_tasks);
472        event.callback_unchecked(move |(), offset, res| {
473            assert!(offset >= 2000 && offset <= 2032);
474            assert!(res.is_ok());
475            assert!(res.unwrap().is_some());
476        });
477    }
478
479    /// Test merged callback with error result
480    #[test]
481    fn test_callback_merged_error() {
482        let parent_buf = Buffer::alloc(4096).unwrap();
483        let mut event = IOEvent::<()>::new(0, parent_buf, IOAction::Read, 3000);
484        event.set_error(Errno::IO.raw_os_error()); // IO error
485
486        let mut sub_tasks = SegList::new();
487
488        sub_tasks.push(IOEventMerged { buf: Buffer::alloc(16).unwrap(), args: Some(()) });
489
490        sub_tasks.push(IOEventMerged { buf: Buffer::alloc(16).unwrap(), args: Some(()) });
491
492        event.set_merged_tasks(Buffer::alloc(48).unwrap(), sub_tasks);
493        event.callback_unchecked(|(), offset, res| {
494            assert!(offset == 3000 || offset == 3016);
495            assert!(res.is_err());
496            assert_eq!(res.err().unwrap(), Errno::IO);
497        });
498    }
499
500    /// Test short read in merged callback
501    #[test]
502    fn test_callback_merged_short_read() {
503        let offsets = Arc::new([AtomicI64::new(0), AtomicI64::new(0)]);
504        let offsets_clone = offsets.clone();
505
506        // Parent buffer with 32 bytes
507        let parent_buf = Buffer::alloc(32).unwrap();
508        let mut event = IOEvent::<()>::new(0, parent_buf, IOAction::Read, 4000);
509        event.set_copied(24); // Short read: only 24 bytes (16 + 8)
510
511        let mut sub_tasks = SegList::new();
512
513        sub_tasks.push(IOEventMerged { buf: Buffer::alloc(16).unwrap(), args: Some(()) });
514
515        sub_tasks.push(IOEventMerged { buf: Buffer::alloc(16).unwrap(), args: Some(()) });
516
517        let parent_buf = match std::mem::replace(&mut event.buf_or_len, BufOrLen::Len(0)) {
518            BufOrLen::Buffer(buf) => buf,
519            BufOrLen::Len(_) => panic!("expected buffer"),
520        };
521
522        event.set_merged_tasks(parent_buf, sub_tasks);
523        event.callback_unchecked(move |(), offset, res| {
524            let idx = (offset - 4000) / 16;
525            offsets_clone[idx as usize].store(offset, Ordering::SeqCst);
526            assert!(res.is_ok());
527            assert!(res.unwrap().is_some());
528        });
529
530        // Verify
531        assert_eq!(offsets[0].load(Ordering::SeqCst), 4000);
532        assert_eq!(offsets[1].load(Ordering::SeqCst), 4016);
533    }
534}
io_engine/tasks.rs

io_engine/
tasks.rs
