io_engine/
tasks.rs

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

io_engine/
tasks.rs
