turbokv 0.4.0

A fast, embedded key-value store with BTreeMap-like API.
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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163

use crate::core::error::Result;
use byteorder::{LittleEndian, WriteBytesExt};
use bytes::{BufMut, Bytes, BytesMut};

pub struct BlockBuilder {
    buffer: BytesMut,
    offsets: Vec<u32>,
    last_key: Option<Bytes>,
    max_size: usize,
}

impl BlockBuilder {
    pub fn new(max_size: usize) -> Self {
        Self {
            buffer: BytesMut::with_capacity(max_size),
            offsets: Vec::new(),
            last_key: None,
            max_size,
        }
    }

    pub fn add(&mut self, key: &[u8], value: &[u8]) -> bool {
        let entry_size = 4 + key.len() + 4 + value.len(); // key_len + key + value_len + value

        // Check if adding this entry would exceed max size
        // Always allow at least one entry
        if !self.is_empty() && self.buffer.len() + entry_size > self.max_size {
            return false;
        }

        // Record offset
        self.offsets.push(self.buffer.len() as u32);

        // Write key length and key
        self.buffer.put_u32_le(key.len() as u32);
        self.buffer.put_slice(key);

        // Write value length and value
        self.buffer.put_u32_le(value.len() as u32);
        self.buffer.put_slice(value);

        // Update last key
        self.last_key = Some(Bytes::copy_from_slice(key));

        true
    }

    /// Check if the block is empty
    pub fn is_empty(&self) -> bool {
        self.offsets.is_empty()
    }

    /// Get the last key in the block
    pub fn last_key(&self) -> Option<Bytes> {
        self.last_key.clone()
    }

    /// Get current size of the block
    pub fn size(&self) -> usize {
        self.buffer.len()
    }

    /// Finish building the block and return the data
    pub fn finish(&mut self) -> Vec<u8> {
        let mut result = Vec::with_capacity(self.buffer.len() + self.offsets.len() * 4 + 4);

        // Write data
        result.extend_from_slice(&self.buffer);

        // Write offsets
        for offset in &self.offsets {
            result.write_u32::<LittleEndian>(*offset).unwrap();
        }

        // Write number of entries
        result
            .write_u32::<LittleEndian>(self.offsets.len() as u32)
            .unwrap();

        // Reset for reuse
        self.buffer.clear();
        self.offsets.clear();
        self.last_key = None;

        result
    }
}

/// Builder for SSTable index
pub struct IndexBuilder {
    entries: Vec<IndexEntry>,
    pending_key: Option<Bytes>,
}

#[derive(Debug, Clone)]
pub(crate) struct IndexEntry {
    pub(crate) last_key: Bytes,
    pub(crate) block_offset: u64,
    pub(crate) block_size: u32,
}

impl IndexBuilder {
    /// Create a new index builder
    pub fn new() -> Self {
        Self {
            entries: Vec::new(),
            pending_key: None,
        }
    }

    /// Set a pending key for the next block
    pub fn set_pending_key(&mut self, key: &[u8]) {
        self.pending_key = Some(Bytes::copy_from_slice(key));
    }

    /// Check if there's a pending key
    pub fn has_pending_key(&self) -> bool {
        self.pending_key.is_some()
    }

    /// Add an index entry
    pub fn add(&mut self, last_key: &[u8], block_offset: u64, block_size: u32) -> Result<()> {
        self.entries.push(IndexEntry {
            last_key: Bytes::copy_from_slice(last_key),
            block_offset,
            block_size,
        });
        self.pending_key = None;
        Ok(())
    }

    /// Finish building the index and return the data
    pub fn finish(&mut self) -> Vec<u8> {
        let mut buffer = Vec::new();

        // Write each index entry
        for entry in &self.entries {
            // Write key length and key
            buffer
                .write_u32::<LittleEndian>(entry.last_key.len() as u32)
                .unwrap();
            buffer.extend_from_slice(&entry.last_key);

            // Write offset and size
            buffer
                .write_u64::<LittleEndian>(entry.block_offset)
                .unwrap();
            buffer.write_u32::<LittleEndian>(entry.block_size).unwrap();
        }

        // Write number of entries
        buffer
            .write_u32::<LittleEndian>(self.entries.len() as u32)
            .unwrap();

        buffer
    }

    /// Get the index entries (for reading)
    pub(crate) fn entries(&self) -> &[IndexEntry] {
        &self.entries
    }
}