Skip to main content

cqlite_core/storage/write_engine/
memtable.rs

1//! In-memory write buffer (memtable)
2//!
3//! Stores mutations in memory using a BTreeMap for partition and clustering ordering.
4//! Flushes to L0 SSTable when size threshold is reached.
5//!
6//! The memtable maintains mutations in token-sorted order (via DecoratedKey) and tracks
7//! approximate memory usage to trigger flushes at a configurable threshold.
8
9use crate::error::Result;
10use crate::storage::write_engine::mutation::{DecoratedKey, Mutation};
11use std::collections::BTreeMap;
12
13/// In-memory write buffer
14///
15/// Stores mutations in memory with token-based ordering. Each partition can have
16/// multiple mutations (e.g., multiple rows with different clustering keys).
17#[derive(Debug)]
18pub struct Memtable {
19    /// Partition-level storage: token-ordered map of mutations
20    data: BTreeMap<DecoratedKey, Vec<Mutation>>,
21    /// Approximate size in bytes
22    size_bytes: usize,
23    /// Approximate row count (total mutations across all partitions)
24    row_count: usize,
25    /// Creation timestamp (Unix epoch microseconds)
26    created_at: i64,
27}
28
29impl Memtable {
30    /// Create a new memtable
31    pub fn new() -> Self {
32        Self {
33            data: BTreeMap::new(),
34            size_bytes: 0,
35            row_count: 0,
36            created_at: Self::current_timestamp_micros(),
37        }
38    }
39
40    /// Insert a mutation into the memtable
41    ///
42    /// Mutations are grouped by partition key (DecoratedKey). Multiple mutations
43    /// for the same partition are stored as a vector.
44    pub fn insert(&mut self, _mutation: Mutation) -> Result<()> {
45        // Calculate decorated key from partition key
46        // Note: This requires schema, but mutation doesn't store it.
47        // For now, we expect the mutation to be pre-validated and the key
48        // to be extractable. In practice, the caller will need to provide
49        // the decorated key or schema context.
50        //
51        // WORKAROUND: Since Mutation doesn't store DecoratedKey directly,
52        // and calculating it requires schema, we need to pass the key separately.
53        // For Issue #362, we'll implement a public API that accepts DecoratedKey.
54        //
55        // This is a design limitation that will be addressed in the full WriteEngine.
56        // For now, insert_with_key() is the primary API.
57
58        // This method is kept for API compatibility but requires rethinking.
59        // We'll implement the core logic in insert_with_key() below.
60        Err(crate::error::Error::InvalidInput(
61            "Use insert_with_key() - decorated key must be provided with mutation".to_string(),
62        ))
63    }
64
65    /// Insert a mutation with an explicit decorated key
66    ///
67    /// This is the primary insertion API. The caller is responsible for computing
68    /// the decorated key from the partition key using the table schema.
69    #[tracing::instrument(name = "memtable.insert", level = "debug", skip(self, key, mutation))]
70    pub fn insert_with_key(&mut self, key: DecoratedKey, mutation: Mutation) -> Result<()> {
71        // Calculate mutation size (conservative estimate). This is the SAME
72        // computation the admission gate consults via `estimate_mutation_size`,
73        // so accounting and admission can never drift (issue #1625).
74        let mutation_size = Self::mutation_size(&mutation);
75
76        // Get or create mutation list for this partition
77        let mutations = self.data.entry(key).or_default();
78
79        // Add mutation
80        mutations.push(mutation);
81        self.row_count = self.row_count.saturating_add(1);
82        // `saturating_add`: `mutation_size` can legitimately be `usize::MAX`
83        // (the estimator fails closed at the node cap, issue #1625). Admission
84        // rejects over-limit mutations, but a direct `Memtable` user (bypassing
85        // admission) or a `WriteEngine` with `memtable_hard_limit == usize::MAX`
86        // could otherwise panic in debug (overflow check) or wrap in release —
87        // the ledger update MUST be self-safe.
88        self.size_bytes = self.size_bytes.saturating_add(mutation_size);
89
90        Ok(())
91    }
92
93    /// Get all mutations for a given partition key
94    pub fn get(&self, key: &DecoratedKey) -> Option<&[Mutation]> {
95        self.data.get(key).map(|v| v.as_slice())
96    }
97
98    /// Check if memtable is empty
99    pub fn is_empty(&self) -> bool {
100        self.data.is_empty()
101    }
102
103    /// Get current size in bytes (approximate)
104    pub fn size_bytes(&self) -> usize {
105        self.size_bytes
106    }
107
108    /// Get approximate row count
109    pub fn row_count(&self) -> usize {
110        self.row_count
111    }
112
113    /// Check if memtable should be flushed
114    pub fn should_flush(&self, threshold_bytes: usize) -> bool {
115        self.size_bytes >= threshold_bytes
116    }
117
118    /// Get creation timestamp (microseconds since Unix epoch)
119    pub fn created_at(&self) -> i64 {
120        self.created_at
121    }
122
123    /// Iterate over all partitions and their mutations
124    ///
125    /// Returns an iterator over (DecoratedKey, mutations) pairs in token order.
126    pub fn iter(&self) -> impl Iterator<Item = (&DecoratedKey, &[Mutation])> {
127        self.data.iter().map(|(k, v)| (k, v.as_slice()))
128    }
129
130    /// Clear all data from the memtable
131    ///
132    /// Used after successful flush to SSTable.
133    pub fn clear(&mut self) {
134        self.data.clear();
135        self.size_bytes = 0;
136        self.row_count = 0;
137        // Keep created_at unchanged - represents original creation time
138    }
139
140    /// Upper bound on the number of value nodes the ITERATIVE size estimator
141    /// visits before it fails closed (issue #1625).
142    ///
143    /// The estimator walks a value with an explicit heap worklist (no recursion,
144    /// so no stack-overflow risk and no depth cap collapsing deep children to a
145    /// floor). To bound worst-case work on a pathologically huge/deep value, the
146    /// traversal stops after visiting this many nodes and returns a
147    /// CONSERVATIVE-LARGE estimate (`usize::MAX`) that is GUARANTEED to exceed any
148    /// hard limit, so admission REJECTS the value rather than under-counting it.
149    /// Failing closed on the pathological case is the correct behavior for a DoS
150    /// guard. `1_000_000` is far beyond any legitimate mutation's node count yet
151    /// keeps the traversal cheap.
152    const MAX_ESTIMATE_NODES: usize = 1_000_000;
153
154    /// Estimate the number of bytes a mutation would add to this memtable.
155    ///
156    /// Returns the SAME value that [`Memtable::insert_with_key`] adds to
157    /// `size_bytes`, so the write-engine admission gate (issue #1625) and the
158    /// running size accounting agree by construction — both funnel through the
159    /// private [`Memtable::mutation_size`] computation, so there is no drift.
160    pub(crate) fn estimate_mutation_size(&self, m: &Mutation) -> usize {
161        Self::mutation_size(m)
162    }
163
164    /// Estimate the size of a mutation in bytes
165    ///
166    /// Conservative estimate includes:
167    /// - Fixed overhead per mutation (48 bytes for struct fields)
168    /// - Partition key size (key bytes)
169    /// - Clustering key size (if present)
170    /// - Cell operation sizes (column names + values)
171    fn mutation_size(mutation: &Mutation) -> usize {
172        // Base struct overhead. All accumulation uses `saturating_add` because
173        // `estimate_value_size` may return `usize::MAX` (fail-closed) for a
174        // pathological value (issue #1625); a plain `+` would panic under debug
175        // overflow checks.
176        let mut size: usize = 48;
177
178        // Partition key size
179        for (col_name, value) in &mutation.partition_key.columns {
180            size = size.saturating_add(col_name.len());
181            size = size.saturating_add(Self::estimate_value_size(value));
182        }
183
184        // Clustering key size
185        if let Some(ref clustering_key) = mutation.clustering_key {
186            for (col_name, value) in &clustering_key.columns {
187                size = size.saturating_add(col_name.len());
188                size = size.saturating_add(Self::estimate_value_size(value));
189            }
190        }
191
192        // Cell operations
193        for op in &mutation.operations {
194            size = size.saturating_add(Self::estimate_operation_size(op));
195        }
196
197        size
198    }
199
200    /// Estimate the size in bytes a CQL value would add to the memtable.
201    ///
202    /// Implemented as a BOUNDED ITERATIVE traversal (issue #1625): an explicit
203    /// worklist of `&Value` replaces the previous recursive estimator. Because
204    /// traversal state is not on the call stack there is NO stack-overflow risk,
205    /// so there is NO depth cap and therefore NO collapsing of deeply nested
206    /// children to a conservative floor — a large scalar buried arbitrarily deep
207    /// (e.g. `List([List([List([Text(128KB)])])])`) is counted at its real heap
208    /// size, closing the hard-limit bypass.
209    ///
210    /// The worklist is a stack-backed [`SmallVec`] with 32 inline slots, so the
211    /// common case (shallow/normal mutations — what the #1660 write-path
212    /// allocation budget test exercises) performs the whole traversal with ZERO
213    /// heap allocation on the admission/accounting hot path. Only pathologically
214    /// deep/wide values spill the worklist to the heap.
215    ///
216    /// Scalars contribute their real heap size (`Text`/`Blob`/`Varint`/`Inet`/
217    /// `Json`/`Decimal` byte length; fixed widths otherwise). Collections, maps,
218    /// UDTs, tuples and frozen values contribute their per-container overhead and
219    /// push their children (maps push keys AND values; UDTs count field names and
220    /// push field values) onto the worklist.
221    ///
222    /// To bound worst-case work on a pathologically huge/deep value, traversal
223    /// stops after visiting [`MAX_ESTIMATE_NODES`](Self::MAX_ESTIMATE_NODES) and
224    /// FAILS CLOSED, returning `usize::MAX` so admission REJECTS the value rather
225    /// than under-counting it. All accumulation uses `saturating_add`, so the
226    /// estimate can never wrap around a small value.
227    fn estimate_value_size(value: &crate::types::Value) -> usize {
228        use crate::types::Value;
229        use smallvec::SmallVec;
230
231        let mut total: usize = 0;
232        let mut visited: usize = 0;
233        // Stack-backed worklist of borrowed values still to be measured. The 32
234        // inline slots cover normal nesting/width, so shallow/normal mutations
235        // (the #1660 write-path allocation budget case) traverse with ZERO heap
236        // allocation; only pathological deep/wide values spill to the heap.
237        let mut worklist: SmallVec<[&Value; 32]> = SmallVec::new();
238        worklist.push(value);
239
240        // Would scheduling `incoming` more children push the total node count
241        // past the cap? `visited` (popped so far) + `pending` (already queued) +
242        // `incoming` is the upper bound on nodes this traversal will touch. This
243        // is checked BEFORE enqueuing so a single flat collection with far more
244        // than `MAX_ESTIMATE_NODES` elements can never grow the worklist
245        // proportional to its element count — the DoS guard fails closed WITHOUT
246        // the huge allocation (issue #1625).
247        let would_exceed_cap = |visited: usize, pending: usize, incoming: usize| -> bool {
248            visited.saturating_add(pending).saturating_add(incoming) > Self::MAX_ESTIMATE_NODES
249        };
250
251        while let Some(v) = worklist.pop() {
252            visited += 1;
253            if visited > Self::MAX_ESTIMATE_NODES {
254                // Pathological value: fail closed so admission rejects it.
255                return usize::MAX;
256            }
257
258            match v {
259                Value::Null => {}
260                Value::Boolean(_) | Value::TinyInt(_) => total = total.saturating_add(1),
261                Value::SmallInt(_) => total = total.saturating_add(2),
262                Value::Integer(_) | Value::Float32(_) | Value::Date(_) => {
263                    total = total.saturating_add(4)
264                }
265                Value::BigInt(_)
266                | Value::Counter(_)
267                | Value::Timestamp(_)
268                | Value::Time(_)
269                | Value::Float(_) => total = total.saturating_add(8),
270                Value::Uuid(_) | Value::Duration { .. } => total = total.saturating_add(16),
271                Value::Text(s) => total = total.saturating_add(s.len()),
272                Value::Blob(bytes) | Value::Varint(bytes) | Value::Inet(bytes) => {
273                    total = total.saturating_add(bytes.len())
274                }
275                Value::Decimal { scale: _, unscaled } => {
276                    total = total.saturating_add(4).saturating_add(unscaled.len())
277                }
278                Value::Json(json) => total = total.saturating_add(json.to_string().len()),
279                Value::Tombstone(_) => total = total.saturating_add(24),
280                Value::List(items) | Value::Set(items) | Value::Tuple(items) => {
281                    total = total.saturating_add(16);
282                    if would_exceed_cap(visited, worklist.len(), items.len()) {
283                        return usize::MAX;
284                    }
285                    worklist.extend(items.iter());
286                }
287                Value::Map(entries) => {
288                    total = total.saturating_add(16);
289                    // Each entry enqueues both a key and a value.
290                    let incoming = entries.len().saturating_mul(2);
291                    if would_exceed_cap(visited, worklist.len(), incoming) {
292                        return usize::MAX;
293                    }
294                    for (k, val) in entries {
295                        worklist.push(k);
296                        worklist.push(val);
297                    }
298                }
299                Value::Udt(udt) => {
300                    total = total.saturating_add(16);
301                    // Upper bound: at most one child per field (fields with a
302                    // value); check before touching any field so a wide UDT
303                    // cannot balloon the worklist.
304                    if would_exceed_cap(visited, worklist.len(), udt.fields.len()) {
305                        return usize::MAX;
306                    }
307                    for field in &udt.fields {
308                        total = total.saturating_add(field.name.len());
309                        if let Some(fv) = field.value.as_ref() {
310                            worklist.push(fv);
311                        }
312                    }
313                }
314                Value::Frozen(inner) => {
315                    total = total.saturating_add(8);
316                    if would_exceed_cap(visited, worklist.len(), 1) {
317                        return usize::MAX;
318                    }
319                    worklist.push(inner);
320                }
321            }
322        }
323
324        total
325    }
326
327    /// Estimate the size of a cell operation
328    fn estimate_operation_size(
329        op: &crate::storage::write_engine::mutation::CellOperation,
330    ) -> usize {
331        use crate::storage::write_engine::mutation::CellOperation;
332
333        // `saturating_add` throughout: `estimate_value_size` may return
334        // `usize::MAX` (fail-closed) for a pathological value (issue #1625).
335        match op {
336            CellOperation::Write { column, value } => column
337                .len()
338                .saturating_add(Self::estimate_value_size(value))
339                .saturating_add(8), // +8 for overhead
340            CellOperation::WriteWithTtl { column, value, .. } => {
341                // TTL cells: same as Write + 4 bytes for TTL + 4 bytes for local_deletion_time
342                column
343                    .len()
344                    .saturating_add(Self::estimate_value_size(value))
345                    .saturating_add(16)
346            }
347            CellOperation::Delete { column, .. } => column.len().saturating_add(8),
348            CellOperation::DeleteRow => 8,
349            // Epic #899: per-element complex ops. Each carries a column name,
350            // the preserved cell path, an optional value, and temporal metadata.
351            CellOperation::WriteComplexElement {
352                column,
353                cell_path,
354                value,
355                ..
356            } => column
357                .len()
358                .saturating_add(cell_path.len())
359                .saturating_add(value.as_ref().map(Self::estimate_value_size).unwrap_or(0))
360                .saturating_add(16), // flags + ts/ldt/ttl deltas + length prefixes overhead
361            CellOperation::ComplexDeletion { column, .. } => column.len().saturating_add(16),
362        }
363    }
364
365    /// Test-only: force the tracked approximate size, used to exercise the
366    /// admission gate's `saturating_add` overflow guard near `usize::MAX`
367    /// (issue #1625) — a size unreachable through real inserts.
368    #[cfg(test)]
369    pub(crate) fn set_size_bytes_for_test(&mut self, size: usize) {
370        self.size_bytes = size;
371    }
372
373    /// Get current timestamp in microseconds since Unix epoch
374    fn current_timestamp_micros() -> i64 {
375        std::time::SystemTime::now()
376            .duration_since(std::time::UNIX_EPOCH)
377            .unwrap_or_default()
378            .as_micros() as i64
379    }
380}
381
382impl Default for Memtable {
383    fn default() -> Self {
384        Self::new()
385    }
386}
387
388#[cfg(test)]
389mod tests {
390    use super::*;
391    use crate::storage::write_engine::mutation::{
392        CellOperation, ClusteringKey, PartitionKey, TableId,
393    };
394    use crate::types::Value;
395
396    fn create_test_mutation(
397        id: i32,
398        name: &str,
399        clustering_val: Option<i64>,
400    ) -> (DecoratedKey, Mutation) {
401        let table_id = TableId::new("test_ks", "test_table");
402        let partition_key = PartitionKey::single("id", Value::Integer(id));
403
404        // Calculate decorated key from partition key bytes
405        let key_bytes = id.to_be_bytes().to_vec();
406        let decorated_key = DecoratedKey::from_key_bytes(key_bytes).unwrap();
407
408        let clustering_key =
409            clustering_val.map(|val| ClusteringKey::single("ts", Value::BigInt(val)));
410
411        let operations = vec![CellOperation::Write {
412            column: "name".to_string(),
413            value: Value::text(name.to_string()),
414        }];
415
416        let mutation = Mutation::new(
417            table_id,
418            partition_key,
419            clustering_key,
420            operations,
421            1234567890,
422            None,
423        );
424
425        (decorated_key, mutation)
426    }
427
428    #[test]
429    fn test_memtable_new() {
430        let memtable = Memtable::new();
431        assert!(memtable.is_empty());
432        assert_eq!(memtable.size_bytes(), 0);
433        assert_eq!(memtable.row_count(), 0);
434        assert!(memtable.created_at() > 0);
435    }
436
437    #[test]
438    fn test_memtable_insert_and_get() {
439        let mut memtable = Memtable::new();
440
441        let (key, mutation) = create_test_mutation(1, "Alice", None);
442        memtable.insert_with_key(key.clone(), mutation).unwrap();
443
444        assert!(!memtable.is_empty());
445        assert_eq!(memtable.row_count(), 1);
446        assert!(memtable.size_bytes() > 0);
447
448        // Retrieve mutation
449        let mutations = memtable.get(&key).unwrap();
450        assert_eq!(mutations.len(), 1);
451        assert_eq!(mutations[0].table.table, "test_table");
452    }
453
454    #[test]
455    fn test_memtable_multiple_mutations_same_partition() {
456        let mut memtable = Memtable::new();
457
458        // Insert multiple mutations for same partition (different clustering keys)
459        let (key, mutation1) = create_test_mutation(1, "Alice", Some(1000));
460        let (_, mutation2) = create_test_mutation(1, "Alice Updated", Some(2000));
461
462        memtable.insert_with_key(key.clone(), mutation1).unwrap();
463        memtable.insert_with_key(key.clone(), mutation2).unwrap();
464
465        assert_eq!(memtable.row_count(), 2);
466
467        // Both mutations should be stored for this partition
468        let mutations = memtable.get(&key).unwrap();
469        assert_eq!(mutations.len(), 2);
470    }
471
472    #[test]
473    fn test_memtable_multiple_partitions() {
474        let mut memtable = Memtable::new();
475
476        let (key1, mutation1) = create_test_mutation(1, "Alice", None);
477        let (key2, mutation2) = create_test_mutation(2, "Bob", None);
478        let (key3, mutation3) = create_test_mutation(3, "Charlie", None);
479
480        memtable.insert_with_key(key1, mutation1).unwrap();
481        memtable.insert_with_key(key2, mutation2).unwrap();
482        memtable.insert_with_key(key3, mutation3).unwrap();
483
484        assert_eq!(memtable.row_count(), 3);
485        assert!(!memtable.is_empty());
486    }
487
488    #[test]
489    fn test_memtable_token_ordering() {
490        let mut memtable = Memtable::new();
491
492        // Insert in non-sorted order
493        let (key3, mutation3) = create_test_mutation(300, "Charlie", None);
494        let (key1, mutation1) = create_test_mutation(100, "Alice", None);
495        let (key2, mutation2) = create_test_mutation(200, "Bob", None);
496
497        memtable.insert_with_key(key3.clone(), mutation3).unwrap();
498        memtable.insert_with_key(key1.clone(), mutation1).unwrap();
499        memtable.insert_with_key(key2.clone(), mutation2).unwrap();
500
501        // Verify iteration returns partitions in token order
502        let keys: Vec<_> = memtable.iter().map(|(k, _)| k.token).collect();
503        assert_eq!(keys.len(), 3);
504
505        // Keys should be sorted by token
506        assert!(keys.windows(2).all(|w| w[0] <= w[1]));
507    }
508
509    #[test]
510    fn test_memtable_size_tracking() {
511        let mut memtable = Memtable::new();
512
513        let initial_size = memtable.size_bytes();
514        assert_eq!(initial_size, 0);
515
516        // Insert mutation
517        let (key, mutation) = create_test_mutation(1, "Alice", None);
518        memtable.insert_with_key(key, mutation).unwrap();
519
520        // Size should increase
521        assert!(memtable.size_bytes() > initial_size);
522        let size_after_insert = memtable.size_bytes();
523
524        // Insert another mutation - size should increase more
525        let (key2, mutation2) = create_test_mutation(2, "Bob with a longer name", None);
526        memtable.insert_with_key(key2, mutation2).unwrap();
527
528        assert!(memtable.size_bytes() > size_after_insert);
529    }
530
531    #[test]
532    fn test_memtable_should_flush() {
533        let mut memtable = Memtable::new();
534
535        // Should not flush when empty
536        assert!(!memtable.should_flush(1024));
537
538        // Insert mutations until threshold
539        for i in 0..100 {
540            let (key, mutation) = create_test_mutation(i, "Test data", None);
541            memtable.insert_with_key(key, mutation).unwrap();
542        }
543
544        // Should flush if size exceeds threshold
545        let current_size = memtable.size_bytes();
546        assert!(memtable.should_flush(current_size - 1));
547        assert!(!memtable.should_flush(current_size + 1000));
548    }
549
550    #[test]
551    fn test_memtable_clear() {
552        let mut memtable = Memtable::new();
553
554        let created_at = memtable.created_at();
555
556        // Insert some data
557        let (key, mutation) = create_test_mutation(1, "Alice", None);
558        memtable.insert_with_key(key, mutation).unwrap();
559
560        assert!(!memtable.is_empty());
561        assert!(memtable.size_bytes() > 0);
562        assert!(memtable.row_count() > 0);
563
564        // Clear
565        memtable.clear();
566
567        assert!(memtable.is_empty());
568        assert_eq!(memtable.size_bytes(), 0);
569        assert_eq!(memtable.row_count(), 0);
570        assert_eq!(memtable.created_at(), created_at); // Timestamp unchanged
571    }
572
573    #[test]
574    fn test_memtable_iterator() {
575        let mut memtable = Memtable::new();
576
577        // Insert multiple partitions
578        let (key1, mutation1) = create_test_mutation(1, "Alice", None);
579        let (key2, mutation2) = create_test_mutation(2, "Bob", None);
580
581        memtable.insert_with_key(key1.clone(), mutation1).unwrap();
582        memtable.insert_with_key(key2.clone(), mutation2).unwrap();
583
584        // Iterate and verify
585        let mut count = 0;
586        for (key, mutations) in memtable.iter() {
587            assert!(!mutations.is_empty());
588            assert!([key1.token, key2.token].contains(&key.token));
589            count += 1;
590        }
591
592        assert_eq!(count, 2);
593    }
594
595    #[test]
596    fn test_memtable_empty_check() {
597        let mut memtable = Memtable::new();
598        assert!(memtable.is_empty());
599
600        let (key, mutation) = create_test_mutation(1, "Alice", None);
601        memtable.insert_with_key(key, mutation).unwrap();
602        assert!(!memtable.is_empty());
603
604        memtable.clear();
605        assert!(memtable.is_empty());
606    }
607
608    #[test]
609    fn test_memtable_size_estimates() {
610        // Test size estimation for different value types
611        let small_text = Value::text("hi".to_string());
612        let large_text = Value::text("a".repeat(1000));
613        let integer = Value::Integer(42);
614        let uuid = Value::Uuid([0u8; 16]);
615
616        assert_eq!(Memtable::estimate_value_size(&small_text), 2);
617        assert_eq!(Memtable::estimate_value_size(&large_text), 1000);
618        assert_eq!(Memtable::estimate_value_size(&integer), 4);
619        assert_eq!(Memtable::estimate_value_size(&uuid), 16);
620    }
621
622    #[test]
623    fn test_memtable_collection_size_estimates() {
624        // List
625        let list = Value::List(vec![
626            Value::Integer(1),
627            Value::Integer(2),
628            Value::Integer(3),
629        ]);
630        let size = Memtable::estimate_value_size(&list);
631        assert!(size >= 12); // 3 * 4 bytes + overhead
632
633        // Set
634        let set = Value::Set(vec![
635            Value::text("a".to_string()),
636            Value::text("b".to_string()),
637        ]);
638        let size = Memtable::estimate_value_size(&set);
639        assert!(size >= 2); // 2 * 1 byte + overhead
640
641        // Map
642        let map = Value::Map(vec![
643            (Value::Integer(1), Value::text("one".to_string())),
644            (Value::Integer(2), Value::text("two".to_string())),
645        ]);
646        let size = Memtable::estimate_value_size(&map);
647        assert!(size >= 11); // 2 * (4 + 3) bytes + overhead
648    }
649
650    #[test]
651    fn test_memtable_realistic_flush_threshold() {
652        let mut memtable = Memtable::new();
653
654        // Target: ~10K mutations before 64MB flush (conservative estimate)
655        // Average mutation size should be < 6.4KB
656        let flush_threshold = 64 * 1024 * 1024; // 64MB
657
658        // Insert 10K typical mutations
659        for i in 0..10_000 {
660            let (key, mutation) = create_test_mutation(
661                i,
662                "Typical user data with moderate length name",
663                Some(i as i64),
664            );
665            memtable.insert_with_key(key, mutation).unwrap();
666        }
667
668        let final_size = memtable.size_bytes();
669        println!(
670            "10K mutations size: {} bytes ({} KB)",
671            final_size,
672            final_size / 1024
673        );
674
675        // Should be well under 64MB for 10K mutations
676        assert!(final_size < flush_threshold);
677
678        // Verify avg size per mutation is reasonable
679        let avg_size = final_size / 10_000;
680        println!("Average mutation size: {} bytes", avg_size);
681        assert!(avg_size > 0);
682        assert!(avg_size < 10_000); // Should be less than 10KB per mutation
683    }
684
685    #[test]
686    fn test_memtable_get_nonexistent_key() {
687        let memtable = Memtable::new();
688        let key = DecoratedKey::new(12345, vec![0, 0, 0, 99]);
689
690        assert!(memtable.get(&key).is_none());
691    }
692
693    #[test]
694    fn test_memtable_insert_deprecated_api() {
695        let mut memtable = Memtable::new();
696
697        let table_id = TableId::new("test_ks", "test_table");
698        let partition_key = PartitionKey::single("id", Value::Integer(1));
699        let operations = vec![CellOperation::Write {
700            column: "name".to_string(),
701            value: Value::text("Alice".to_string()),
702        }];
703
704        let mutation = Mutation::new(table_id, partition_key, None, operations, 1234567890, None);
705
706        // Deprecated insert() should return error
707        let result = memtable.insert(mutation);
708        assert!(result.is_err());
709    }
710
711    #[test]
712    fn test_memtable_nested_collection_depth_limit() {
713        // Issue #1625: a deeply nested list wrapping a tiny scalar is handled by
714        // the ITERATIVE estimator without stack overflow and counted ACCURATELY
715        // (no depth cap, no floor). 40 lists (16 each) around Integer(42) (4).
716        let mut nested_value = Value::Integer(42);
717        for _ in 0..40 {
718            nested_value = Value::List(vec![nested_value]);
719        }
720
721        let size = Memtable::estimate_value_size(&nested_value);
722        assert_eq!(
723            size,
724            40 * 16 + 4,
725            "deep list of a tiny scalar must be counted accurately, not floored"
726        );
727    }
728
729    #[test]
730    fn test_memtable_nested_map_depth_limit() {
731        // Issue #1625: deep map nesting counted accurately. 35 maps, each +16
732        // overhead + Integer key (4), innermost Text("bottom") (6). No floor.
733        let mut nested_value = Value::text("bottom".to_string());
734        for _ in 0..35 {
735            nested_value = Value::Map(vec![(Value::Integer(1), nested_value)]);
736        }
737
738        let size = Memtable::estimate_value_size(&nested_value);
739        assert_eq!(
740            size,
741            35 * (16 + 4) + 6,
742            "deep map of a tiny scalar must be counted accurately, not floored"
743        );
744    }
745
746    #[test]
747    fn test_memtable_nested_udt_depth_limit() {
748        use crate::types::{UdtField, UdtValue};
749
750        // Issue #1625: deep UDT nesting counted accurately. 35 UDTs, each +16
751        // overhead + field name "field" (5), innermost Integer(1) (4). No floor.
752        let mut nested_value = Value::Integer(1);
753        for i in 0..35 {
754            let udt = UdtValue {
755                type_name: format!("type_{}", i),
756                keyspace: "test_ks".to_string(),
757                fields: vec![UdtField {
758                    name: "field".to_string(),
759                    value: Some(nested_value),
760                }],
761            };
762            nested_value = Value::Udt(Box::new(udt));
763        }
764
765        let size = Memtable::estimate_value_size(&nested_value);
766        assert_eq!(
767            size,
768            35 * (16 + 5) + 4,
769            "deep UDT of a tiny scalar must be counted accurately, not floored"
770        );
771    }
772
773    #[test]
774    fn test_memtable_frozen_nested_depth_limit() {
775        // Issue #1625: deep Frozen nesting must NOT stack-overflow (iterative)
776        // and is counted accurately. 40 frozen (8 each) around Integer(99) (4).
777        let mut nested_value = Value::Integer(99);
778        for _ in 0..40 {
779            nested_value = Value::Frozen(Box::new(nested_value));
780        }
781
782        let size = Memtable::estimate_value_size(&nested_value);
783        assert_eq!(
784            size,
785            40 * 8 + 4,
786            "deep frozen of a tiny scalar must be counted accurately, not floored"
787        );
788    }
789
790    #[test]
791    fn test_memtable_mixed_nested_collections() {
792        use crate::types::{UdtField, UdtValue};
793
794        // Create a complex nested structure mixing different types
795        let mut nested_value = Value::text("base".to_string());
796
797        // Alternate between different collection types
798        for i in 0..50 {
799            nested_value = match i % 5 {
800                0 => Value::List(vec![nested_value]),
801                1 => Value::Set(vec![nested_value]),
802                2 => Value::Map(vec![(Value::Integer(i), nested_value)]),
803                3 => Value::Tuple(vec![nested_value]),
804                4 => Value::Udt(Box::new(UdtValue {
805                    type_name: format!("type_{}", i),
806                    keyspace: "test_ks".to_string(),
807                    fields: vec![UdtField {
808                        name: "f".to_string(),
809                        value: Some(nested_value),
810                    }],
811                })),
812                _ => unreachable!(),
813            };
814        }
815
816        // Should handle mixed nesting without panic
817        let size = Memtable::estimate_value_size(&nested_value);
818        assert!(size > 0);
819    }
820
821    #[test]
822    fn test_memtable_depth_limit_exact_boundary() {
823        // Issue #1625: with the iterative estimator there is no depth cap, so
824        // adding a wrapper level increases the estimate by exactly one List's
825        // overhead (16) — no discontinuity/jump to a floor at any depth.
826        let mut nested_value = Value::Integer(1);
827        for _ in 0..32 {
828            nested_value = Value::List(vec![nested_value]);
829        }
830
831        let size = Memtable::estimate_value_size(&nested_value);
832        assert_eq!(size, 32 * 16 + 4);
833
834        // One more wrapper: exactly +16, not a jump to any conservative floor.
835        nested_value = Value::List(vec![nested_value]);
836        let size_over = Memtable::estimate_value_size(&nested_value);
837        assert_eq!(size_over, size + 16);
838    }
839
840    #[test]
841    fn test_estimate_mutation_size_matches_insert_accounting() {
842        // Issue #1625: the admission gate and the running size accounting must
843        // agree by construction — `estimate_mutation_size(&m)` must equal the
844        // delta `size_bytes()` gains from inserting the same mutation.
845        let mut memtable = Memtable::new();
846        let (key, mutation) = create_test_mutation(7, "some data here", Some(42));
847
848        let before = memtable.size_bytes();
849        let predicted = memtable.estimate_mutation_size(&mutation);
850        memtable.insert_with_key(key, mutation).unwrap();
851        let actual_delta = memtable.size_bytes() - before;
852
853        assert_eq!(
854            predicted, actual_delta,
855            "estimate_mutation_size must equal the size delta insert applies"
856        );
857        assert!(predicted > 0);
858    }
859
860    #[test]
861    fn test_deep_wide_collection_counted_accurately() {
862        // Issue #1625: a WIDE collection buried under many wrapper lists must be
863        // counted at its REAL byte size regardless of depth. The iterative
864        // estimator has no depth cap, so the 500 × 200-byte strings (~100KB) at
865        // the bottom are summed exactly — not collapsed to any floor.
866        let wide = Value::List((0..500).map(|_| Value::text("y".repeat(200))).collect());
867        let mut nested = wide;
868        for _ in 0..32 {
869            nested = Value::List(vec![nested]);
870        }
871
872        let size = Memtable::estimate_value_size(&nested);
873        let real_payload = 500 * 200; // 100_000 bytes of text
874        assert!(
875            size >= real_payload,
876            "wide collection must count real bytes at any depth (got {size})"
877        );
878        // Accurate, not wildly inflated: real payload + bounded per-container
879        // overhead only.
880        assert!(
881            size < real_payload + 10 * 1024,
882            "estimate must be accurate, not floor-inflated (got {size})"
883        );
884    }
885
886    #[test]
887    fn test_deep_narrow_collection_large_scalar_not_undercounted() {
888        // Issue #1625 (roborev finding): a deep NARROW collection wrapping a
889        // large DIRECT scalar must not be systematically under-counted. Wrap a
890        // `List([Text(128KB)])` in 32 single-element lists.
891        //
892        // Pre-fix: any node past the depth cap collapsed to a ~1KB floor — small
893        // enough to slip a 64KB gate. Post-fix: the iterative estimator counts
894        // the 128KB scalar at its real heap size regardless of depth.
895        let big = 128 * 1024;
896        let mut nested = Value::List(vec![Value::text("x".repeat(big))]);
897        for _ in 0..32 {
898            nested = Value::List(vec![nested]);
899        }
900
901        let size = Memtable::estimate_value_size(&nested);
902        assert!(
903            size >= big,
904            "deep narrow collection with a large scalar must count the scalar's \
905             real heap size, not the old ~1KB floor (got {size}, expected >= {big})"
906        );
907    }
908
909    #[test]
910    fn test_large_scalar_buried_below_old_cap_counted() {
911        // Issue #1625 (3rd iteration): a large scalar buried MANY levels below
912        // where the old depth cap (32) sat must be counted at real size. Wrap a
913        // 128KB text in 40 single-element lists — well past the old cap — and
914        // confirm the iterative estimator still sees the full 128KB.
915        let big = 128 * 1024;
916        let mut nested = Value::text("z".repeat(big));
917        for _ in 0..40 {
918            nested = Value::List(vec![nested]);
919        }
920
921        let size = Memtable::estimate_value_size(&nested);
922        assert!(
923            size >= big,
924            "scalar buried below the old depth cap must be counted at real size \
925             (got {size}, expected >= {big})"
926        );
927    }
928
929    #[test]
930    fn test_pathological_node_cap_returns_usize_max() {
931        // Issue #1625: exceeding MAX_ESTIMATE_NODES must FAIL CLOSED (usize::MAX),
932        // never under-count — and never overflow or hang. A flat list of
933        // 1_000_001 integers is a single fast allocation that trips the cap.
934        let value = Value::List((0..1_000_001i32).map(Value::Integer).collect());
935        let size = Memtable::estimate_value_size(&value);
936        assert_eq!(
937            size,
938            usize::MAX,
939            "hitting the node cap must fail closed with usize::MAX (got {size})"
940        );
941    }
942
943    #[test]
944    fn test_insert_with_pathological_value_saturates_ledger_no_panic() {
945        // Issue #1625 (roborev finding 1): a direct `Memtable` user bypasses the
946        // WriteEngine admission gate, so `insert_with_key` MUST be self-safe when
947        // the estimator returns `usize::MAX` (node cap fail-closed). The ledger
948        // update must saturate — never panic (debug overflow) or wrap (release).
949        let mut memtable = Memtable::new();
950
951        // A flat list past the node cap makes `mutation_size` == usize::MAX.
952        let pathological = Value::List(
953            (0..(Memtable::MAX_ESTIMATE_NODES as i32 + 5))
954                .map(Value::Integer)
955                .collect(),
956        );
957        let table_id = TableId::new("test_ks", "test_table");
958        let partition_key = PartitionKey::single("id", Value::Integer(1));
959        let key = DecoratedKey::from_key_bytes(1i32.to_be_bytes().to_vec()).unwrap();
960        let operations = vec![CellOperation::Write {
961            column: "big".to_string(),
962            value: pathological,
963        }];
964        let mutation = Mutation::new(table_id, partition_key, None, operations, 1, None);
965
966        // Sanity: the estimate really is usize::MAX for this mutation.
967        assert_eq!(memtable.estimate_mutation_size(&mutation), usize::MAX);
968
969        // Must not panic; ledger saturates at usize::MAX (no wrap to a small value).
970        memtable.insert_with_key(key, mutation).unwrap();
971        assert_eq!(memtable.size_bytes(), usize::MAX);
972    }
973
974    #[test]
975    fn test_insert_saturates_when_ledger_already_near_max() {
976        // Issue #1625 (roborev finding 1): incrementing an already-huge ledger by
977        // a normal mutation size must saturate, not wrap around to a small value.
978        let mut memtable = Memtable::new();
979        memtable.set_size_bytes_for_test(usize::MAX - 3);
980
981        let (key, mutation) = create_test_mutation(1, "Alice", None);
982        memtable.insert_with_key(key, mutation).unwrap();
983
984        assert_eq!(
985            memtable.size_bytes(),
986            usize::MAX,
987            "ledger must saturate at usize::MAX, never wrap"
988        );
989    }
990
991    #[test]
992    fn test_wide_collection_fails_closed_before_enqueuing_children() {
993        // Issue #1625 (roborev finding 2): a single flat collection whose element
994        // count exceeds MAX_ESTIMATE_NODES must fail closed (usize::MAX) at the
995        // ENQUEUE check — the worklist is never grown proportional to the element
996        // count. Verified for every enqueue site (list, map, UDT, frozen).
997        use crate::types::{UdtField, UdtValue};
998
999        let over = Memtable::MAX_ESTIMATE_NODES + 5;
1000
1001        // List / Set / Tuple enqueue site.
1002        let list = Value::List((0..over as i32).map(Value::Integer).collect());
1003        assert_eq!(Memtable::estimate_value_size(&list), usize::MAX);
1004
1005        // Map enqueue site (keys + values).
1006        let map = Value::Map(
1007            (0..over as i32)
1008                .map(|i| (Value::Integer(i), Value::Integer(i)))
1009                .collect(),
1010        );
1011        assert_eq!(Memtable::estimate_value_size(&map), usize::MAX);
1012
1013        // UDT enqueue site (field values).
1014        let udt = Value::Udt(Box::new(UdtValue {
1015            type_name: "t".to_string(),
1016            keyspace: "ks".to_string(),
1017            fields: (0..over)
1018                .map(|i| UdtField {
1019                    name: String::new(),
1020                    value: Some(Value::Integer(i as i32)),
1021                })
1022                .collect(),
1023        }));
1024        assert_eq!(Memtable::estimate_value_size(&udt), usize::MAX);
1025    }
1026
1027    #[test]
1028    fn test_memtable_shallow_collections_unaffected() {
1029        // Verify shallow collections are not affected by depth limit
1030
1031        // Simple list
1032        let simple_list = Value::List(vec![
1033            Value::Integer(1),
1034            Value::Integer(2),
1035            Value::Integer(3),
1036        ]);
1037        let size = Memtable::estimate_value_size(&simple_list);
1038        assert_eq!(size, 12 + 16); // 3 * 4 bytes + overhead
1039
1040        // Nested but shallow (3 levels)
1041        let shallow_nested =
1042            Value::List(vec![Value::List(vec![Value::List(vec![Value::Integer(
1043                1,
1044            )])])]);
1045        let size = Memtable::estimate_value_size(&shallow_nested);
1046        assert!(size > 0);
1047        assert!(size < 1024); // Should not use conservative estimate
1048    }
1049}