restructure code

Author: mox692

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,
@@ -89,6 +91,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

@@ -42,8 +42,7 @@ pub async fn write(path: impl AsRef<Path>, contents: impl AsRef<[u8]>) -> io::Re
 
 #[cfg(all(tokio_uring, feature = "rt", feature = "fs", target_os = "linux"))]
 async fn write_uring(path: &Path, contents: OwnedBuf) -> io::Result<()> {
-    use crate::{fs::OpenOptions, runtime::driver::op::Op};
-    use std::os::fd::AsFd;
+    use crate::{fs::OpenOptions, io::uring::utils::SharedFd, runtime::driver::op::Op};
 
     let file = OpenOptions::new()
         .write(true)
@@ -52,15 +51,16 @@ async fn write_uring(path: &Path, contents: OwnedBuf) -> io::Result<()> {
         .open(path)
         .await?;
 
+    let fd = SharedFd::try_from(file).expect("unexpected in-flight operation detected");
+
     let mut pos = 0;
-    let fd = file.as_fd();
-    let mut buf = contents.as_ref();
+    let mut buf = contents;
     while !buf.is_empty() {
-        let n = Op::write_at(fd, buf, pos)?.await? as usize;
+        let n = Op::write_at(fd.clone(), buf.clone(), pos)?.await? as usize;
         if n == 0 {
             return Err(io::ErrorKind::WriteZero.into());
         }
-        buf = &buf[n..];
+        buf.advance(n);
         pos += n as u64;
     }
 

tokio/src/io/uring/utils.rs

@@ -1,6 +1,35 @@
+use std::os::fd::{AsRawFd, OwnedFd, RawFd};
 use std::os::unix::ffi::OsStrExt;
+use std::sync::Arc;
 use std::{ffi::CString, io, path::Path};
 
 pub(crate) fn cstr(p: &Path) -> io::Result<CString> {
     Ok(CString::new(p.as_os_str().as_bytes())?)
 }
+
+/// A reference-counted file descriptor.
+#[derive(Clone, Debug)]
+pub(crate) struct SharedFd {
+    inner: Arc<OwnedFd>,
+}
+
+impl SharedFd {
+    pub(crate) fn new(fd: OwnedFd) -> SharedFd {
+        SharedFd {
+            inner: Arc::new(fd),
+        }
+    }
+}
+
+impl AsRawFd for SharedFd {
+    fn as_raw_fd(&self) -> RawFd {
+        self.inner.as_raw_fd()
+    }
+}
+
+impl TryFrom<crate::fs::File> for SharedFd {
+    type Error = crate::fs::File;
+    fn try_from(value: crate::fs::File) -> Result<Self, Self::Error> {
+        Ok(SharedFd::new(OwnedFd::from(value.try_into_std()?)))
+    }
+}

tokio/src/io/uring/write.rs

@@ -1,12 +1,16 @@
-use crate::runtime::driver::op::{CancelData, Cancellable, Completable, CqeResult, Op};
-use io_uring::{opcode, types};
-use std::{
-    cmp, io,
-    os::fd::{AsRawFd, BorrowedFd},
+use crate::{
+    io::uring::utils::SharedFd,
+    runtime::driver::op::{CancelData, Cancellable, Completable, CqeResult, Op},
+    util::as_ref::OwnedBuf,
 };
+use io_uring::{opcode, types};
+use std::{io, os::fd::AsRawFd};
 
 #[derive(Debug)]
-pub(crate) struct Write;
+pub(crate) struct Write {
+    _fd: SharedFd,
+    _buf: OwnedBuf,
+}
 
 impl Completable for Write {
     // The number of bytes written.
@@ -24,17 +28,18 @@ impl Cancellable for Write {
 }
 
 impl Op<Write> {
-    pub(crate) fn write_at(fd: BorrowedFd<'_>, buf: &[u8], offset: u64) -> io::Result<Self> {
+    pub(crate) fn write_at(fd: SharedFd, buf: OwnedBuf, 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 = cmp::min(buf.len(), u32::MAX as usize) as u32;
+        let len: u32 = std::cmp::min(buf.len(), u32::MAX as usize) as u32;
 
-        let sqe = opcode::Write::new(types::Fd(fd.as_raw_fd()), buf.as_ptr(), len)
+        let sqe = opcode::Write::new(types::Fd(fd.as_raw_fd()), buf.as_ref().as_ptr(), len)
             .offset(offset)
             .build();
 
-        // SAFETY: TODO
-        let op = unsafe { Op::new(sqe, Write) };
+        // SAFETY: `fd` and `buf` are owned by this `Op`, ensuring that these params are valid
+        // until `Op::drop` gets called.
+        let op = unsafe { Op::new(sqe, Write { _buf: buf, _fd: fd }) };
         Ok(op)
     }
 }

tokio/src/runtime/driver/op.rs

@@ -17,7 +17,11 @@ pub(crate) enum CancelData {
         // so `#[allow(dead_code)]` is needed.
         #[allow(dead_code)] Open,
     ),
-    Write(Write),
+    Write(
+        // This field isn't accessed directly, but it holds cancellation data,
+        // so `#[allow(dead_code)]` is needed.
+        #[allow(dead_code)] Write,
+    ),
 }
 
 #[derive(Debug)]

tokio/src/util/as_ref.rs

@@ -1,16 +1,69 @@
+use std::sync::Arc;
+
+#[cfg(feature = "io-util")]
+use bytes::Buf;
+
 use super::typeid;
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub(crate) enum OwnedBuf {
-    Vec(Vec<u8>),
+    Vec(VecBuf),
     #[cfg(feature = "io-util")]
     Bytes(bytes::Bytes),
 }
 
+#[allow(unused)]
+impl OwnedBuf {
+    pub(crate) fn len(&self) -> usize {
+        match self {
+            OwnedBuf::Vec(vec) => vec.len(),
+            #[cfg(feature = "io-util")]
+            OwnedBuf::Bytes(b) => b.len(),
+        }
+    }
+
+    pub(crate) fn advance(&mut self, n: usize) {
+        match self {
+            OwnedBuf::Vec(vec) => vec.advance(n),
+            #[cfg(feature = "io-util")]
+            OwnedBuf::Bytes(b) => b.advance(n),
+        }
+    }
+
+    pub(crate) fn is_empty(&self) -> bool {
+        match self {
+            Self::Vec(vec) => vec.is_empty(),
+            #[cfg(feature = "io-util")]
+            Self::Bytes(bytes) => bytes.is_empty(),
+        }
+    }
+}
+
+/// A Vec-like, reference-counted buffer.
+#[derive(Debug, Clone)]
+pub(crate) struct VecBuf {
+    data: Arc<[u8]>,
+    pos: usize,
+}
+
+impl VecBuf {
+    fn len(&self) -> usize {
+        self.data.len().saturating_sub(self.pos)
+    }
+
+    fn advance(&mut self, n: usize) {
+        self.pos = self.pos.saturating_add(n);
+    }
+
+    fn is_empty(&self) -> bool {
+        self.data.len().saturating_sub(self.pos) == 0
+    }
+}
+
 impl AsRef<[u8]> for OwnedBuf {
     fn as_ref(&self) -> &[u8] {
         match self {
-            Self::Vec(vec) => vec,
+            Self::Vec(VecBuf { data, pos }) => &data[*pos..],
             #[cfg(feature = "io-util")]
             Self::Bytes(bytes) => bytes,
         }
@@ -19,12 +72,22 @@ impl AsRef<[u8]> for OwnedBuf {
 
 pub(crate) fn upgrade<B: AsRef<[u8]>>(buf: B) -> OwnedBuf {
     let buf = match unsafe { typeid::try_transmute::<B, Vec<u8>>(buf) } {
-        Ok(vec) => return OwnedBuf::Vec(vec),
+        Ok(vec) => {
+            return OwnedBuf::Vec(VecBuf {
+                data: vec.into(),
+                pos: 0,
+            });
+        }
         Err(original_buf) => original_buf,
     };
 
     let buf = match unsafe { typeid::try_transmute::<B, String>(buf) } {
-        Ok(string) => return OwnedBuf::Vec(string.into_bytes()),
+        Ok(string) => {
+            return OwnedBuf::Vec(VecBuf {
+                data: string.into_bytes().into(),
+                pos: 0,
+            });
+        }
         Err(original_buf) => original_buf,
     };
 
@@ -34,5 +97,8 @@ pub(crate) fn upgrade<B: AsRef<[u8]>>(buf: B) -> OwnedBuf {
         Err(original_buf) => original_buf,
     };
 
-    OwnedBuf::Vec(buf.as_ref().to_owned())
+    OwnedBuf::Vec(VecBuf {
+        data: buf.as_ref().to_owned().into(),
+        pos: 0,
+    })
 }