fs: support io_uring in fs::write (#7567)

tokio/src/fs/mocks.rs

@@ -2,6 +2,8 @@
 use mockall::mock;
 
 use crate::sync::oneshot;
+#[cfg(all(test, unix))]
+use std::os::fd::{AsRawFd, FromRawFd, OwnedFd};
 use std::{
     cell::RefCell,
     collections::VecDeque,
@@ -96,6 +98,14 @@ impl Write for &'_ MockFile {
     }
 }
 
+#[cfg(all(test, unix))]
+impl From<MockFile> for OwnedFd {
+    #[inline]
+    fn from(file: MockFile) -> OwnedFd {
+        unsafe { OwnedFd::from_raw_fd(file.as_raw_fd()) }
+    }
+}
+
 tokio_thread_local! {
     static QUEUE: RefCell<VecDeque<Box<dyn FnOnce() + Send>>> = RefCell::new(VecDeque::new())
 }

tokio/src/fs/write.rs

@@ -1,4 +1,4 @@
-use crate::fs::asyncify;
+use crate::{fs::asyncify, util::as_ref::OwnedBuf};
 
 use std::{io, path::Path};
 
@@ -24,8 +24,69 @@ use std::{io, path::Path};
 /// # }
 /// ```
 pub async fn write(path: impl AsRef<Path>, contents: impl AsRef<[u8]>) -> io::Result<()> {
-    let path = path.as_ref().to_owned();
+    let path = path.as_ref();
     let contents = crate::util::as_ref::upgrade(contents);
 
+    #[cfg(all(
+        tokio_unstable,
+        feature = "io-uring",
+        feature = "rt",
+        feature = "fs",
+        target_os = "linux"
+    ))]
+    {
+        let handle = crate::runtime::Handle::current();
+        let driver_handle = handle.inner.driver().io();
+        if driver_handle.check_and_init()? {
+            return write_uring(path, contents).await;
+        }
+    }
+
+    write_spawn_blocking(path, contents).await
+}
+
+#[cfg(all(
+    tokio_unstable,
+    feature = "io-uring",
+    feature = "rt",
+    feature = "fs",
+    target_os = "linux"
+))]
+async fn write_uring(path: &Path, mut buf: OwnedBuf) -> io::Result<()> {
+    use crate::{fs::OpenOptions, runtime::driver::op::Op};
+    use std::os::fd::OwnedFd;
+
+    let file = OpenOptions::new()
+        .write(true)
+        .create(true)
+        .truncate(true)
+        .open(path)
+        .await?;
+
+    let mut fd: OwnedFd = file
+        .try_into_std()
+        .expect("unexpected in-flight operation detected")
+        .into();
+
+    let total: usize = buf.as_ref().len();
+    let mut buf_offset: usize = 0;
+    let mut file_offset: u64 = 0;
+    while buf_offset < total {
+        let (n, _buf, _fd) = Op::write_at(fd, buf, buf_offset, file_offset)?.await?;
+        if n == 0 {
+            return Err(io::ErrorKind::WriteZero.into());
+        }
+
+        buf = _buf;
+        fd = _fd;
+        buf_offset += n as usize;
+        file_offset += n as u64;
+    }
+
+    Ok(())
+}
+
+async fn write_spawn_blocking(path: &Path, contents: OwnedBuf) -> io::Result<()> {
+    let path = path.to_owned();
     asyncify(move || std::fs::write(path, contents)).await
 }

tokio/src/io/uring/mod.rs

@@ -1,2 +1,3 @@
 pub(crate) mod open;
 pub(crate) mod utils;
+pub(crate) mod write;

tokio/src/io/uring/write.rs

@@ -0,0 +1,54 @@
+use crate::{
+    runtime::driver::op::{CancelData, Cancellable, Completable, CqeResult, Op},
+    util::as_ref::OwnedBuf,
+};
+use io_uring::{opcode, types};
+use std::{
+    io,
+    os::fd::{AsRawFd, OwnedFd},
+};
+
+#[derive(Debug)]
+pub(crate) struct Write {
+    buf: OwnedBuf,
+    fd: OwnedFd,
+}
+
+impl Completable for Write {
+    type Output = (u32, OwnedBuf, OwnedFd);
+    fn complete(self, cqe: CqeResult) -> io::Result<Self::Output> {
+        Ok((cqe.result?, self.buf, self.fd))
+    }
+}
+
+impl Cancellable for Write {
+    fn cancel(self) -> CancelData {
+        CancelData::Write(self)
+    }
+}
+
+impl Op<Write> {
+    /// Issue a write that starts at `buf_offset` within `buf` and writes some bytes
+    /// into `file` at `file_offset`.
+    pub(crate) fn write_at(
+        fd: OwnedFd,
+        buf: OwnedBuf,
+        buf_offset: usize,
+        file_offset: u64,
+    ) -> io::Result<Self> {
+        // There is a cap on how many bytes we can write in a single uring write operation.
+        // ref: https://github.com/axboe/liburing/discussions/497
+        let len = u32::try_from(buf.as_ref().len() - buf_offset).unwrap_or(u32::MAX);
+
+        let ptr = buf.as_ref()[buf_offset..buf_offset + len as usize].as_ptr();
+
+        let sqe = opcode::Write::new(types::Fd(fd.as_raw_fd()), ptr, len)
+            .offset(file_offset)
+            .build();
+
+        // SAFETY: parameters of the entry, such as `fd` and `buf`, are valid
+        // until this operation completes.
+        let op = unsafe { Op::new(sqe, Write { buf, fd }) };
+        Ok(op)
+    }
+}

tokio/src/runtime/driver/op.rs

@@ -1,4 +1,5 @@
 use crate::io::uring::open::Open;
+use crate::io::uring::write::Write;
 use crate::runtime::Handle;
 use io_uring::cqueue;
 use io_uring::squeue::Entry;
@@ -9,13 +10,13 @@ use std::task::Poll;
 use std::task::Waker;
 use std::{io, mem};
 
+// This field isn't accessed directly, but it holds cancellation data,
+// so `#[allow(dead_code)]` is needed.
+#[allow(dead_code)]
 #[derive(Debug)]
 pub(crate) enum CancelData {
-    Open(
-        // This field isn't accessed directly, but it holds cancellation data,
-        // so `#[allow(dead_code)]` is needed.
-        #[allow(dead_code)] Open,
-    ),
+    Open(Open),
+    Write(Write),
 }
 
 #[derive(Debug)]

tokio/tests/fs_write.rs

@@ -0,0 +1,16 @@
+#![warn(rust_2018_idioms)]
+#![cfg(all(feature = "full", not(target_os = "wasi")))] // WASI does not support all fs operations
+
+use tempfile::tempdir;
+use tokio::fs;
+
+#[tokio::test]
+async fn write() {
+    let dir = tempdir().unwrap();
+    let path = dir.path().join("test.txt");
+
+    fs::write(&path, "Hello, World!").await.unwrap();
+
+    let contents = fs::read_to_string(&path).await.unwrap();
+    assert_eq!(contents, "Hello, World!");
+}