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    #[test]
330    fn wait_read() {
331        if !is_uring_stable() {
332            return;
333        }
334
335        async fn go(ex: &Arc<RawExecutor<UringReactor>>) {
336            let f = File::open("/dev/zero").unwrap();
337            let source = UringSource::new(f, ex).unwrap();
338            source.wait_readable().await.unwrap();
339        }
340
341        let ex = RawExecutor::<UringReactor>::new().unwrap();
342        ex.run_until(go(&ex)).unwrap();
343    }
344
345    struct State {
346        should_quit: bool,
347        waker: Option<Waker>,
348    }
349
350    impl State {
351        fn wake(&mut self) {
352            self.should_quit = true;
353            let waker = self.waker.take();
354
355            if let Some(waker) = waker {
356                waker.wake();
357            }
358        }
359    }
360
361    struct Quit {
362        state: Arc<Mutex<State>>,
363    }
364
365    impl Future for Quit {
366        type Output = ();
367
368        fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<()> {
369            let mut state = self.state.lock();
370            if state.should_quit {
371                return Poll::Ready(());
372            }
373
374            state.waker = Some(cx.waker().clone());
375            Poll::Pending
376        }
377    }
378
379    #[cfg(any(target_os = "android", target_os = "linux"))]
380    #[test]
381    fn await_uring_from_poll() {
382        if !is_uring_stable() {
383            return;
384        }
385        // Start a uring operation and then await the result from an FdExecutor.
386        async fn go(source: IoSource<File>) {
387            let v = vec![0xa4u8; 16];
388            let (len, vec) = source
389                .read_to_vec(None, v, Default::default())
390                .await
391                .unwrap();
392            assert_eq!(len, 16);
393            assert!(vec.iter().all(|&b| b == 0));
394        }
395
396        let state = Arc::new(Mutex::new(State {
397            should_quit: false,
398            waker: None,
399        }));
400
401        let uring_ex = Executor::with_executor_kind(ExecutorKindSys::Uring.into()).unwrap();
402        let f = File::open("/dev/zero").unwrap();
403        let source = uring_ex.async_from(f).unwrap();
404
405        let quit = Quit {
406            state: state.clone(),
407        };
408        let handle = std::thread::spawn(move || uring_ex.run_until(quit));
409
410        let poll_ex = Executor::with_executor_kind(ExecutorKindSys::Fd.into()).unwrap();
411        poll_ex.run_until(go(source)).unwrap();
412
413        state.lock().wake();
414        handle.join().unwrap().unwrap();
415    }
416
417    #[cfg(any(target_os = "android", target_os = "linux"))]
418    #[test]
419    fn await_poll_from_uring() {
420        if !is_uring_stable() {
421            return;
422        }
423        // Start a poll operation and then await the result
424        async fn go(source: IoSource<File>) {
425            let v = vec![0x2cu8; 16];
426            let (len, vec) = source
427                .read_to_vec(None, v, Default::default())
428                .await
429                .unwrap();
430            assert_eq!(len, 16);
431            assert!(vec.iter().all(|&b| b == 0));
432        }
433
434        let state = Arc::new(Mutex::new(State {
435            should_quit: false,
436            waker: None,
437        }));
438
439        let poll_ex = Executor::with_executor_kind(ExecutorKindSys::Fd.into()).unwrap();
440        let f = File::open("/dev/zero").unwrap();
441        let source = poll_ex.async_from(f).unwrap();
442
443        let quit = Quit {
444            state: state.clone(),
445        };
446        let handle = std::thread::spawn(move || poll_ex.run_until(quit));
447
448        let uring_ex = Executor::with_executor_kind(ExecutorKindSys::Uring.into()).unwrap();
449        uring_ex.run_until(go(source)).unwrap();
450
451        state.lock().wake();
452        handle.join().unwrap().unwrap();
453    }
454}