diff --git a/tokio/src/fs/mocks.rs b/tokio/src/fs/mocks.rs index 3a8ac2bdd..ae5d7e536 100644 --- a/tokio/src/fs/mocks.rs +++ b/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 for OwnedFd { + #[inline] + fn from(file: MockFile) -> OwnedFd { + unsafe { OwnedFd::from_raw_fd(file.as_raw_fd()) } + } +} + tokio_thread_local! { static QUEUE: RefCell>> = RefCell::new(VecDeque::new()) } diff --git a/tokio/src/fs/write.rs b/tokio/src/fs/write.rs index f5d18e843..543f97fd4 100644 --- a/tokio/src/fs/write.rs +++ b/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, 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 } diff --git a/tokio/src/io/uring/mod.rs b/tokio/src/io/uring/mod.rs index e5ac85af6..4899d0a4a 100644 --- a/tokio/src/io/uring/mod.rs +++ b/tokio/src/io/uring/mod.rs @@ -1,2 +1,3 @@ pub(crate) mod open; pub(crate) mod utils; +pub(crate) mod write; diff --git a/tokio/src/io/uring/write.rs b/tokio/src/io/uring/write.rs new file mode 100644 index 000000000..9332bd007 --- /dev/null +++ b/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 { + Ok((cqe.result?, self.buf, self.fd)) + } +} + +impl Cancellable for Write { + fn cancel(self) -> CancelData { + CancelData::Write(self) + } +} + +impl Op { + /// 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 { + // 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) + } +} diff --git a/tokio/src/runtime/driver/op.rs b/tokio/src/runtime/driver/op.rs index 94afe163a..413dd4a83 100644 --- a/tokio/src/runtime/driver/op.rs +++ b/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)] diff --git a/tokio/tests/fs_write.rs b/tokio/tests/fs_write.rs new file mode 100644 index 000000000..a125e0408 --- /dev/null +++ b/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!"); +}