armdb 0.5.3

#[cfg(target_os = "linux")]
use rustix_uring::{IoUring, cqueue, opcode, squeue, types};

#[cfg(target_os = "linux")]
use crate::error::{DbError, DbResult};

/// io_uring-based writer. Submits writes and fsyncs via the kernel's submission queue.
/// No internal buffering — callers provide the data buffer directly.
#[cfg(target_os = "linux")]
pub struct UringWriter {
    ring: IoUring<squeue::Entry, cqueue::Entry>,
    fd: Option<std::os::unix::io::RawFd>,
    direct: bool,
}

#[cfg(target_os = "linux")]
impl UringWriter {
    pub fn new(sqpoll_idle_ms: Option<u32>) -> DbResult<Self> {
        let mut builder = IoUring::builder();
        if let Some(idle) = sqpoll_idle_ms {
            builder.setup_sqpoll(idle);
            // TODO: builder.setup_sqpoll_cpu(cpu_id) — pin the poller (backlog).
        }
        let ring = builder.build(256).map_err(|e| DbError::Io(e.into()))?;

        Ok(Self {
            ring,
            fd: None,
            direct: false,
        })
    }

    pub fn set_file(&mut self, fd: std::os::unix::io::RawFd) {
        self.fd = Some(fd);
    }

    pub fn set_direct(&mut self, direct: bool) {
        self.direct = direct;
    }

    /// Submit a single write opcode and return the byte count reported by the CQE.
    fn submit_write(&mut self, data: &[u8], offset: u64) -> DbResult<u32> {
        let fd = self.fd.expect("No file set in UringWriter");

        let write_e = opcode::Write::new(types::Fd(fd), data.as_ptr(), data.len() as u32)
            .offset(offset)
            .build()
            .user_data(1);

        unsafe {
            if self.ring.submission().push(&write_e).is_err() {
                self.ring.submit().map_err(|e| DbError::Io(e.into()))?;
                self.ring.submission().push(&write_e).expect("Queue full");
            }
        }

        self.ring
            .submit_and_wait(1)
            .map_err(|e| DbError::Io(e.into()))?;
        let mut cq = self.ring.completion();
        let mut written = 0u32;
        for cqe in &mut cq {
            written = cqe.result().map_err(|e| DbError::Io(e.into()))?;
        }
        Ok(written)
    }

    /// Write data to the file at the given offset. Blocks until completion.
    /// Loops on short successful writes; returns `WriteZero` if a write makes
    /// no progress.
    pub fn write_at(&mut self, data: &[u8], file_offset: u64) -> DbResult<()> {
        debug_assert!(
            data.len() <= u32::MAX as usize,
            "write_at data exceeds u32::MAX"
        );
        if data.is_empty() {
            return Ok(());
        }
        let mut written = 0usize;
        while written < data.len() {
            let remaining = &data[written..];
            let n = self.submit_write(remaining, file_offset + written as u64)?;
            if n == 0 {
                return Err(DbError::Io(std::io::Error::new(
                    std::io::ErrorKind::WriteZero,
                    "uring write returned zero bytes",
                )));
            }
            if self.direct && !(n as usize).is_multiple_of(4096) {
                return Err(DbError::Io(std::io::Error::new(
                    std::io::ErrorKind::InvalidInput,
                    "direct mode requires 4096-byte aligned write length",
                )));
            }
            written += n as usize;
        }
        Ok(())
    }

    pub fn fsync(&mut self) -> DbResult<()> {
        let fd = self.fd.expect("No file set in UringWriter");

        let fsync_e = opcode::Fsync::new(types::Fd(fd)).build().user_data(2);

        unsafe {
            if self.ring.submission().push(&fsync_e).is_err() {
                self.ring.submit().map_err(|e| DbError::Io(e.into()))?;
                self.ring.submission().push(&fsync_e).expect("Queue full");
            }
        }

        self.ring
            .submit_and_wait(1)
            .map_err(|e| DbError::Io(e.into()))?;
        let mut cq = self.ring.completion();
        for cqe in &mut cq {
            if let Err(err) = cqe.result() {
                return Err(DbError::Io(err.into()));
            }
        }

        Ok(())
    }
}

#[cfg(all(target_os = "linux", test))]
mod tests {
    use std::os::unix::io::AsRawFd;

    use crate::error::DbError;
    use crate::io::aligned_buf::AlignedBuf;

    use super::UringWriter;

    #[test]
    fn direct_writer_rejects_unaligned_short_write() {
        let dir = std::env::temp_dir().join(format!("armdb_uring_direct_{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("test.data");

        let file = crate::io::direct::open_serving_write(&path, false).unwrap();

        let mut writer = UringWriter::new(Some(2000)).unwrap();
        writer.set_file(file.as_raw_fd());
        writer.set_direct(true);

        let page = AlignedBuf::zeroed(4096);
        writer.write_at(&page, 0).unwrap();

        let err = writer.write_at(b"hello", 4096).unwrap_err();
        assert!(matches!(err, DbError::Io(_)));

        std::fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn builds_without_sqpoll_and_writes() {
        let dir = std::env::temp_dir().join(format!("armdb_uring_nosqp_{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("nosqp.data");
        let file = crate::io::direct::open_serving_write(&path, false).unwrap();

        let mut writer = UringWriter::new(None).unwrap(); // io_uring, no poller
        writer.set_file(file.as_raw_fd());
        writer.write_at(b"hello world", 0).unwrap();

        std::fs::remove_dir_all(&dir).ok();
    }

    #[test]
    fn builds_with_short_idle() {
        let _w = UringWriter::new(Some(20)).unwrap();
    }
}