mongreldb_core/
memtable.rs

1//! In-memory write buffer (the "memtable").
2//!
3//! Phase 11.2 wires the buffered [`crate::be_tree::BeTree`] (a Bε-tree over the
4//! composite `(RowId, Epoch)` version key) in as the live memtable, replacing
5//! the prototype skip list. A Bε-tree buffers many pending mutations per
6//! internal node and flushes them to one child in bulk, so write amplification
7//! approaches O(1) — the update-amplification win the design calls for. The
8//! composite key keeps multiple versions of a logical row coexisting so a
9//! snapshot read returns the newest version with `committed_epoch <= snapshot`,
10//! which is what makes MVCC correct within the live memtable.
11
12use crate::be_tree::BeTree;
13use crate::epoch::Epoch;
14use crate::rowid::RowId;
15use serde::{Deserialize, Serialize};
16use std::collections::{BTreeMap, HashMap};
17use std::sync::Arc;
18
19/// A cell value in the in-memory path. The flush path re-encodes these into
20/// columnar pages; it is intentionally simple for the prototype.
21#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
22pub enum Value {
23    Null,
24    Bool(bool),
25    Int64(i64),
26    Float64(f64),
27    Bytes(Vec<u8>),
28    Embedding(Vec<f32>),
29    /// Unscaled decimal value (i128). The column's `TypeId::Decimal128`
30    /// carries the precision/scale for formatting.
31    Decimal(i128),
32    /// SQL INTERVAL value: months, days, nanoseconds.
33    Interval {
34        months: i64,
35        days: i32,
36        nanos: i64,
37    },
38    /// RFC 4122 UUID (16 bytes, big-endian for sort order).
39    Uuid([u8; 16]),
40    /// JSON value stored as a UTF-8 byte sequence.
41    Json(Vec<u8>),
42}
43
44impl Value {
45    /// Lexicographically-comparable byte encoding for index keys (PK HOT,
46    /// bitmaps). Big-endian for integers so byte order matches value order.
47    pub fn encode_key(&self) -> Vec<u8> {
48        match self {
49            Value::Null => Vec::new(),
50            Value::Bool(b) => vec![*b as u8],
51            Value::Int64(n) => n.to_be_bytes().to_vec(),
52            Value::Float64(f) => f.to_bits().to_be_bytes().to_vec(),
53            Value::Bytes(b) => b.clone(),
54            Value::Embedding(v) => {
55                let mut out = Vec::with_capacity(v.len() * 4);
56                for x in v {
57                    out.extend_from_slice(&x.to_bits().to_be_bytes());
58                }
59                out
60            }
61            Value::Decimal(d) => d.to_be_bytes().to_vec(),
62            Value::Interval {
63                months,
64                days,
65                nanos,
66            } => {
67                let mut out = Vec::with_capacity(20);
68                out.extend_from_slice(&months.to_be_bytes());
69                out.extend_from_slice(&days.to_be_bytes());
70                out.extend_from_slice(&nanos.to_be_bytes());
71                out
72            }
73            Value::Uuid(b) => b.to_vec(),
74            Value::Json(b) => b.clone(),
75        }
76    }
77
78    pub(crate) fn estimated_bytes(&self) -> u64 {
79        match self {
80            Value::Null => 1,
81            Value::Bool(_) => 1,
82            Value::Int64(_) | Value::Float64(_) => 8,
83            Value::Bytes(bytes) | Value::Json(bytes) => 16 + bytes.len() as u64,
84            Value::Embedding(values) => 16 + (values.len() as u64) * 4,
85            Value::Decimal(_) | Value::Uuid(_) => 16,
86            Value::Interval { .. } => 20,
87        }
88    }
89}
90
91/// One logical row held in the memtable. A `deleted` row is a tombstone.
92#[derive(Debug, Clone, Serialize, Deserialize)]
93pub struct Row {
94    pub row_id: RowId,
95    pub committed_epoch: Epoch,
96    pub columns: HashMap<u16, Value>,
97    pub deleted: bool,
98}
99
100impl Row {
101    pub fn new(row_id: RowId, committed_epoch: Epoch) -> Self {
102        Self {
103            row_id,
104            committed_epoch,
105            columns: HashMap::new(),
106            deleted: false,
107        }
108    }
109
110    pub fn with_column(mut self, column_id: u16, value: Value) -> Self {
111        self.columns.insert(column_id, value);
112        self
113    }
114
115    /// Rough byte estimate for flush-threshold decisions.
116    pub fn estimated_bytes(&self) -> u64 {
117        self.columns
118            .values()
119            .fold(32, |bytes, value| bytes + value.estimated_bytes())
120    }
121}
122
123/// Bε-tree-backed memtable, ordered by `(RowId, Epoch)`. A drop-in replacement
124/// for the prototype skip list: the same MVCC semantics with lower write
125/// amplification (buffered messages flush to children in bulk).
126#[derive(Clone)]
127struct MemtableSegment {
128    tree: BeTree,
129    byte_size: u64,
130}
131
132/// Structurally shared committed overlays plus one small mutable write delta.
133#[derive(Clone)]
134pub struct Memtable {
135    frozen: Arc<Vec<Arc<MemtableSegment>>>,
136    active: MemtableSegment,
137    byte_size: u64,
138}
139
140impl Default for Memtable {
141    fn default() -> Self {
142        Self::new()
143    }
144}
145
146impl Memtable {
147    pub fn new() -> Self {
148        Self {
149            frozen: Arc::new(Vec::new()),
150            active: MemtableSegment {
151                tree: BeTree::new(),
152                byte_size: 0,
153            },
154            byte_size: 0,
155        }
156    }
157
158    /// Append a row version (keyed by `(row_id, committed_epoch)`). Versions are
159    /// never overwritten; the newest visible one wins at read time.
160    pub fn upsert(&mut self, row: Row) {
161        let bytes = row.estimated_bytes();
162        self.byte_size = self.byte_size.saturating_add(bytes);
163        self.active.byte_size = self.active.byte_size.saturating_add(bytes);
164        self.active.tree.insert_row(row);
165    }
166
167    /// Append a tombstone version for `row_id` at `epoch`. The tombstone copies
168    /// the columns from the newest live version so that engine-level HOT cleanup
169    /// can recover the primary-key value during WAL replay.
170    pub fn tombstone(&mut self, row_id: RowId, epoch: Epoch) {
171        let mut columns = HashMap::new();
172        if let Some(live) = self.get(row_id, Epoch(epoch.0.saturating_sub(1))) {
173            columns = live.columns;
174        }
175        let row = Row {
176            row_id,
177            committed_epoch: epoch,
178            columns,
179            deleted: true,
180        };
181        self.upsert(row);
182    }
183
184    /// Read the row at `row_id` visible to `snapshot`: the newest version with
185    /// `epoch <= snapshot`. Returns `None` if that version is a tombstone (or no
186    /// such version exists).
187    pub fn get(&self, row_id: RowId, snapshot_epoch: Epoch) -> Option<Row> {
188        self.get_version(row_id, snapshot_epoch)
189            .and_then(|(_, row)| (!row.deleted).then_some(row))
190    }
191
192    /// Newest version of `row_id` with `epoch <= snapshot`, **including
193    /// tombstones** (as a `Row` with `deleted=true`). Used by the engine to
194    /// merge versions across the memtable and sorted runs.
195    pub fn get_version(&self, row_id: RowId, snapshot_epoch: Epoch) -> Option<(Epoch, Row)> {
196        let mut best = self.active.tree.get_version(row_id, snapshot_epoch);
197        for segment in self.frozen.iter().rev() {
198            let Some(candidate) = segment.tree.get_version(row_id, snapshot_epoch) else {
199                continue;
200            };
201            if best.as_ref().is_none_or(|(epoch, _)| candidate.0 > *epoch) {
202                best = Some(candidate);
203            }
204        }
205        best
206    }
207
208    /// Number of stored versions.
209    pub fn len(&self) -> usize {
210        self.active.tree.mutations()
211            + self
212                .frozen
213                .iter()
214                .map(|segment| segment.tree.mutations())
215                .sum::<usize>()
216    }
217
218    pub fn is_empty(&self) -> bool {
219        self.active.tree.is_empty() && self.frozen.is_empty()
220    }
221
222    pub fn approx_bytes(&self) -> u64 {
223        self.byte_size
224    }
225
226    /// Visible rows at `snapshot`, deduplicated to the newest version per
227    /// `RowId` (tombstones drop their row). Returned in ascending `RowId` order.
228    pub fn visible_rows(&self, snapshot_epoch: Epoch) -> Vec<Row> {
229        self.visible_versions(snapshot_epoch)
230            .into_iter()
231            .filter(|r| !r.deleted)
232            .collect()
233    }
234
235    /// Newest visible version per `RowId` at `snapshot`, **including
236    /// tombstones** (as `Row`s with `deleted=true`). Used by the engine to merge
237    /// versions across the memtable and sorted runs.
238    pub fn visible_versions(&self, snapshot_epoch: Epoch) -> Vec<Row> {
239        let mut by_row: BTreeMap<RowId, Row> = BTreeMap::new();
240        for segment in self
241            .frozen
242            .iter()
243            .map(|segment| &segment.tree)
244            .chain(std::iter::once(&self.active.tree))
245        {
246            for row in segment.versions() {
247                if row.committed_epoch <= snapshot_epoch {
248                    by_row
249                        .entry(row.row_id)
250                        .and_modify(|existing| {
251                            if row.committed_epoch > existing.committed_epoch {
252                                *existing = row.clone();
253                            }
254                        })
255                        .or_insert(row);
256                }
257            }
258        }
259        by_row.into_values().collect()
260    }
261
262    /// Freeze the current write delta so future clones share it by `Arc`.
263    pub(crate) fn seal(&mut self) {
264        if self.active.tree.is_empty() {
265            return;
266        }
267        let active = std::mem::replace(
268            &mut self.active,
269            MemtableSegment {
270                tree: BeTree::new(),
271                byte_size: 0,
272            },
273        );
274        Arc::make_mut(&mut self.frozen).push(Arc::new(active));
275        if self.frozen.len() >= crate::MAX_READ_GENERATION_LAYERS {
276            self.consolidate();
277        }
278    }
279
280    fn consolidate(&mut self) {
281        let mut tree = BeTree::new();
282        for row in self
283            .frozen
284            .iter()
285            .flat_map(|segment| segment.tree.versions())
286        {
287            tree.insert_row(row);
288        }
289        self.frozen = Arc::new(vec![Arc::new(MemtableSegment {
290            tree,
291            byte_size: self.byte_size,
292        })]);
293    }
294
295    #[cfg(test)]
296    pub(crate) fn frozen_layer_count(&self) -> usize {
297        self.frozen.len()
298    }
299
300    /// Drain all versions (for a memtable-to-run flush). Returns them in
301    /// ascending `(RowId, Epoch)` order.
302    pub fn drain_sorted(&mut self) -> Vec<Row> {
303        let mut out = self
304            .frozen
305            .iter()
306            .flat_map(|segment| segment.tree.versions())
307            .chain(self.active.tree.versions())
308            .collect::<Vec<_>>();
309        out.sort_by_key(|row| (row.row_id, row.committed_epoch));
310        self.frozen = Arc::new(Vec::new());
311        self.active = MemtableSegment {
312            tree: BeTree::new(),
313            byte_size: 0,
314        };
315        self.byte_size = 0;
316        out
317    }
318}
319
320#[cfg(test)]
321mod tests {
322    use super::*;
323
324    fn row(id: u64, epoch: u64) -> Row {
325        Row::new(RowId(id), Epoch(epoch)).with_column(1, Value::Int64(id as i64 * 10))
326    }
327
328    #[test]
329    fn upsert_get_and_visibility() {
330        let mut m = Memtable::new();
331        m.upsert(row(1, 5));
332        assert_eq!(m.len(), 1);
333        assert!(m.get(RowId(1), Epoch(5)).is_some());
334        assert!(m.get(RowId(1), Epoch(4)).is_none()); // not yet visible
335        assert!(m.get(RowId(2), Epoch(9)).is_none()); // missing
336    }
337
338    #[test]
339    fn tombstone_supersedes_at_its_epoch() {
340        let mut m = Memtable::new();
341        m.upsert(row(1, 1));
342        // Before the tombstone: the live version is visible.
343        assert!(m.get(RowId(1), Epoch(1)).is_some());
344        m.tombstone(RowId(1), Epoch(2));
345        // At/after the tombstone: hidden.
346        assert!(m.get(RowId(1), Epoch(2)).is_none());
347        assert!(m.get(RowId(1), Epoch(9)).is_none());
348        // A snapshot before the tombstone still sees the live version.
349        assert!(m.get(RowId(1), Epoch(1)).is_some());
350    }
351
352    #[test]
353    fn sealed_generations_share_rows_and_consolidate() {
354        let mut writer = Memtable::new();
355        for id in 0..crate::MAX_READ_GENERATION_LAYERS as u64 + 2 {
356            writer.upsert(row(id, id + 1));
357            writer.seal();
358        }
359        assert!(writer.frozen_layer_count() < crate::MAX_READ_GENERATION_LAYERS);
360        let generation = writer.clone();
361        writer.upsert(row(99, 99));
362        assert!(generation.get(RowId(99), Epoch(99)).is_none());
363        assert!(writer.get(RowId(99), Epoch(99)).is_some());
364    }
365
366    #[test]
367    fn drain_sorted_is_ascending_and_empties() {
368        let mut m = Memtable::new();
369        m.upsert(row(3, 1));
370        m.upsert(row(1, 1));
371        m.upsert(row(2, 1));
372        let out = m.drain_sorted();
373        let ids: Vec<u64> = out.iter().map(|r| r.row_id.0).collect();
374        assert_eq!(ids, vec![1, 2, 3]);
375        assert!(m.is_empty());
376        assert_eq!(m.approx_bytes(), 0);
377    }
378
379    #[test]
380    fn visible_rows_dedups_to_newest_version() {
381        let mut m = Memtable::new();
382        m.upsert(row(1, 1));
383        m.upsert(row(2, 9)); // future relative to snapshot 5
384        m.upsert(row(3, 1));
385        m.upsert(row(1, 3)); // newer version of row 1
386        let ids: Vec<u64> = m
387            .visible_rows(Epoch(5))
388            .iter()
389            .map(|r| r.row_id.0)
390            .collect();
391        assert_eq!(ids, vec![1, 3]);
392    }
393}
mongreldb_core/memtable.rs

mongreldb_core/
memtable.rs
