Skip to main content

nodedb_columnar/
wal_record.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! WAL record types for columnar operations.
4//!
5//! Each mutation (INSERT, DELETE, memtable flush) produces a WAL record
6//! that is written before the mutation is applied. On crash recovery, WAL
7//! records are replayed to reconstruct the memtable, delete bitmaps, and
8//! segment metadata.
9//!
10//! Records are serialized as MessagePack for compact wire representation.
11
12use std::mem::size_of;
13
14use nodedb_types::decode_bounds::checked_decode_capacity;
15use serde::{Deserialize, Serialize};
16use sonic_rs;
17use zerompk::{FromMessagePack, ToMessagePack};
18
19/// A WAL record for a columnar collection operation.
20#[derive(Debug, Clone, Serialize, Deserialize, ToMessagePack, FromMessagePack)]
21#[serde(rename_all = "snake_case")]
22#[non_exhaustive]
23pub enum ColumnarWalRecord {
24    /// A row was inserted into the memtable.
25    ///
26    /// Contains the collection name and the row data as packed binary
27    /// (the columnar wire format, not MessagePack). On replay, the row
28    /// is re-inserted into the memtable.
29    #[serde(rename = "insert_row")]
30    InsertRow {
31        collection: String,
32        /// Row data as packed binary values. Each value is encoded per its
33        /// column type: i64 as 8 LE bytes, f64 as 8 LE bytes, strings as
34        /// length-prefixed UTF-8, etc.
35        row_data: Vec<u8>,
36    },
37
38    /// Rows were marked as deleted in a segment's delete bitmap.
39    ///
40    /// On replay, these row indices are re-applied to the segment's
41    /// delete bitmap.
42    #[serde(rename = "delete_rows")]
43    DeleteRows {
44        collection: String,
45        segment_id: u64,
46        row_indices: Vec<u32>,
47    },
48
49    /// The memtable was flushed to a new segment.
50    ///
51    /// On replay, if the segment file exists, update metadata to include it.
52    /// If it doesn't exist, the flush was interrupted; rows are already in
53    /// the memtable via InsertRow records.
54    #[serde(rename = "memtable_flushed")]
55    MemtableFlushed {
56        collection: String,
57        segment_id: u64,
58        row_count: u64,
59    },
60}
61
62impl ColumnarWalRecord {
63    /// Collection name this record belongs to.
64    pub fn collection(&self) -> &str {
65        match self {
66            Self::InsertRow { collection, .. }
67            | Self::DeleteRows { collection, .. }
68            | Self::MemtableFlushed { collection, .. } => collection,
69        }
70    }
71
72    /// Serialize the record to bytes.
73    pub fn to_bytes(&self) -> Result<Vec<u8>, crate::error::ColumnarError> {
74        zerompk::to_msgpack_vec(self)
75            .map_err(|e| crate::error::ColumnarError::Serialization(e.to_string()))
76    }
77
78    /// Deserialize a record from bytes.
79    pub fn from_bytes(data: &[u8]) -> Result<Self, crate::error::ColumnarError> {
80        zerompk::from_msgpack(data)
81            .map_err(|e| crate::error::ColumnarError::Serialization(e.to_string()))
82    }
83}
84
85/// Encode a row of values into the columnar wire format for WAL records.
86///
87/// Each value is written as: [type_tag: u8][value_bytes].
88/// This is more compact than MessagePack for typed columns and enables
89/// direct replay into the memtable without schema interpretation overhead.
90pub fn encode_row_for_wal(
91    values: &[nodedb_types::value::Value],
92) -> Result<Vec<u8>, crate::error::ColumnarError> {
93    use nodedb_types::value::Value;
94
95    let mut buf = Vec::with_capacity(values.len() * 10); // Rough estimate.
96
97    for value in values {
98        match value {
99            Value::Null => buf.push(0),
100            Value::Integer(v) => {
101                buf.push(1);
102                buf.extend_from_slice(&v.to_le_bytes());
103            }
104            Value::Float(v) => {
105                buf.push(2);
106                buf.extend_from_slice(&v.to_le_bytes());
107            }
108            Value::Bool(v) => {
109                buf.push(3);
110                buf.push(*v as u8);
111            }
112            Value::String(s) => {
113                buf.push(4);
114                let bytes = s.as_bytes();
115                buf.extend_from_slice(&(bytes.len() as u32).to_le_bytes());
116                buf.extend_from_slice(bytes);
117            }
118            Value::Bytes(b) => {
119                buf.push(5);
120                buf.extend_from_slice(&(b.len() as u32).to_le_bytes());
121                buf.extend_from_slice(b);
122            }
123            Value::DateTime(dt) => {
124                buf.push(6);
125                buf.extend_from_slice(&dt.micros.to_le_bytes());
126            }
127            Value::Decimal(d) => {
128                buf.push(7);
129                buf.extend_from_slice(&d.serialize());
130            }
131            Value::Uuid(s) => {
132                buf.push(8);
133                let bytes = s.as_bytes();
134                buf.extend_from_slice(&(bytes.len() as u32).to_le_bytes());
135                buf.extend_from_slice(bytes);
136            }
137            Value::Array(arr) if arr.iter().all(|v| matches!(v, Value::Float(_))) => {
138                // An all-float array: tag(9) + count(u32) + f64 values, so
139                // every element decodes to the value it was.
140                buf.push(9);
141                buf.extend_from_slice(&(arr.len() as u32).to_le_bytes());
142                for v in arr {
143                    if let Value::Float(f) = v {
144                        buf.extend_from_slice(&f.to_le_bytes());
145                    }
146                }
147            }
148            _ => {
149                // Any other array, geometry and other complex types: JSON
150                // bytes.
151                buf.push(10);
152                let json = sonic_rs::to_vec(value).map_err(|e| {
153                    crate::error::ColumnarError::Serialization(format!(
154                        "failed to serialize value as JSON: {e}"
155                    ))
156                })?;
157                buf.extend_from_slice(&(json.len() as u32).to_le_bytes());
158                buf.extend_from_slice(&json);
159            }
160        }
161    }
162
163    Ok(buf)
164}
165
166/// Maximum length for a variable-length field in a WAL record (256 MiB).
167/// Prevents OOM from crafted/corrupt records with bogus length prefixes.
168const MAX_FIELD_LEN: usize = 256 * 1024 * 1024;
169
170/// The error for a WAL row that stops decoding at byte `offset`.
171fn corrupt(offset: usize, reason: impl Into<String>) -> crate::error::ColumnarError {
172    crate::error::ColumnarError::WalRowCorrupt {
173        offset,
174        reason: reason.into(),
175    }
176}
177
178/// Read exactly `N` bytes from `data` at `cursor` as an array, advancing
179/// cursor. Returns `Err` if not enough bytes remain.
180fn read_array<const N: usize>(
181    data: &[u8],
182    cursor: &mut usize,
183    context: &str,
184) -> Result<[u8; N], crate::error::ColumnarError> {
185    let at = *cursor;
186    let slice = read_slice(data, cursor, N, context)?;
187    slice
188        .try_into()
189        .map_err(|_| corrupt(at, format!("truncated {context}")))
190}
191
192/// Read exactly `n` bytes from `data` at `cursor`, advancing cursor.
193/// Returns `Err` if not enough bytes remain.
194fn read_slice<'a>(
195    data: &'a [u8],
196    cursor: &mut usize,
197    n: usize,
198    context: &str,
199) -> Result<&'a [u8], crate::error::ColumnarError> {
200    let end = cursor
201        .checked_add(n)
202        .ok_or_else(|| corrupt(*cursor, format!("overflow in {context}")))?;
203    if end > data.len() {
204        return Err(corrupt(
205            *cursor,
206            format!(
207                "truncated {context}: need {n} bytes, have {}",
208                data.len().saturating_sub(*cursor)
209            ),
210        ));
211    }
212    let slice = &data[*cursor..end];
213    *cursor = end;
214    Ok(slice)
215}
216
217/// Read a u32 length prefix, validate it against MAX_FIELD_LEN, then read
218/// that many bytes. Returns the payload slice.
219fn read_length_prefixed<'a>(
220    data: &'a [u8],
221    cursor: &mut usize,
222    context: &str,
223) -> Result<&'a [u8], crate::error::ColumnarError> {
224    let at = *cursor;
225    let len = u32::from_le_bytes(read_array::<4>(data, cursor, context)?) as usize;
226    if len > MAX_FIELD_LEN {
227        return Err(corrupt(
228            at,
229            format!("{context} length {len} exceeds maximum {MAX_FIELD_LEN}"),
230        ));
231    }
232    read_slice(data, cursor, len, context)
233}
234
235/// Read a length-prefixed UTF-8 string. Bytes that are not UTF-8 are an
236/// error, never replacement characters.
237fn read_utf8(
238    data: &[u8],
239    cursor: &mut usize,
240    context: &str,
241) -> Result<String, crate::error::ColumnarError> {
242    let at = *cursor;
243    let bytes = read_length_prefixed(data, cursor, context)?;
244    String::from_utf8(bytes.to_vec())
245        .map_err(|e| corrupt(at, format!("{context} is not UTF-8: {e}")))
246}
247
248/// Decode a row from the columnar wire format back into Values.
249///
250/// `Err(WalRowCorrupt)` when the bytes do not decode. The corruption is
251/// reported here, where it is detected.
252pub fn decode_row_from_wal(
253    data: &[u8],
254) -> Result<Vec<nodedb_types::value::Value>, crate::error::ColumnarError> {
255    decode_row(data).inspect_err(crate::diag::wal_row_corrupt)
256}
257
258/// The body of [`decode_row_from_wal`].
259fn decode_row(data: &[u8]) -> Result<Vec<nodedb_types::value::Value>, crate::error::ColumnarError> {
260    use nodedb_types::value::Value;
261
262    let mut values = Vec::new();
263    let mut cursor = 0;
264
265    while cursor < data.len() {
266        let tag_at = cursor;
267        let [tag] = read_array::<1>(data, &mut cursor, "tag")?;
268
269        let value = match tag {
270            0 => Value::Null,
271            1 => Value::Integer(i64::from_le_bytes(read_array(data, &mut cursor, "i64")?)),
272            2 => Value::Float(f64::from_le_bytes(read_array(data, &mut cursor, "f64")?)),
273            3 => {
274                let [b] = read_array::<1>(data, &mut cursor, "bool")?;
275                Value::Bool(b != 0)
276            }
277            4 => Value::String(read_utf8(data, &mut cursor, "string")?),
278            5 => Value::Bytes(read_length_prefixed(data, &mut cursor, "bytes")?.to_vec()),
279            8 => Value::Uuid(read_utf8(data, &mut cursor, "uuid")?),
280            6 => {
281                let micros = i64::from_le_bytes(read_array(data, &mut cursor, "timestamp")?);
282                Value::DateTime(nodedb_types::datetime::NdbDateTime::from_micros(micros))
283            }
284            7 => Value::Decimal(rust_decimal::Decimal::deserialize(read_array(
285                data,
286                &mut cursor,
287                "decimal",
288            )?)),
289            9 => {
290                let count_at = cursor;
291                let count =
292                    u32::from_le_bytes(read_array(data, &mut cursor, "vector count")?) as usize;
293                let remaining_values = data.len().saturating_sub(cursor) / 8;
294                let max_count = (MAX_FIELD_LEN / 8).min(remaining_values);
295                if count > max_count {
296                    return Err(corrupt(
297                        count_at,
298                        format!("vector count {count} exceeds maximum {max_count}"),
299                    ));
300                }
301                let capacity = checked_decode_capacity(
302                    count,
303                    size_of::<nodedb_types::value::Value>(),
304                    data.len().saturating_sub(cursor),
305                    8,
306                    max_count,
307                    usize::MAX,
308                )
309                .ok_or_else(|| {
310                    corrupt(count_at, "vector count exceeds decode allocation bounds")
311                })?;
312                let mut arr = Vec::with_capacity(capacity);
313                for _ in 0..count {
314                    let f = f64::from_le_bytes(read_array(data, &mut cursor, "vector f64")?);
315                    arr.push(Value::Float(f));
316                }
317                Value::Array(arr)
318            }
319            10 => {
320                let json_at = cursor;
321                let json_bytes = read_length_prefixed(data, &mut cursor, "json")?;
322                sonic_rs::from_slice(json_bytes)
323                    .map_err(|e| corrupt(json_at, format!("json value does not decode: {e}")))?
324            }
325            _ => return Err(corrupt(tag_at, format!("unknown WAL value tag: {tag}"))),
326        };
327
328        values.push(value);
329    }
330
331    Ok(values)
332}
333
334#[cfg(test)]
335mod tests {
336    use nodedb_types::datetime::NdbDateTime;
337    use nodedb_types::value::Value;
338
339    use super::*;
340
341    #[test]
342    fn rejects_huge_vector_count_with_tiny_payload_before_allocation() {
343        // tag + one value whose vector count claims nearly u32::MAX values.
344        let mut bytes = vec![9];
345        bytes.extend_from_slice(&u32::MAX.to_le_bytes());
346        assert!(matches!(
347            decode_row_from_wal(&bytes),
348            Err(crate::error::ColumnarError::WalRowCorrupt { .. })
349        ));
350    }
351
352    #[test]
353    fn a_string_that_is_not_utf8_is_refused() {
354        let mut bytes = vec![4];
355        bytes.extend_from_slice(&2u32.to_le_bytes());
356        bytes.extend_from_slice(&[0xC3, 0x28]);
357        assert!(matches!(
358            decode_row_from_wal(&bytes),
359            Err(crate::error::ColumnarError::WalRowCorrupt { offset: 1, .. })
360        ));
361    }
362
363    #[test]
364    fn a_uuid_that_is_not_utf8_is_refused() {
365        let mut bytes = vec![8];
366        bytes.extend_from_slice(&1u32.to_le_bytes());
367        bytes.push(0xFF);
368        assert!(matches!(
369            decode_row_from_wal(&bytes),
370            Err(crate::error::ColumnarError::WalRowCorrupt { .. })
371        ));
372    }
373
374    #[test]
375    fn a_json_value_that_does_not_decode_is_refused() {
376        let json = b"{not json";
377        let mut bytes = vec![10];
378        bytes.extend_from_slice(&(json.len() as u32).to_le_bytes());
379        bytes.extend_from_slice(json);
380        assert!(matches!(
381            decode_row_from_wal(&bytes),
382            Err(crate::error::ColumnarError::WalRowCorrupt { .. })
383        ));
384    }
385
386    #[test]
387    fn arrays_decode_to_the_values_they_were() {
388        let values = vec![
389            Value::Array(vec![Value::Float(0.1), Value::Float(-2.5)]),
390            Value::Array(vec![Value::String("a".into()), Value::String("b".into())]),
391        ];
392        let encoded = encode_row_for_wal(&values).expect("encode");
393        assert_eq!(decode_row_from_wal(&encoded).expect("decode"), values);
394    }
395
396    #[test]
397    fn wal_record_roundtrip() {
398        let records = vec![
399            ColumnarWalRecord::InsertRow {
400                collection: "test".into(),
401                row_data: vec![1, 2, 3],
402            },
403            ColumnarWalRecord::DeleteRows {
404                collection: "test".into(),
405                segment_id: 0,
406                row_indices: vec![5, 10, 15],
407            },
408            ColumnarWalRecord::MemtableFlushed {
409                collection: "test".into(),
410                segment_id: 3,
411                row_count: 1024,
412            },
413        ];
414
415        for record in &records {
416            let bytes = record.to_bytes().expect("serialize");
417            let restored = ColumnarWalRecord::from_bytes(&bytes).expect("deserialize");
418            assert_eq!(restored.collection(), record.collection());
419        }
420    }
421
422    #[test]
423    fn row_wire_format_roundtrip() {
424        let values = vec![
425            Value::Integer(42),
426            Value::Float(0.75),
427            Value::Bool(true),
428            Value::String("hello".into()),
429            Value::Bytes(vec![0xDE, 0xAD]),
430            Value::DateTime(NdbDateTime::from_micros(1_700_000_000)),
431            Value::Decimal(rust_decimal::Decimal::new(314, 2)),
432            Value::Uuid("550e8400-e29b-41d4-a716-446655440000".into()),
433            Value::Null,
434            Value::Array(vec![Value::Float(1.0), Value::Float(2.0)]),
435        ];
436
437        let encoded = encode_row_for_wal(&values).expect("encode");
438        let decoded = decode_row_from_wal(&encoded).expect("decode");
439
440        assert_eq!(decoded.len(), values.len());
441        assert_eq!(decoded[0], Value::Integer(42));
442        assert_eq!(decoded[1], Value::Float(0.75));
443        assert_eq!(decoded[2], Value::Bool(true));
444        assert_eq!(decoded[3], Value::String("hello".into()));
445        assert_eq!(decoded[4], Value::Bytes(vec![0xDE, 0xAD]));
446        assert_eq!(
447            decoded[5],
448            Value::DateTime(NdbDateTime::from_micros(1_700_000_000))
449        );
450        assert_eq!(
451            decoded[7],
452            Value::Uuid("550e8400-e29b-41d4-a716-446655440000".into())
453        );
454        assert_eq!(decoded[8], Value::Null);
455    }
456
457    #[test]
458    fn decode_truncated_i64_returns_error() {
459        // Tag 1 (i64) requires 8 payload bytes; supply none.
460        // A slice index `data[cursor..cursor+8]` here would panic with an
461        // index out-of-bounds. `try_into()` must return the
462        // Serialization error instead.
463        let result = decode_row_from_wal(&[1]);
464        assert!(
465            result.is_err(),
466            "truncated i64 payload must return Err, not panic"
467        );
468    }
469
470    #[test]
471    fn decode_truncated_string_returns_error() {
472        // Tag 4 (string): length prefix says 255 bytes but the slice ends
473        // immediately after the 4-byte length field. The read of
474        // `data[cursor..cursor+255]` must error, not panic.
475        let input = {
476            let mut v = vec![4u8]; // tag = string
477            v.extend_from_slice(&255u32.to_le_bytes()); // len = 255
478            // no payload bytes follow
479            v
480        };
481        let result = decode_row_from_wal(&input);
482        assert!(
483            result.is_err(),
484            "truncated string payload must return Err, not panic"
485        );
486    }
487
488    #[test]
489    fn decode_huge_vector_count_returns_error() {
490        // Tag 9 (vector array): count = 0x7FFFFFFF. After reading the count,
491        // the very first iteration tries to read 4 bytes of f32 from an empty
492        // slice. The loop must error out cleanly there, before any allocation
493        // proportional to count is attempted, rather than panic.
494        let input = {
495            let mut v = vec![9u8]; // tag = vector array
496            v.extend_from_slice(&0x7FFF_FFFFu32.to_le_bytes()); // count
497            // no f32 bytes follow
498            v
499        };
500        let result = decode_row_from_wal(&input);
501        assert!(
502            result.is_err(),
503            "huge vector count with no payload must return Err, not panic or OOM"
504        );
505    }
506
507    #[test]
508    fn decode_truncated_decimal_returns_error() {
509        // Tag 7 (Decimal) requires 16 bytes; supply only 4.
510        // `data[cursor..cursor+16]` must error, not panic.
511        let input = {
512            let mut v = vec![7u8]; // tag = decimal
513            v.extend_from_slice(&[0u8; 4]); // only 4 bytes, need 16
514            v
515        };
516        let result = decode_row_from_wal(&input);
517        assert!(
518            result.is_err(),
519            "truncated decimal payload must return Err, not panic"
520        );
521    }
522}