ronantakizawa/codereview-bench huggingface dataset:
Repository: tokio-rs/tokio
File: tokio/src/sync/notify.rs (Rust)

Code before:
            // be unlinked from the list in `NotifyWaitersList` destructor.
            wakers.wake_all();

            // Acquire the lock again.
            waiters = self.waiters.lock();
        }

        // Release the lock before notifying
        drop(waiters);

        wakers.wake_all();
    }

    pub(crate) fn lock(&self) -> NotifyGuard<'_> {
        let guarded_waiters = self.waiters.lock();

        // The state must be loaded while the lock is held. The state may only
        // transition out of WAITING while the lock is held.
        let current_state = self.state.load(SeqCst);

        NotifyGuard {
            guarded_notify: self,
            guarded_waiters,
            current_state,
        }
    }
}

impl Default for Notify {
    fn default() -> Notify {
        Notify::new()
    }
}

impl UnwindSafe for Notify {}
impl RefUnwindSafe for Notify {}

fn notify_locked(
    waiters: &mut WaitList,
    state: &AtomicUsize,
    curr: usize,
    strategy: NotifyOneStrategy,
) -> Option<Waker> {
    match get_state(curr) {
        EMPTY | NOTIFIED => {
            let res = state.compare_exchange(curr, set_state(curr, NOTIFIED), SeqCst, SeqCst);

            match res {
                Ok(_) => None,
                Err(actual) => {
                    let actual_state = get_state(actual);

Review comment (suggestion):
Even though this is just an internal API, I still feel that adding a `lock` method to `Notify` would confuse peoples who are developing the tokio itself.

Code after:
            // be unlinked from the list in `NotifyWaitersList` destructor.
            wakers.wake_all();

            // Acquire the lock again.
            waiters = self.waiters.lock();
        }

        // Release the lock before notifying
        drop(waiters);

        wakers.wake_all();
    }

    pub(crate) fn lock_waiter_list(&self) -> NotifyGuard<'_> {
        let guarded_waiters = self.waiters.lock();

        // The state must be loaded while the lock is held. The state may only
        // transition out of WAITING while the lock is held.
        let current_state = self.state.load(SeqCst);

        NotifyGuard {
            guarded_notify: self,
            guarded_waiters,
            current_state,
        }
    }
}

impl Default for Notify {
    fn default() -> Notify {
        Notify::new()
    }
}

impl UnwindSafe for Notify {}
impl RefUnwindSafe for Notify {}

fn notify_locked(
    waiters: &mut WaitList,
    state: &AtomicUsize,
    curr: usize,
    strategy: NotifyOneStrategy,
) -> Option<Waker> {
    match get_state(curr) {
        EMPTY | NOTIFIED => {
            let res = state.compare_exchange(curr, set_state(curr, NOTIFIED), SeqCst, SeqCst);

            match res {
                Ok(_) => None,
                Err(actual) => {
                    let actual_state = get_state(actual);
Repository: tokio-rs/tokio
File: tokio/src/fs/write.rs (Rust)

Code before:
///
/// # Examples
///
/// ```no_run
/// use tokio::fs;
///
/// # async fn dox() -> std::io::Result<()> {
/// fs::write("foo.txt", b"Hello world!").await?;
/// # Ok(())
/// # }
/// ```
pub async fn write(path: impl AsRef<Path>, contents: impl AsRef<[u8]>) -> io::Result<()> {
    #[cfg(all(tokio_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.as_ref()).await;
        }
    }

    write_spawn_blocking(path, contents).await
}

#[cfg(all(tokio_uring, feature = "rt", feature = "fs", target_os = "linux"))]
async fn write_uring(path: impl AsRef<Path>, mut buf: &[u8]) -> io::Result<()> {
    use crate::{fs::OpenOptions, runtime::driver::op::Op};
    use std::os::fd::AsFd;

    let file = OpenOptions::new()
        .write(true)
        .create(true)
        .truncate(true)
        .open(path.as_ref())
        .await?;

    let mut pos = 0;
    let fd = file.as_fd();
    while !buf.is_empty() {
        let n = Op::write_at(fd, buf, pos)?.await? as usize;
        if n == 0 {
            return Err(io::ErrorKind::WriteZero.into());
        }
        buf = &buf[n..];
        pos += n as u64;
    }

    Ok(())
}

async fn write_spawn_blocking<P, C>(path: P, contents: C) -> io::Result<()>

Review comment (suggestion):
I would change both methods to take `&Path` and `crate::util::as_ref::OwnedBuf`, to avoid generics in the actual implementations.

Code after:
///
/// # Examples
///
/// ```no_run
/// use tokio::fs;
///
/// # async fn dox() -> std::io::Result<()> {
/// fs::write("foo.txt", b"Hello world!").await?;
/// # Ok(())
/// # }
/// ```
pub async fn write(path: impl AsRef<Path>, contents: impl AsRef<[u8]>) -> io::Result<()> {
    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?;
Repository: tokio-rs/tokio
File: tokio/src/io/uring/write.rs (Rust)

Code before:

    fn complete(self, cqe: CqeResult) -> io::Result<Self::Output> {
        cqe.result
    }
}

impl Cancellable for Write {
    fn cancel(self) -> CancelData {
        CancelData::Write(self)
    }
}

impl Op<Write> {
    pub(crate) fn write_at(fd: BorrowedFd<'_>, buf: &[u8], 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 sqe = opcode::Write::new(types::Fd(fd.as_raw_fd()), buf.as_ptr(), len)
            .offset(offset)
            .build();

        // SAFETY: `fd` and `buf` is valid until this future completes.
        // If this operation is cancelled, drop of `Op` occurs before dropping
        // `fd` or `buf`, ensuring the safety contract holds.
        let op = unsafe { Op::new(sqe, Write) };
        Ok(op)
    }
}


Review comment (suggestion):
This logic I don't buy. Using `mem::forget` on the future it's possible to cause UAF. Some form of ownership is required.

Code after:
    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)
    }
}
Repository: tokio-rs/tokio
File: tokio/src/io/uring/write.rs (Rust)

Code before:
use std::{io, os::fd::AsRawFd};

// This holds a strong ref to the resources such as `fd`, `buf`,
// preventing these resources from being dropped while the operation is in-flight.
#[derive(Debug)]
pub(crate) struct Write {
    _fd: SharedFd,
    _buf: OwnedBuf,
}

impl Completable for Write {
    // The number of bytes written.
    type Output = u32;
    fn complete(self, cqe: CqeResult) -> io::Result<Self::Output> {
        cqe.result
    }
}

impl Cancellable for Write {
    fn cancel(self) -> CancelData {
        CancelData::Write(self)
    }
}

impl Op<Write> {
    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 = std::cmp::min(buf.len(), u32::MAX as usize) as u32;

        let sqe = opcode::Write::new(types::Fd(fd.as_raw_fd()), buf.as_ref().as_ptr(), len)
            .offset(offset)
            .build();

        // 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)
    }
}


