cros_async/sys/linux/
uring_source.rs

1// Copyright 2020 The ChromiumOS Authors
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use std::ops::Deref;
6use std::ops::DerefMut;
7use std::sync::Arc;
8
9use base::sys::FallocateMode;
10use base::AsRawDescriptor;
11
12use super::uring_executor::RegisteredSource;
13use super::uring_executor::Result;
14use super::uring_executor::UringReactor;
15use crate::common_executor::RawExecutor;
16use crate::mem::BackingMemory;
17use crate::mem::MemRegion;
18use crate::mem::VecIoWrapper;
19use crate::AsyncResult;
20use crate::IoOptions;
21
22/// `UringSource` wraps FD backed IO sources for use with io_uring. It is a thin wrapper around
23/// registering an IO source with the uring that provides an `IoSource` implementation.
24pub struct UringSource<F: AsRawDescriptor> {
25    registered_source: RegisteredSource,
26    source: F,
27}
28
29impl<F: AsRawDescriptor> UringSource<F> {
30    /// Creates a new `UringSource` that wraps the given `io_source` object.
31    pub fn new(io_source: F, ex: &Arc<RawExecutor<UringReactor>>) -> Result<UringSource<F>> {
32        let r = ex.reactor.register_source(ex, &io_source)?;
33        Ok(UringSource {
34            registered_source: r,
35            source: io_source,
36        })
37    }
38
39    /// Reads from the iosource at `file_offset` and fill the given `vec`.
40    pub async fn read_to_vec(
41        &self,
42        file_offset: Option<u64>,
43        vec: Vec<u8>,
44        _options: IoOptions,
45    ) -> AsyncResult<(usize, Vec<u8>)> {
46        let buf = Arc::new(VecIoWrapper::from(vec));
47        let op = self.registered_source.start_read_to_mem(
48            file_offset,
49            buf.clone(),
50            [MemRegion {
51                offset: 0,
52                len: buf.len(),
53            }],
54        )?;
55        let len = op.await?;
56        let bytes = if let Ok(v) = Arc::try_unwrap(buf) {
57            v.into()
58        } else {
59            panic!("too many refs on buf");
60        };
61
62        Ok((len as usize, bytes))
63    }
64
65    /// Wait for the FD of `self` to be readable.
66    pub async fn wait_readable(&self) -> AsyncResult<()> {
67        let op = self.registered_source.poll_fd_readable()?;
68        op.await?;
69        Ok(())
70    }
71
72    /// Reads to the given `mem` at the given offsets from the file starting at `file_offset`.
73    pub async fn read_to_mem(
74        &self,
75        file_offset: Option<u64>,
76        mem: Arc<dyn BackingMemory + Send + Sync>,
77        mem_offsets: impl IntoIterator<Item = MemRegion>,
78        _options: IoOptions,
79    ) -> AsyncResult<usize> {
80        let op = self
81            .registered_source
82            .start_read_to_mem(file_offset, mem, mem_offsets)?;
83        let len = op.await?;
84        Ok(len as usize)
85    }
86
87    /// Writes from the given `vec` to the file starting at `file_offset`.
88    pub async fn write_from_vec(
89        &self,
90        file_offset: Option<u64>,
91        vec: Vec<u8>,
92        _options: IoOptions,
93    ) -> AsyncResult<(usize, Vec<u8>)> {
94        let buf = Arc::new(VecIoWrapper::from(vec));
95        let op = self.registered_source.start_write_from_mem(
96            file_offset,
97            buf.clone(),
98            [MemRegion {
99                offset: 0,
100                len: buf.len(),
101            }],
102        )?;
103        let len = op.await?;
104        let bytes = if let Ok(v) = Arc::try_unwrap(buf) {
105            v.into()
106        } else {
107            panic!("too many refs on buf");
108        };
109
110        Ok((len as usize, bytes))
111    }
112
113    /// Writes from the given `mem` from the given offsets to the file starting at `file_offset`.
114    pub async fn write_from_mem(
115        &self,
116        file_offset: Option<u64>,
117        mem: Arc<dyn BackingMemory + Send + Sync>,
118        mem_offsets: impl IntoIterator<Item = MemRegion>,
119        _options: IoOptions,
120    ) -> AsyncResult<usize> {
121        let op = self
122            .registered_source
123            .start_write_from_mem(file_offset, mem, mem_offsets)?;
124        let len = op.await?;
125        Ok(len as usize)
126    }
127
128    /// Deallocates the given range of a file.
129    pub async fn punch_hole(&self, file_offset: u64, len: u64) -> AsyncResult<()> {
130        let op = self.registered_source.start_fallocate(
131            file_offset,
132            len,
133            FallocateMode::PunchHole.into(),
134        )?;
135        let _ = op.await?;
136        Ok(())
137    }
138
139    /// Fills the given range with zeroes.
140    pub async fn write_zeroes_at(&self, file_offset: u64, len: u64) -> AsyncResult<()> {
141        let op = self.registered_source.start_fallocate(
142            file_offset,
143            len,
144            FallocateMode::ZeroRange.into(),
145        )?;
146        let _ = op.await?;
147        Ok(())
148    }
149
150    /// Sync all completed write operations to the backing storage.
151    pub async fn fsync(&self) -> AsyncResult<()> {
152        let op = self.registered_source.start_fsync()?;
153        let _ = op.await?;
154        Ok(())
155    }
156
157    /// Sync all data of completed write operations to the backing storage. Currently, the
158    /// implementation is equivalent to fsync.
159    pub async fn fdatasync(&self) -> AsyncResult<()> {
160        // Currently io_uring does not implement fdatasync. Fall back to fsync.
161        // TODO(b/281609112): Implement real fdatasync with io_uring.
162        self.fsync().await
163    }
164
165    /// Yields the underlying IO source.
166    pub fn into_source(self) -> F {
167        self.source
168    }
169
170    /// Provides a mutable ref to the underlying IO source.
171    pub fn as_source(&self) -> &F {
172        &self.source
173    }
174
175    /// Provides a ref to the underlying IO source.
176    pub fn as_source_mut(&mut self) -> &mut F {
177        &mut self.source
178    }
179}
180
181impl<F: AsRawDescriptor> Deref for UringSource<F> {
182    type Target = F;
183
184    fn deref(&self) -> &Self::Target {
185        &self.source
186    }
187}
188
189impl<F: AsRawDescriptor> DerefMut for UringSource<F> {
190    fn deref_mut(&mut self) -> &mut Self::Target {
191        &mut self.source
192    }
193}
194
195// NOTE: Prefer adding tests to io_source.rs if not backend specific.
196#[cfg(test)]
197mod tests {
198    use std::fs::File;
199    use std::future::Future;
200    use std::pin::Pin;
201    use std::task::Context;
202    use std::task::Poll;
203    use std::task::Waker;
204
205    use sync::Mutex;
206
207    use super::super::uring_executor::is_uring_stable;
208    use super::super::UringSource;
209    use super::*;
210    use crate::sys::linux::ExecutorKindSys;
211    use crate::Executor;
212    use crate::ExecutorTrait;
213    use crate::IoSource;
214
215    async fn read_u64<T: AsRawDescriptor>(source: &UringSource<T>) -> u64 {
216        // Init a vec that translates to u64::max;
217        let u64_mem = vec![0xffu8; std::mem::size_of::<u64>()];
218        let (ret, u64_mem) = source
219            .read_to_vec(None, u64_mem, Default::default())
220            .await
221            .unwrap();
222        assert_eq!(ret, std::mem::size_of::<u64>());
223        let mut val = 0u64.to_ne_bytes();
224        val.copy_from_slice(&u64_mem);
225        u64::from_ne_bytes(val)
226    }
227
228    #[test]
229    fn event() {
230        if !is_uring_stable() {
231            return;
232        }
233
234        use base::Event;
235        use base::EventExt;
236
237        async fn write_event(ev: Event, wait: Event, ex: &Arc<RawExecutor<UringReactor>>) {
238            let wait = UringSource::new(wait, ex).unwrap();
239            ev.write_count(55).unwrap();
240            read_u64(&wait).await;
241            ev.write_count(66).unwrap();
242            read_u64(&wait).await;
243            ev.write_count(77).unwrap();
244            read_u64(&wait).await;
245        }
246
247        async fn read_events(ev: Event, signal: Event, ex: &Arc<RawExecutor<UringReactor>>) {
248            let source = UringSource::new(ev, ex).unwrap();
249            assert_eq!(read_u64(&source).await, 55);
250            signal.signal().unwrap();
251            assert_eq!(read_u64(&source).await, 66);
252            signal.signal().unwrap();
253            assert_eq!(read_u64(&source).await, 77);
254            signal.signal().unwrap();
255        }
256
257        let event = Event::new().unwrap();
258        let signal_wait = Event::new().unwrap();
259        let ex = RawExecutor::<UringReactor>::new().unwrap();
260        let write_task = write_event(
261            event.try_clone().unwrap(),
262            signal_wait.try_clone().unwrap(),
263            &ex,
264        );
265        let read_task = read_events(event, signal_wait, &ex);
266        ex.run_until(futures::future::join(read_task, write_task))
267            .unwrap();
268    }
269
270    #[test]
271    fn pend_on_pipe() {
272        if !is_uring_stable() {
273            return;
274        }
275
276        use std::io::Write;
277
278        use futures::future::Either;
279
280        async fn do_test(ex: &Arc<RawExecutor<UringReactor>>) {
281            let (read_source, mut w) = base::pipe().unwrap();
282            let source = UringSource::new(read_source, ex).unwrap();
283            let done = Box::pin(async { 5usize });
284            let pending = Box::pin(read_u64(&source));
285            match futures::future::select(pending, done).await {
286                Either::Right((5, pending)) => {
287                    // Write to the pipe so that the kernel will release the memory associated with
288                    // the uring read operation.
289                    w.write_all(&[0]).expect("failed to write to pipe");
290                    ::std::mem::drop(pending);
291                }
292                _ => panic!("unexpected select result"),
293            };
294        }
295
296        let ex = RawExecutor::<UringReactor>::new().unwrap();
297        ex.run_until(do_test(&ex)).unwrap();
298    }
299
300    #[test]
301    fn range_error() {
302        if !is_uring_stable() {
303            return;
304        }
305
306        async fn go(ex: &Arc<RawExecutor<UringReactor>>) {
307            let f = File::open("/dev/zero").unwrap();
308            let source = UringSource::new(f, ex).unwrap();
309            let v = vec![0x55u8; 64];
310            let vw = Arc::new(VecIoWrapper::from(v));
311            let ret = source
312                .read_to_mem(
313                    None,
314                    Arc::<VecIoWrapper>::clone(&vw),
315                    [MemRegion {
316                        offset: 32,
317                        len: 33,
318                    }],
319                    Default::default(),
320                )
321                .await;
322            assert!(ret.is_err());
323        }
324
325        let ex = RawExecutor::<UringReactor>::new().unwrap();
326        ex.run_until(go(&ex)).unwrap();
327    }
328
329    // Single-region ops use `IORING_OP_READ`/`IORING_OP_WRITE` while multi-region ops use
330    // readv/writev; both must move the same bytes to the same places.
331    #[test]
332    fn single_and_multi_region_io() {
333        if !is_uring_stable() {
334            return;
335        }
336
337        const EXPECTED: &[u8; 24] = b"0123456789abcdef89ab0123";
338
339        async fn go(ex: &Arc<RawExecutor<UringReactor>>, f: File) {
340            let source = UringSource::new(f, ex).unwrap();
341
342            let src = Arc::new(VecIoWrapper::from(b"0123456789abcdef".to_vec()));
343            let written = source
344                .write_from_mem(
345                    Some(0),
346                    Arc::<VecIoWrapper>::clone(&src),
347                    [MemRegion { offset: 0, len: 16 }],
348                    Default::default(),
349                )
350                .await
351                .unwrap();
352            assert_eq!(written, 16);
353
354            // Append "89ab0123" from two disjoint regions of the same buffer.
355            let written = source
356                .write_from_mem(
357                    Some(16),
358                    Arc::<VecIoWrapper>::clone(&src),
359                    [
360                        MemRegion { offset: 8, len: 4 },
361                        MemRegion { offset: 0, len: 4 },
362                    ],
363                    Default::default(),
364                )
365                .await
366                .unwrap();
367            assert_eq!(written, 8);
368
369            for regions in [
370                vec![MemRegion { offset: 0, len: 24 }],
371                vec![
372                    MemRegion { offset: 0, len: 10 },
373                    MemRegion {
374                        offset: 10,
375                        len: 14,
376                    },
377                ],
378            ] {
379                let dst = Arc::new(VecIoWrapper::from(vec![0u8; EXPECTED.len()]));
380                let read = source
381                    .read_to_mem(
382                        Some(0),
383                        Arc::<VecIoWrapper>::clone(&dst),
384                        regions,
385                        Default::default(),
386                    )
387                    .await
388                    .unwrap();
389                assert_eq!(read, EXPECTED.len());
390                let dst = Arc::try_unwrap(dst).unwrap_or_else(|_| panic!("too many refs on dst"));
391                assert_eq!(Vec::<u8>::from(dst), EXPECTED);
392            }
393        }
394
395        let f = tempfile::tempfile().unwrap();
396        let ex = RawExecutor::<UringReactor>::new().unwrap();
397        ex.run_until(go(&ex, f)).unwrap();
398    }
399
400    #[test]
401    fn wait_read() {
402        if !is_uring_stable() {
403            return;
404        }
405
406        async fn go(ex: &Arc<RawExecutor<UringReactor>>) {
407            let f = File::open("/dev/zero").unwrap();
408            let source = UringSource::new(f, ex).unwrap();
409            source.wait_readable().await.unwrap();
410        }
411
412        let ex = RawExecutor::<UringReactor>::new().unwrap();
413        ex.run_until(go(&ex)).unwrap();
414    }
415
416    struct State {
417        should_quit: bool,
418        waker: Option<Waker>,
419    }
420
421    impl State {
422        fn wake(&mut self) {
423            self.should_quit = true;
424            let waker = self.waker.take();
425
426            if let Some(waker) = waker {
427                waker.wake();
428            }
429        }
430    }
431
432    struct Quit {
433        state: Arc<Mutex<State>>,
434    }
435
436    impl Future for Quit {
437        type Output = ();
438
439        fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<()> {
440            let mut state = self.state.lock();
441            if state.should_quit {
442                return Poll::Ready(());
443            }
444
445            state.waker = Some(cx.waker().clone());
446            Poll::Pending
447        }
448    }
449
450    #[cfg(any(target_os = "android", target_os = "linux"))]
451    #[test]
452    fn await_uring_from_poll() {
453        if !is_uring_stable() {
454            return;
455        }
456        // Start a uring operation and then await the result from an FdExecutor.
457        async fn go(source: IoSource<File>) {
458            let v = vec![0xa4u8; 16];
459            let (len, vec) = source
460                .read_to_vec(None, v, Default::default())
461                .await
462                .unwrap();
463            assert_eq!(len, 16);
464            assert!(vec.iter().all(|&b| b == 0));
465        }
466
467        let state = Arc::new(Mutex::new(State {
468            should_quit: false,
469            waker: None,
470        }));
471
472        let uring_ex = Executor::with_executor_kind(ExecutorKindSys::Uring.into()).unwrap();
473        let f = File::open("/dev/zero").unwrap();
474        let source = uring_ex.async_from(f).unwrap();
475
476        let quit = Quit {
477            state: state.clone(),
478        };
479        let handle = std::thread::spawn(move || uring_ex.run_until(quit));
480
481        let poll_ex = Executor::with_executor_kind(ExecutorKindSys::Fd.into()).unwrap();
482        poll_ex.run_until(go(source)).unwrap();
483
484        state.lock().wake();
485        handle.join().unwrap().unwrap();
486    }
487
488    #[cfg(any(target_os = "android", target_os = "linux"))]
489    #[test]
490    fn await_poll_from_uring() {
491        if !is_uring_stable() {
492            return;
493        }
494        // Start a poll operation and then await the result
495        async fn go(source: IoSource<File>) {
496            let v = vec![0x2cu8; 16];
497            let (len, vec) = source
498                .read_to_vec(None, v, Default::default())
499                .await
500                .unwrap();
501            assert_eq!(len, 16);
502            assert!(vec.iter().all(|&b| b == 0));
503        }
504
505        let state = Arc::new(Mutex::new(State {
506            should_quit: false,
507            waker: None,
508        }));
509
510        let poll_ex = Executor::with_executor_kind(ExecutorKindSys::Fd.into()).unwrap();
511        let f = File::open("/dev/zero").unwrap();
512        let source = poll_ex.async_from(f).unwrap();
513
514        let quit = Quit {
515            state: state.clone(),
516        };
517        let handle = std::thread::spawn(move || poll_ex.run_until(quit));
518
519        let uring_ex = Executor::with_executor_kind(ExecutorKindSys::Uring.into()).unwrap();
520        uring_ex.run_until(go(source)).unwrap();
521
522        state.lock().wake();
523        handle.join().unwrap().unwrap();
524    }
525}