cros_async/sys/linux/
tokio_source.rs

1// Copyright 2024 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::io;
6use std::os::fd::AsRawFd;
7use std::os::fd::OwnedFd;
8use std::os::fd::RawFd;
9use std::sync::Arc;
10
11use base::add_fd_flags;
12use base::clone_descriptor;
13use base::linux::fallocate;
14use base::linux::FallocateMode;
15use base::AsRawDescriptor;
16use base::VolatileSlice;
17use tokio::io::unix::AsyncFd;
18
19use crate::mem::MemRegion;
20use crate::AsyncError;
21use crate::AsyncResult;
22use crate::BackingMemory;
23use crate::IoOptions;
24
25#[derive(Debug, thiserror::Error)]
26pub enum Error {
27    #[error("Failed to copy the FD for the polling context: '{0}'")]
28    DuplicatingFd(base::Error),
29    #[error("Failed to punch hole in file: '{0}'.")]
30    Fallocate(base::Error),
31    #[error("Failed to fdatasync: '{0}'")]
32    Fdatasync(io::Error),
33    #[error("Failed to fsync: '{0}'")]
34    Fsync(io::Error),
35    #[error("Failed to join task: '{0}'")]
36    Join(tokio::task::JoinError),
37    #[error("Cannot wait on file descriptor")]
38    NonWaitable,
39    #[error("Failed to read: '{0}'")]
40    Read(io::Error),
41    #[error("Failed to set nonblocking: '{0}'")]
42    SettingNonBlocking(base::Error),
43    #[error("Tokio Async FD error: '{0}'")]
44    TokioAsyncFd(io::Error),
45    #[error("Failed to write: '{0}'")]
46    Write(io::Error),
47}
48
49impl From<Error> for io::Error {
50    fn from(e: Error) -> Self {
51        use Error::*;
52        match e {
53            DuplicatingFd(e) => e.into(),
54            Fallocate(e) => e.into(),
55            Fdatasync(e) => e,
56            Fsync(e) => e,
57            Join(e) => io::Error::other(e),
58            NonWaitable => io::Error::other(e),
59            Read(e) => e,
60            SettingNonBlocking(e) => e.into(),
61            TokioAsyncFd(e) => e,
62            Write(e) => e,
63        }
64    }
65}
66
67enum FdType {
68    Async(AsyncFd<Arc<OwnedFd>>),
69    Blocking(Arc<OwnedFd>),
70}
71
72impl AsRawFd for FdType {
73    fn as_raw_fd(&self) -> RawFd {
74        match self {
75            FdType::Async(async_fd) => async_fd.as_raw_fd(),
76            FdType::Blocking(blocking) => blocking.as_raw_fd(),
77        }
78    }
79}
80
81impl From<Error> for AsyncError {
82    fn from(e: Error) -> AsyncError {
83        AsyncError::SysVariants(e.into())
84    }
85}
86
87fn do_fdatasync(raw: Arc<OwnedFd>) -> io::Result<()> {
88    let fd = raw.as_raw_fd();
89    // SAFETY: we partially own `raw`
90    match unsafe { libc::fdatasync(fd) } {
91        0 => Ok(()),
92        _ => Err(io::Error::last_os_error()),
93    }
94}
95
96fn do_fsync(raw: Arc<OwnedFd>) -> io::Result<()> {
97    let fd = raw.as_raw_fd();
98    // SAFETY: we partially own `raw`
99    match unsafe { libc::fsync(fd) } {
100        0 => Ok(()),
101        _ => Err(io::Error::last_os_error()),
102    }
103}
104
105fn do_read_vectored(
106    raw: Arc<OwnedFd>,
107    file_offset: Option<u64>,
108    io_vecs: &[VolatileSlice],
109) -> io::Result<usize> {
110    let ptr = io_vecs.as_ptr() as *const libc::iovec;
111    let len = io_vecs.len() as i32;
112    let fd = raw.as_raw_fd();
113    let res = match file_offset {
114        // SAFETY: we partially own `raw`, `io_vecs` is validated
115        Some(off) => unsafe { libc::preadv64(fd, ptr, len, off as libc::off64_t) },
116        // SAFETY: we partially own `raw`, `io_vecs` is validated
117        None => unsafe { libc::readv(fd, ptr, len) },
118    };
119    match res {
120        r if r >= 0 => Ok(res as usize),
121        _ => Err(io::Error::last_os_error()),
122    }
123}
124fn do_read(raw: Arc<OwnedFd>, file_offset: Option<u64>, buf: &mut [u8]) -> io::Result<usize> {
125    let fd = raw.as_raw_fd();
126    let ptr = buf.as_mut_ptr() as *mut libc::c_void;
127    let res = match file_offset {
128        // SAFETY: we partially own `raw`, `ptr` has space up to vec.len()
129        Some(off) => unsafe { libc::pread64(fd, ptr, buf.len(), off as libc::off64_t) },
130        // SAFETY: we partially own `raw`, `ptr` has space up to vec.len()
131        None => unsafe { libc::read(fd, ptr, buf.len()) },
132    };
133    match res {
134        r if r >= 0 => Ok(res as usize),
135        _ => Err(io::Error::last_os_error()),
136    }
137}
138
139fn do_write(raw: Arc<OwnedFd>, file_offset: Option<u64>, buf: &[u8]) -> io::Result<usize> {
140    let fd = raw.as_raw_fd();
141    let ptr = buf.as_ptr() as *const libc::c_void;
142    let res = match file_offset {
143        // SAFETY: we partially own `raw`, `ptr` has data up to vec.len()
144        Some(off) => unsafe { libc::pwrite64(fd, ptr, buf.len(), off as libc::off64_t) },
145        // SAFETY: we partially own `raw`, `ptr` has data up to vec.len()
146        None => unsafe { libc::write(fd, ptr, buf.len()) },
147    };
148    match res {
149        r if r >= 0 => Ok(res as usize),
150        _ => Err(io::Error::last_os_error()),
151    }
152}
153
154fn do_write_vectored(
155    raw: Arc<OwnedFd>,
156    file_offset: Option<u64>,
157    io_vecs: &[VolatileSlice],
158) -> io::Result<usize> {
159    let ptr = io_vecs.as_ptr() as *const libc::iovec;
160    let len = io_vecs.len() as i32;
161    let fd = raw.as_raw_fd();
162    let res = match file_offset {
163        // SAFETY: we partially own `raw`, `io_vecs` is validated
164        Some(off) => unsafe { libc::pwritev64(fd, ptr, len, off as libc::off64_t) },
165        // SAFETY: we partially own `raw`, `io_vecs` is validated
166        None => unsafe { libc::writev(fd, ptr, len) },
167    };
168    match res {
169        r if r >= 0 => Ok(res as usize),
170        _ => Err(io::Error::last_os_error()),
171    }
172}
173
174pub struct TokioSource<T> {
175    fd: FdType,
176    inner: T,
177    runtime: tokio::runtime::Handle,
178}
179impl<T: AsRawDescriptor> TokioSource<T> {
180    pub fn new(inner: T, runtime: tokio::runtime::Handle) -> Result<TokioSource<T>, Error> {
181        let _guard = runtime.enter(); // Required for AsyncFd
182        let safe_fd = clone_descriptor(&inner).map_err(Error::DuplicatingFd)?;
183        let fd_arc: Arc<OwnedFd> = Arc::new(safe_fd.into());
184        let fd = match AsyncFd::new(fd_arc.clone()) {
185            Ok(async_fd) => {
186                add_fd_flags(async_fd.get_ref().as_raw_descriptor(), libc::O_NONBLOCK)
187                    .map_err(Error::SettingNonBlocking)?;
188                FdType::Async(async_fd)
189            }
190            Err(e) if e.kind() == io::ErrorKind::PermissionDenied => FdType::Blocking(fd_arc),
191            Err(e) => return Err(Error::TokioAsyncFd(e)),
192        };
193        Ok(TokioSource { fd, inner, runtime })
194    }
195
196    pub fn as_source(&self) -> &T {
197        &self.inner
198    }
199
200    pub fn as_source_mut(&mut self) -> &mut T {
201        &mut self.inner
202    }
203
204    fn clone_fd(&self) -> Arc<OwnedFd> {
205        match &self.fd {
206            FdType::Async(async_fd) => async_fd.get_ref().clone(),
207            FdType::Blocking(blocking) => blocking.clone(),
208        }
209    }
210
211    pub async fn fdatasync(&self) -> AsyncResult<()> {
212        let fd = self.clone_fd();
213        Ok(self
214            .runtime
215            .spawn_blocking(move || do_fdatasync(fd))
216            .await
217            .map_err(Error::Join)?
218            .map_err(Error::Fdatasync)?)
219    }
220
221    pub async fn fsync(&self) -> AsyncResult<()> {
222        let fd = self.clone_fd();
223        Ok(self
224            .runtime
225            .spawn_blocking(move || do_fsync(fd))
226            .await
227            .map_err(Error::Join)?
228            .map_err(Error::Fsync)?)
229    }
230
231    pub fn into_source(self) -> T {
232        self.inner
233    }
234
235    pub async fn read_to_vec(
236        &self,
237        file_offset: Option<u64>,
238        mut vec: Vec<u8>,
239        _options: IoOptions,
240    ) -> AsyncResult<(usize, Vec<u8>)> {
241        Ok(match &self.fd {
242            FdType::Async(async_fd) => {
243                let res = async_fd
244                    .async_io(tokio::io::Interest::READABLE, |fd| {
245                        do_read(fd.clone(), file_offset, &mut vec)
246                    })
247                    .await
248                    .map_err(AsyncError::Io)?;
249                (res, vec)
250            }
251            FdType::Blocking(blocking) => {
252                let fd = blocking.clone();
253                self.runtime
254                    .spawn_blocking(move || {
255                        let size = do_read(fd, file_offset, &mut vec)?;
256                        Ok((size, vec))
257                    })
258                    .await
259                    .map_err(Error::Join)?
260                    .map_err(Error::Read)?
261            }
262        })
263    }
264
265    pub async fn read_to_mem(
266        &self,
267        file_offset: Option<u64>,
268        mem: Arc<dyn BackingMemory + Send + Sync>,
269        mem_offsets: impl IntoIterator<Item = MemRegion>,
270        _options: IoOptions,
271    ) -> AsyncResult<usize> {
272        let mem_offsets_vec: Vec<MemRegion> = mem_offsets.into_iter().collect();
273        Ok(match &self.fd {
274            FdType::Async(async_fd) => {
275                let iovecs = mem_offsets_vec
276                    .into_iter()
277                    .filter_map(|mem_range| mem.get_volatile_slice(mem_range).ok())
278                    .collect::<Vec<VolatileSlice>>();
279                async_fd
280                    .async_io(tokio::io::Interest::READABLE, |fd| {
281                        do_read_vectored(fd.clone(), file_offset, &iovecs)
282                    })
283                    .await
284                    .map_err(AsyncError::Io)?
285            }
286            FdType::Blocking(blocking) => {
287                let fd = blocking.clone();
288                self.runtime
289                    .spawn_blocking(move || {
290                        let iovecs = mem_offsets_vec
291                            .into_iter()
292                            .filter_map(|mem_range| mem.get_volatile_slice(mem_range).ok())
293                            .collect::<Vec<VolatileSlice>>();
294                        do_read_vectored(fd, file_offset, &iovecs)
295                    })
296                    .await
297                    .map_err(Error::Join)?
298                    .map_err(Error::Read)?
299            }
300        })
301    }
302
303    pub async fn punch_hole(&self, file_offset: u64, len: u64) -> AsyncResult<()> {
304        let fd = self.clone_fd();
305        Ok(self
306            .runtime
307            .spawn_blocking(move || fallocate(&*fd, FallocateMode::PunchHole, file_offset, len))
308            .await
309            .map_err(Error::Join)?
310            .map_err(Error::Fallocate)?)
311    }
312
313    pub async fn wait_readable(&self) -> AsyncResult<()> {
314        match &self.fd {
315            FdType::Async(async_fd) => async_fd
316                .readable()
317                .await
318                .map_err(crate::AsyncError::Io)?
319                .retain_ready(),
320            FdType::Blocking(_) => return Err(Error::NonWaitable.into()),
321        }
322        Ok(())
323    }
324
325    pub async fn write_from_mem(
326        &self,
327        file_offset: Option<u64>,
328        mem: Arc<dyn BackingMemory + Send + Sync>,
329        mem_offsets: impl IntoIterator<Item = MemRegion>,
330        _options: IoOptions,
331    ) -> AsyncResult<usize> {
332        let mem_offsets_vec: Vec<MemRegion> = mem_offsets.into_iter().collect();
333        Ok(match &self.fd {
334            FdType::Async(async_fd) => {
335                let iovecs = mem_offsets_vec
336                    .into_iter()
337                    .filter_map(|mem_range| mem.get_volatile_slice(mem_range).ok())
338                    .collect::<Vec<VolatileSlice>>();
339                async_fd
340                    .async_io(tokio::io::Interest::WRITABLE, |fd| {
341                        do_write_vectored(fd.clone(), file_offset, &iovecs)
342                    })
343                    .await
344                    .map_err(AsyncError::Io)?
345            }
346            FdType::Blocking(blocking) => {
347                let fd = blocking.clone();
348                self.runtime
349                    .spawn_blocking(move || {
350                        let iovecs = mem_offsets_vec
351                            .into_iter()
352                            .filter_map(|mem_range| mem.get_volatile_slice(mem_range).ok())
353                            .collect::<Vec<VolatileSlice>>();
354                        do_write_vectored(fd, file_offset, &iovecs)
355                    })
356                    .await
357                    .map_err(Error::Join)?
358                    .map_err(Error::Read)?
359            }
360        })
361    }
362
363    pub async fn write_from_vec(
364        &self,
365        file_offset: Option<u64>,
366        vec: Vec<u8>,
367        _options: IoOptions,
368    ) -> AsyncResult<(usize, Vec<u8>)> {
369        Ok(match &self.fd {
370            FdType::Async(async_fd) => {
371                let res = async_fd
372                    .async_io(tokio::io::Interest::WRITABLE, |fd| {
373                        do_write(fd.clone(), file_offset, &vec)
374                    })
375                    .await
376                    .map_err(AsyncError::Io)?;
377                (res, vec)
378            }
379            FdType::Blocking(blocking) => {
380                let fd = blocking.clone();
381                self.runtime
382                    .spawn_blocking(move || {
383                        let size = do_write(fd.clone(), file_offset, &vec)?;
384                        Ok((size, vec))
385                    })
386                    .await
387                    .map_err(Error::Join)?
388                    .map_err(Error::Read)?
389            }
390        })
391    }
392
393    pub async fn write_zeroes_at(&self, file_offset: u64, len: u64) -> AsyncResult<()> {
394        let fd = self.clone_fd();
395        Ok(self
396            .runtime
397            .spawn_blocking(move || fallocate(&*fd, FallocateMode::ZeroRange, file_offset, len))
398            .await
399            .map_err(Error::Join)?
400            .map_err(Error::Fallocate)?)
401    }
402}
cros_async/sys/linux/tokio_source.rs

cros_async/sys/linux/
tokio_source.rs