Review comment (suggestion):
I think you can avoid cloning `OwnedBuf` by returning it here.
```suggestion
    pub(crate) fn write_at(fd: SharedFd, buf: OwnedBuf, offset: u64) -> io::Result<(Self, OwnedBuf)> {
```

Code after:
    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)
    }
}
Repository: tokio-rs/tokio
File: tokio/src/fs/write.rs (Rust)

Code before:
    }

    write_spawn_blocking(path, contents).await
}

#[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, io::uring::utils::SharedFd, runtime::driver::op::Op};

    let file = OpenOptions::new()
        .write(true)
        .create(true)
        .truncate(true)
        .open(path)
        .await?;

    let fd = SharedFd::try_from(file).expect("unexpected in-flight operation detected");

    let mut pos = 0;
    let mut buf = contents;
    while !buf.is_empty() {
        let n = Op::write_at(fd.clone(), buf.clone(), pos)?.await? as usize;
        if n == 0 {
            return Err(io::ErrorKind::WriteZero.into());
        }
        buf.advance(n);
        pos += 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
}


Review comment (performance):
I think `advance` is pretty expensive for a `Vec<u8>`, so we probably want `write_at` to instead take an offset and have it add the offset to the pointer internally.

Code after:
        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(())
}
Repository: tokio-rs/tokio
File: tokio/src/io/uring/write.rs (Rust)

Code before:
    // The number of bytes written.
    type Output = u32;
    fn complete(self, cqe: CqeResult) -> io::Result<Self::Output> {
        cqe.result
    }
}

impl Cancellable for Write {
    fn cancel(self) -> CancelData {
        CancelData::Write(self)
    }
}

impl Op<Write> {
    pub(crate) fn write_at(fd: BorrowedFd<'_>, buf: &[u8], 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 sqe = opcode::Write::new(types::Fd(fd.as_raw_fd()), buf.as_ptr(), len)
            .offset(offset)
            .build();

        // SAFETY: `fd` and `buf` are owned by caller function, ensuring these params are
        // valid until operation completes.
        let op = unsafe { Op::new(sqe, Write) };
        Ok(op)
    }
}


Review comment (suggestion):
We can do it this way, but `write_at` will need to be unsafe if you have a requirement on the caller like this.

But I have another question: How does cancellation work here? If this io_uring operation is cancelled, then `buf` must not be freed until io_uring confirms that cancellation has happened.

Code after:
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)
    }
}
Repository: tokio-rs/tokio
File: tokio/src/fs/write.rs (Rust)

Code before:
    write_spawn_blocking(path, contents).await
}

#[cfg(all(tokio_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.len();
    let mut offset: usize = 0;
    while offset < total {
        // 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 = std::cmp::min(total - offset, u32::MAX as usize) as u32;

        let (n, _buf, _fd) = Op::write_at(fd, buf, offset, len, offset as u64)?.await?;
        if n == 0 {
            return Err(io::ErrorKind::WriteZero.into());
        }

        buf = _buf;
        fd = _fd;
        offset += n as usize;
    }

    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
}


Review comment (refactor):
Could also make sense to move this logic inside `Op::write_at` and have it accept usize. After all, it's indistinguishable to the caller from any other situation that writes the bytes partially. That way it's more natural for the caller as the caller expects usize for both input and output length.
```suggestion
        let len = u32::try_from(total - offset).unwrap_or(u32::MAX);
```

Code after:
        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<()> {
Repository: tokio-rs/tokio
File: tokio/src/fs/write.rs (Rust)

Code before:
        }
    }

    write_spawn_blocking(path, contents).await
}

#[cfg(all(tokio_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.len();
    let mut offset: usize = 0;
    while offset < total {
        // 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 = std::cmp::min(total - offset, u32::MAX as usize) as u32;

        let (n, _buf, _fd) = Op::write_at(fd, buf, offset, len, offset as u64)?.await?;
        if n == 0 {
            return Err(io::ErrorKind::WriteZero.into());
        }

        buf = _buf;
        fd = _fd;
        offset += n as usize;
    }

    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
}


Review comment (suggestion):
Since `write_at` accepts u64 for file positions, let's use that here too.

Code after:
    ))]
    {
        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(())