armdb 0.2.0

sharded bitcask key-value storage optimized for NVMe
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109

#[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>,
}

#[cfg(target_os = "linux")]
impl UringWriter {
    pub fn new() -> DbResult<Self> {
        let ring = IoUring::builder()
            .setup_sqpoll(2000)
            // TODO:
            // .setup_sqpoll_cpu(cpu_id)
            .build(256)
            .map_err(|e| DbError::Io(e.into()))?;

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

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

    /// 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",
                )));
            }
            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(())
    }
}