aiofut/
lib.rs

1//! Straightforward Linux AIO using Futures/async/await.
2//!
3//! # Example
4//!
5//! Use aiofut to schedule writes to a file:
6//!
7//! ```rust
8//! use futures::{executor::LocalPool, future::FutureExt, task::LocalSpawnExt};
9//! use aiofut::AIOBuilder;
10//! use std::os::unix::io::AsRawFd;
11//! let mut aiomgr = AIOBuilder::default().build().unwrap();
12//! let file = std::fs::OpenOptions::new()
13//!     .read(true)
14//!     .write(true)
15//!     .create(true)
16//!     .truncate(true)
17//!     .open("test")
18//!     .unwrap();
19//! let fd = file.as_raw_fd();
20//! // keep all returned futures in a vector
21//! let ws = vec![(0, "hello"), (5, "world"), (2, "xxxx")]
22//!     .into_iter()
23//!     .map(|(off, s)| aiomgr.write(fd, off, s.as_bytes().into(), None))
24//!     .collect::<Vec<_>>();
25//! // here we use futures::executor::LocalPool to poll all futures
26//! let mut pool = LocalPool::new();
27//! let spawner = pool.spawner();
28//! for w in ws.into_iter() {
29//!     let h = spawner.spawn_local_with_handle(w).unwrap().map(|r| {
30//!         println!("wrote {} bytes", r.0.unwrap());
31//!     });
32//!     spawner.spawn_local(h).unwrap();
33//! }
34//! pool.run();
35//! ```
36
37mod abi;
38use parking_lot::Mutex;
39use std::collections::{hash_map, HashMap};
40use std::os::unix::io::RawFd;
41use std::pin::Pin;
42use std::os::raw::c_long;
43use libc::time_t;
44use std::sync::{
45    atomic::{AtomicPtr, AtomicUsize, Ordering},
46    Arc,
47};
48
49const LIBAIO_EAGAIN: libc::c_int = -libc::EAGAIN;
50const LIBAIO_ENOMEM: libc::c_int = -libc::ENOMEM;
51const LIBAIO_ENOSYS: libc::c_int = -libc::ENOSYS;
52
53#[derive(Debug)]
54pub enum Error {
55    MaxEventsTooLarge,
56    LowKernelRes,
57    NotSupported,
58    OtherError,
59}
60
61// NOTE: I assume it io_context_t is thread-safe, no?
62struct AIOContext(abi::IOContextPtr);
63unsafe impl Sync for AIOContext {}
64unsafe impl Send for AIOContext {}
65
66impl std::ops::Deref for AIOContext {
67    type Target = abi::IOContextPtr;
68    fn deref(&self) -> &abi::IOContextPtr {
69        &self.0
70    }
71}
72
73impl AIOContext {
74    fn new(maxevents: u32) -> Result<Self, Error> {
75        let mut ctx = std::ptr::null_mut();
76        unsafe {
77            match abi::io_setup(maxevents as libc::c_int, &mut ctx) {
78                0 => Ok(()),
79                LIBAIO_EAGAIN => Err(Error::MaxEventsTooLarge),
80                LIBAIO_ENOMEM => Err(Error::LowKernelRes),
81                LIBAIO_ENOSYS => Err(Error::NotSupported),
82                _ => Err(Error::OtherError),
83            }
84            .and_then(|_| Ok(AIOContext(ctx)))
85        }
86    }
87}
88
89impl Drop for AIOContext {
90    fn drop(&mut self) {
91        unsafe {
92            assert_eq!(abi::io_destroy(self.0), 0);
93        }
94    }
95}
96
97/// Represent the necessary data for an AIO operation. Memory-safe when moved.
98pub struct AIO {
99    // hold the buffer used by iocb
100    data: Option<Box<[u8]>>,
101    iocb: AtomicPtr<abi::IOCb>,
102    id: u64,
103}
104
105impl AIO {
106    fn new(
107        id: u64,
108        fd: RawFd,
109        off: u64,
110        data: Box<[u8]>,
111        priority: u16,
112        flags: u32,
113        opcode: abi::IOCmd,
114    ) -> Self {
115        let mut iocb = Box::new(abi::IOCb::default());
116        iocb.aio_fildes = fd as u32;
117        iocb.aio_lio_opcode = opcode as u16;
118        iocb.aio_reqprio = priority;
119        iocb.aio_buf = data.as_ptr() as u64;
120        iocb.aio_nbytes = data.len() as u64;
121        iocb.aio_offset = off;
122        iocb.aio_flags = flags;
123        iocb.aio_data = id;
124        let iocb = AtomicPtr::new(Box::into_raw(iocb));
125        let data = Some(data);
126        AIO { iocb, id, data }
127    }
128}
129
130impl Drop for AIO {
131    fn drop(&mut self) {
132        unsafe {
133            drop(Box::from_raw(self.iocb.load(Ordering::Acquire)));
134        }
135    }
136}
137
138/// The result of an AIO operation: the number of bytes written on success,
139/// or the errno on failure.
140pub type AIOResult = (Result<usize, i32>, Box<[u8]>);
141
142/// Represents a scheduled (future) asynchronous I/O operation, which gets executed (resolved)
143/// automatically.
144pub struct AIOFuture {
145    notifier: Arc<AIONotifier>,
146    aio_id: u64,
147}
148
149impl AIOFuture {
150    pub fn get_id(&self) -> u64 {
151        self.aio_id
152    }
153}
154
155impl std::future::Future for AIOFuture {
156    type Output = AIOResult;
157    fn poll(
158        self: Pin<&mut Self>,
159        cx: &mut std::task::Context,
160    ) -> std::task::Poll<Self::Output> {
161        if let Some(ret) = self.notifier.poll(self.aio_id, cx.waker()) {
162            std::task::Poll::Ready(ret)
163        } else {
164            std::task::Poll::Pending
165        }
166    }
167}
168
169impl Drop for AIOFuture {
170    fn drop(&mut self) {
171        self.notifier.dropped(self.aio_id)
172    }
173}
174
175enum AIOState {
176    FutureInit(AIO, bool),
177    FuturePending(AIO, std::task::Waker, bool),
178    FutureDone(AIOResult),
179}
180
181/// The state machine for finished AIO operations and wakes up the futures.
182pub struct AIONotifier {
183    waiting: Mutex<HashMap<u64, AIOState>>,
184    npending: AtomicUsize,
185    io_ctx: AIOContext,
186    #[cfg(feature = "emulated-failure")]
187    emul_fail: Option<EmulatedFailureShared>,
188}
189
190impl AIONotifier {
191    fn register_notify(&self, id: u64, state: AIOState) {
192        let mut waiting = self.waiting.lock();
193        assert!(waiting.insert(id, state).is_none());
194    }
195
196    fn dropped(&self, id: u64) {
197        let mut waiting = self.waiting.lock();
198        match waiting.entry(id) {
199            hash_map::Entry::Occupied(mut e) => match e.get_mut() {
200                AIOState::FutureInit(_, dropped) => *dropped = true,
201                AIOState::FuturePending(_, _, dropped) => *dropped = true,
202                AIOState::FutureDone(_) => {
203                    e.remove();
204                }
205            },
206            _ => (),
207        }
208    }
209
210    fn poll(&self, id: u64, waker: &std::task::Waker) -> Option<AIOResult> {
211        let mut waiting = self.waiting.lock();
212        match waiting.entry(id) {
213            hash_map::Entry::Occupied(e) => {
214                let v = e.remove();
215                match v {
216                    AIOState::FutureInit(aio, _) => {
217                        waiting.insert(
218                            id,
219                            AIOState::FuturePending(aio, waker.clone(), false),
220                        );
221                        None
222                    }
223                    AIOState::FuturePending(aio, waker, dropped) => {
224                        waiting.insert(
225                            id,
226                            AIOState::FuturePending(aio, waker, dropped),
227                        );
228                        None
229                    }
230                    AIOState::FutureDone(res) => Some(res),
231                }
232            }
233            _ => unreachable!(),
234        }
235    }
236
237    fn finish(&self, id: u64, res: i64) {
238        let mut w = self.waiting.lock();
239        self.npending.fetch_sub(1, Ordering::Relaxed);
240        match w.entry(id) {
241            hash_map::Entry::Occupied(e) => match e.remove() {
242                AIOState::FutureInit(mut aio, dropped) => {
243                    if !dropped {
244                        let data = aio.data.take().unwrap();
245                        w.insert(
246                            id,
247                            AIOState::FutureDone(if res >= 0 {
248                                (Ok(res as usize), data)
249                            } else {
250                                (Err(-res as i32), data)
251                            }),
252                        );
253                    }
254                }
255                AIOState::FuturePending(mut aio, waker, dropped) => {
256                    if !dropped {
257                        let data = aio.data.take().unwrap();
258                        w.insert(
259                            id,
260                            AIOState::FutureDone(if res >= 0 {
261                                (Ok(res as usize), data)
262                            } else {
263                                (Err(-res as i32), data)
264                            }),
265                        );
266                        waker.wake();
267                    }
268                }
269                AIOState::FutureDone(ret) => {
270                    w.insert(id, AIOState::FutureDone(ret));
271                }
272            },
273            _ => unreachable!(),
274        }
275    }
276}
277
278pub struct AIOBuilder {
279    max_events: u32,
280    max_nwait: u16,
281    max_nbatched: usize,
282    timeout: Option<u32>,
283    #[cfg(feature = "emulated-failure")]
284    emul_fail: Option<EmulatedFailureShared>,
285}
286
287impl Default for AIOBuilder {
288    fn default() -> Self {
289        AIOBuilder {
290            max_events: 128,
291            max_nwait: 128,
292            max_nbatched: 128,
293            timeout: None,
294            #[cfg(feature = "emulated-failure")]
295            emul_fail: None,
296        }
297    }
298}
299
300impl AIOBuilder {
301    /// Maximum concurrent async IO operations.
302    pub fn max_events(&mut self, v: u32) -> &mut Self {
303        self.max_events = v;
304        self
305    }
306
307    /// Maximum complete IOs per poll.
308    pub fn max_nwait(&mut self, v: u16) -> &mut Self {
309        self.max_nwait = v;
310        self
311    }
312
313    /// Maximum number of IOs per submission.
314    pub fn max_nbatched(&mut self, v: usize) -> &mut Self {
315        self.max_nbatched = v;
316        self
317    }
318
319    /// Timeout for a polling iteration (default is None).
320    pub fn timeout(&mut self, sec: u32) -> &mut Self {
321        self.timeout = Some(sec);
322        self
323    }
324
325    #[cfg(feature = "emulated-failure")]
326    pub fn emulated_failure(&mut self, ef: EmulatedFailureShared) -> &mut Self {
327        self.emul_fail = Some(ef);
328        self
329    }
330
331    /// Build an AIOManager object based on the configuration (and auto-start the background IO
332    /// scheduling thread).
333    pub fn build(&mut self) -> Result<AIOManager, Error> {
334        let (scheduler_in, scheduler_out) =
335            new_batch_scheduler(self.max_nbatched);
336        let (exit_s, exit_r) = crossbeam_channel::bounded(0);
337
338        let notifier = Arc::new(AIONotifier {
339            io_ctx: AIOContext::new(self.max_events)?,
340            waiting: Mutex::new(HashMap::new()),
341            npending: AtomicUsize::new(0),
342            #[cfg(feature = "emulated-failure")]
343            emul_fail: self.emul_fail.as_ref().map(|ef| ef.clone()),
344        });
345        let mut aiomgr = AIOManager {
346            notifier,
347            listener: None,
348            scheduler_in,
349            exit_s,
350        };
351        aiomgr.start(scheduler_out, exit_r, self.max_nwait, self.timeout)?;
352        Ok(aiomgr)
353    }
354}
355
356pub trait EmulatedFailure: Send {
357    fn tick(&mut self) -> Option<i64>;
358}
359
360pub type EmulatedFailureShared = Arc<Mutex<dyn EmulatedFailure>>;
361
362/// Manager all AIOs.
363pub struct AIOManager {
364    notifier: Arc<AIONotifier>,
365    scheduler_in: AIOBatchSchedulerIn,
366    listener: Option<std::thread::JoinHandle<()>>,
367    exit_s: crossbeam_channel::Sender<()>,
368}
369
370impl AIOManager {
371    fn start(
372        &mut self,
373        mut scheduler_out: AIOBatchSchedulerOut,
374        exit_r: crossbeam_channel::Receiver<()>,
375        max_nwait: u16,
376        timeout: Option<u32>,
377    ) -> Result<(), Error> {
378        let n = self.notifier.clone();
379        self.listener = Some(std::thread::spawn(move || {
380            let mut timespec = timeout.and_then(|sec: u32| {
381                Some(libc::timespec {
382                    tv_sec: sec as time_t,
383                    tv_nsec: 0,
384                })
385            });
386            let mut ongoing = 0;
387            loop {
388                // try to quiesce
389                if ongoing == 0 && scheduler_out.is_empty() {
390                    let mut sel = crossbeam_channel::Select::new();
391                    sel.recv(&exit_r);
392                    sel.recv(&scheduler_out.get_receiver());
393                    if sel.ready() == 0 {
394                        exit_r.recv().unwrap();
395                        break
396                    }
397                }
398                // submit as many aios as possible
399                loop {
400                    let nacc = scheduler_out.submit(&n);
401                    ongoing += nacc;
402                    if nacc == 0 {
403                        break
404                    }
405                }
406                // no need to wait if there is no progress
407                if ongoing == 0 {
408                    continue
409                }
410                // then block on any finishing aios
411                let mut events =
412                    vec![abi::IOEvent::default(); max_nwait as usize];
413                let ret = unsafe {
414                    abi::io_getevents(
415                        *n.io_ctx,
416                        1,
417                        max_nwait as c_long,
418                        events.as_mut_ptr(),
419                        timespec
420                            .as_mut()
421                            .and_then(|t| Some(t as *mut libc::timespec))
422                            .unwrap_or(std::ptr::null_mut()),
423                    )
424                };
425                // TODO: AIO fatal error handling
426                // avoid empty slice
427                if ret == 0 {
428                    continue
429                }
430                assert!(ret > 0);
431                ongoing -= ret as usize;
432                for ev in events[..ret as usize].iter() {
433                    #[cfg(not(feature = "emulated-failure"))]
434                    n.finish(ev.data as u64, ev.res);
435                    #[cfg(feature = "emulated-failure")]
436                    {
437                        let mut res = ev.res;
438                        if let Some(emul_fail) = n.emul_fail.as_ref() {
439                            let mut ef = emul_fail.lock();
440                            if let Some(e) = ef.tick() {
441                                res = e
442                            }
443                        }
444                        n.finish(ev.data as u64, res);
445                    }
446                }
447            }
448        }));
449        Ok(())
450    }
451
452    pub fn read(
453        &self,
454        fd: RawFd,
455        offset: u64,
456        length: usize,
457        priority: Option<u16>,
458    ) -> AIOFuture {
459        let priority = priority.unwrap_or(0);
460        let mut data = Vec::new();
461        data.resize(length, 0);
462        let data = data.into_boxed_slice();
463        let aio = AIO::new(
464            self.scheduler_in.next_id(),
465            fd,
466            offset,
467            data,
468            priority,
469            0,
470            abi::IOCmd::PRead,
471        );
472        self.scheduler_in.schedule(aio, &self.notifier)
473    }
474
475    pub fn write(
476        &self,
477        fd: RawFd,
478        offset: u64,
479        data: Box<[u8]>,
480        priority: Option<u16>,
481    ) -> AIOFuture {
482        let priority = priority.unwrap_or(0);
483        let aio = AIO::new(
484            self.scheduler_in.next_id(),
485            fd,
486            offset,
487            data,
488            priority,
489            0,
490            abi::IOCmd::PWrite,
491        );
492        self.scheduler_in.schedule(aio, &self.notifier)
493    }
494
495    /// Get a copy of the current data in the buffer.
496    pub fn copy_data(&self, aio_id: u64) -> Option<Vec<u8>> {
497        let w = self.notifier.waiting.lock();
498        w.get(&aio_id).and_then(|state| {
499            Some(
500                match state {
501                    AIOState::FutureInit(aio, _) => {
502                        &**aio.data.as_ref().unwrap()
503                    }
504                    AIOState::FuturePending(aio, _, _) => {
505                        &**aio.data.as_ref().unwrap()
506                    }
507                    AIOState::FutureDone(res) => &res.1,
508                }
509                .to_vec(),
510            )
511        })
512    }
513
514    /// Get the number of pending AIOs (approximation).
515    pub fn get_npending(&self) -> usize {
516        self.notifier.npending.load(Ordering::Relaxed)
517    }
518}
519
520impl Drop for AIOManager {
521    fn drop(&mut self) {
522        self.exit_s.send(()).unwrap();
523        self.listener.take().unwrap().join().unwrap();
524    }
525}
526
527pub struct AIOBatchSchedulerIn {
528    queue_in: crossbeam_channel::Sender<AtomicPtr<abi::IOCb>>,
529    last_id: std::cell::Cell<u64>,
530}
531
532pub struct AIOBatchSchedulerOut {
533    queue_out: crossbeam_channel::Receiver<AtomicPtr<abi::IOCb>>,
534    max_nbatched: usize,
535    leftover: Vec<AtomicPtr<abi::IOCb>>,
536}
537
538impl AIOBatchSchedulerIn {
539    fn schedule(&self, aio: AIO, notifier: &Arc<AIONotifier>) -> AIOFuture {
540        let fut = AIOFuture {
541            notifier: notifier.clone(),
542            aio_id: aio.id,
543        };
544        let iocb = aio.iocb.load(Ordering::Acquire);
545        notifier.register_notify(aio.id, AIOState::FutureInit(aio, false));
546        self.queue_in.send(AtomicPtr::new(iocb)).unwrap();
547        notifier.npending.fetch_add(1, Ordering::Relaxed);
548        fut
549    }
550
551    fn next_id(&self) -> u64 {
552        let id = self.last_id.get();
553        self.last_id.set(id.wrapping_add(1));
554        id
555    }
556}
557
558impl AIOBatchSchedulerOut {
559    fn get_receiver(
560        &self,
561    ) -> &crossbeam_channel::Receiver<AtomicPtr<abi::IOCb>> {
562        &self.queue_out
563    }
564    fn is_empty(&self) -> bool {
565        self.leftover.len() == 0
566    }
567    fn submit(&mut self, notifier: &AIONotifier) -> usize {
568        let mut quota = self.max_nbatched;
569        let mut pending = self
570            .leftover
571            .iter()
572            .map(|p| p.load(Ordering::Acquire))
573            .collect::<Vec<_>>();
574        if pending.len() < quota {
575            quota -= pending.len();
576            while let Ok(iocb) = self.queue_out.try_recv() {
577                pending.push(iocb.load(Ordering::Acquire));
578                quota -= 1;
579                if quota == 0 {
580                    break
581                }
582            }
583        }
584        if pending.len() == 0 {
585            return 0
586        }
587        let mut ret = unsafe {
588            abi::io_submit(
589                *notifier.io_ctx,
590                pending.len() as c_long,
591                (&mut pending).as_mut_ptr(),
592            )
593        };
594        if ret < 0 && ret == LIBAIO_EAGAIN {
595            ret = 0
596        }
597        let nacc = ret as usize;
598        self.leftover = (&pending[nacc..])
599            .iter()
600            .map(|p| AtomicPtr::new(*p))
601            .collect::<Vec<_>>();
602        nacc
603    }
604}
605
606/// Create the scheduler that submits AIOs in batches.
607fn new_batch_scheduler(
608    max_nbatched: usize,
609) -> (AIOBatchSchedulerIn, AIOBatchSchedulerOut) {
610    let (queue_in, queue_out) = crossbeam_channel::unbounded();
611    let bin = AIOBatchSchedulerIn {
612        queue_in,
613        last_id: std::cell::Cell::new(0),
614    };
615    let bout = AIOBatchSchedulerOut {
616        queue_out,
617        max_nbatched,
618        leftover: Vec::new(),
619    };
620    (bin, bout)
621}
aiofut/lib.rs

aiofut/
lib.rs
