fix(zero-rust): keep every replica row findable after a B+tree merge (#177)

When a delete emptied an internal node's first child, the next child became
the first with its own, higher lower bound. Inside the node that bound is
never read, but once the node merged into its left neighbor it was, and
lookups of keys between the two bounds went to the neighbor's last child:
the replica held a row that get and has missed. A mutator or a poke then
couldn't delete the row, and putting it again panicked in MemorySource
("Row already exists"). JS deletes and re-inserts the row. A merge now
gives the right node's first child the bound the parent had for it.

Tests: tests/replica_tree.rs reproduces it through mutators and through
pokes; a KvMap fuzz of range deletes and refills, on nodes of 8 in unit
tests, checks that every key is findable and within the bounds lookups use.
Conformance unchanged: 63,453 cases, 0 failures.

Amp-Thread-ID: https://ampcode.com/threads/T-319b90a3-6ed9-4da2-8727-8dc76a1cf67a

Co-authored-by: Amp <amp@ampcode.com>

Author: tmm

.changeset/zero-rust-replica-tree.md

@@ -0,0 +1,5 @@
+---
+'zero-rust': patch
+---
+
+Fix the replica losing rows: after a delete emptied the first child of a node of its B+tree and the node later merged into its neighbor, the client couldn't find a row it still held. A mutator or a poke then couldn't delete the row, and putting it again panicked ("Row already exists").

rust/zero-client/src/client/replica_store.rs

@@ -8,8 +8,12 @@ use std::cmp::Ordering;
 use std::collections::{BTreeMap, HashMap, HashSet};
 use std::rc::Rc;
 
-/// The most entries or children a node of a [`KvMap`] holds.
+/// The most entries or children a node of a [`KvMap`] holds. Unit tests
+/// use small nodes, so a few thousand keys make deep trees and many merges.
+#[cfg(not(test))]
 const MAX_NODE: usize = 64;
+#[cfg(test)]
+const MAX_NODE: usize = 8;
 /// A node with fewer after a delete merges with a neighbor that has room.
 const MIN_NODE: usize = MAX_NODE / 4;
 
@@ -112,11 +116,17 @@ impl Node {
                 } else if len < MIN_NODE && children.len() > 1 {
                     let left = if i + 1 < children.len() { i } else { i - 1 };
                     if children[left].1.len() + children[left + 1].1.len() <= MAX_NODE {
-                        let (_, right) = children.remove(left + 1);
+                        let (bound, right) = children.remove(left + 1);
                         let right = Rc::unwrap_or_clone(right);
                         match (Rc::make_mut(&mut children[left].1), right) {
                             (Node::Leaf(a), Node::Leaf(b)) => a.extend(b),
-                            (Node::Internal(a), Node::Internal(b)) => a.extend(b),
+                            // The first child's bound goes unused until now,
+                            // and is stale once the child before it was
+                            // removed: it gets the bound the right node had.
+                            (Node::Internal(a), Node::Internal(mut b)) => {
+                                b[0].0 = bound;
+                                a.extend(b);
+                            }
                             _ => unreachable!("the children of a node are at one depth"),
                         }
                     }
@@ -717,4 +727,89 @@ mod tests {
             assert!(copy.iter().eq(model.iter()));
         }
     }
+
+    /// Checks that each node's keys are within the bounds lookups use: at
+    /// or above its bound (a first child's, its parent's), below the next.
+    fn check_bounds(map: &KvMap) {
+        fn walk(node: &Node, low: &str, high: Option<&str>) {
+            match node {
+                Node::Leaf(entries) => {
+                    for e in entries {
+                        assert!(e.0.as_str() >= low, "{} below the bound {low}", e.0);
+                        assert!(
+                            high.is_none_or(|h| e.0.as_str() < h),
+                            "{} at or above the next bound",
+                            e.0
+                        );
+                    }
+                }
+                Node::Internal(children) => {
+                    for (i, (bound, child)) in children.iter().enumerate() {
+                        // A first child's own bound is never read.
+                        let low = if i == 0 { low } else { bound.as_str() };
+                        let high = children.get(i + 1).map(|(b, _)| b.as_str()).or(high);
+                        walk(child, low, high);
+                    }
+                }
+            }
+        }
+        if let Some(root) = &map.root {
+            walk(root, "", None);
+        }
+    }
+
+    /// Range deletes, which empty whole nodes, and refills of the gaps
+    /// they leave, at three levels: every key the map iterates is found.
+    /// A merged node once kept a stale bound and lost keys this way.
+    #[test]
+    fn kv_map_range_deletes_keep_every_key_findable() {
+        const SEEDS: u64 = 20;
+        const SPACE: u64 = 2000;
+        const ROUNDS: usize = 200;
+        const WIDTH: u64 = 200;
+        for seed in 1..=SEEDS {
+            let mut state = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1;
+            let mut random = |n: u64| {
+                state ^= state << 13;
+                state ^= state >> 7;
+                state ^= state << 17;
+                state % n
+            };
+            let key = |i: u64| format!("e/{i:06}");
+            let mut map = KvMap::default();
+            let mut model = BTreeMap::new();
+            for i in 0..SPACE {
+                if random(3) > 0 {
+                    map.set(key(i * 10), Value::Null);
+                    model.insert(key(i * 10), Value::Null);
+                }
+            }
+            for round in 0..ROUNDS {
+                if random(3) == 0 {
+                    let start = random(SPACE) * 10;
+                    let end = start + random(WIDTH) * 10;
+                    let keys: Vec<String> = model
+                        .range(key(start)..key(end))
+                        .map(|(k, _)| k.clone())
+                        .collect();
+                    for k in keys {
+                        assert!(map.delete(&k), "seed {seed} round {round}: delete {k}");
+                        model.remove(&k);
+                    }
+                } else {
+                    let base = random(SPACE) * 10;
+                    for j in 0..random(40) {
+                        map.set(key(base + j), Value::Null);
+                        model.insert(key(base + j), Value::Null);
+                    }
+                }
+                assert_eq!(map.len, model.len(), "seed {seed} round {round}");
+                assert!(map.iter().eq(model.iter()), "seed {seed} round {round}");
+                check_bounds(&map);
+                for k in model.keys() {
+                    assert!(map.has(k), "seed {seed} round {round}: lost {k}");
+                }
+            }
+        }
+    }
 }

rust/zero-client/tests/replica_tree.rs

@@ -0,0 +1,233 @@
+//! The replica's B+tree (`KvMap` in `client/replica_store.rs`) once lost
+//! keys. When a delete emptied the first child of an internal node, the child
+//! was removed and the next one became the first, keeping its own lower
+//! bound, above the bound the parent had for the node. Inside the node the
+//! first bound is never read, so a key between the two went to that child.
+//! Once the node merged into its left neighbor, that bound was read: lookups
+//! of such a key went to the neighbor's last child, so `get` and `has` missed
+//! a row the map still held. A mutator or a poke then couldn't delete the
+//! row, and putting it again added a second entry and an add of a row the
+//! IVM source had, which panicked ("Row already exists").
+//!
+//! The steps are the shortest construction found: ascending inserts make
+//! leaves of 32 and a root over two internal nodes, P0 and P1. JS runs the
+//! same mutations (through the conformance harness's `harden.ops`) and
+//! deletes the row, then shows it again after the insert.
+
+use std::cell::Cell;
+use std::rc::Rc;
+use std::sync::Arc;
+use zero_client::ast::System;
+use zero_client::client::{
+    MaterializedView, Mutator, SocketFactory, Transaction, Zero, ZeroOptions,
+};
+use zero_client::json::{self, Object, Value};
+use zero_client::query::AnyQuery;
+use zero_client::schema::Schema;
+use zero_client::testing::{ManualScheduler, MockSocketFactory, RecordingLogSink};
+use zero_client::ttl::Ttl;
+
+fn schema() -> Arc<Schema> {
+    let schema = json::parse(
+        r#"{"tables": {"issue": {"name": "issue", "columns": {
+            "id": {"type": "string"}, "title": {"type": "string"}}, "primaryKey": ["id"]}}}"#,
+    )
+    .unwrap();
+    Arc::new(Schema::from_json(&schema).unwrap())
+}
+
+fn id(i: u32) -> String {
+    format!("i{i:06}")
+}
+
+/// `ids` as the mutator's arguments.
+fn args(ids: impl IntoIterator<Item = u32>) -> Option<Value> {
+    Some(Value::Array(
+        ids.into_iter().map(|i| Value::from(id(i))).collect(),
+    ))
+}
+
+struct Client {
+    zero: Zero,
+    scheduler: Rc<ManualScheduler>,
+}
+
+impl Client {
+    /// A local client with a view of every issue.
+    fn new() -> (Client, MaterializedView) {
+        let scheduler = ManualScheduler::new(1_700_000_000_000.0);
+        let insert = Mutator::new("insert", |tx: &Transaction, args, _: &()| {
+            for v in args.unwrap().as_array().unwrap() {
+                let mut row = Object::new();
+                row.insert("id", v.clone());
+                row.insert("title", Value::from("t"));
+                tx.table("issue").insert(&row)?;
+            }
+            Ok(None)
+        });
+        let delete = Mutator::new("delete", |tx: &Transaction, args, _: &()| {
+            for v in args.unwrap().as_array().unwrap() {
+                let mut row = Object::new();
+                row.insert("id", v.clone());
+                tx.table("issue").delete(&row)?;
+            }
+            Ok(None)
+        });
+        let zero = Zero::new(
+            schema(),
+            (),
+            vec![insert, delete],
+            ZeroOptions {
+                log_sink: Some(Rc::new(RecordingLogSink::default())),
+                ..ZeroOptions::default()
+            },
+            scheduler.clone(),
+            None,
+        )
+        .unwrap();
+        scheduler.run_checkpoints();
+        let view = zero
+            .materialize(
+                &AnyQuery::new(schema(), "issue", System::Client),
+                &Ttl::default(),
+            )
+            .unwrap();
+        scheduler.run_checkpoints();
+        (Client { zero, scheduler }, view)
+    }
+
+    fn mutate(&self, name: &str, ids: impl IntoIterator<Item = u32>) {
+        self.zero.mutate(name, args(ids));
+        self.scheduler.run_checkpoints();
+    }
+
+    /// Leaves `i010300` where the B+tree couldn't find it.
+    fn lose_a_row(&self) {
+        // Leaves of 32 keys; P1's first leaf has 10240..=10550, its second
+        // 10560..=10870.
+        self.mutate("insert", (0..2081).map(|i| i * 10));
+        // The second leaf full (64), so the first can't merge into it.
+        self.mutate("insert", (0..32).map(|i| 10561 + i * 10));
+        // The first leaf emptied and removed: the second, bound 10560, is
+        // P1's first child, and P1's bound is 10240.
+        self.mutate("delete", (0..32).map(|i| 10240 + i * 10));
+        self.mutate("insert", [10300]);
+        // P1 shrunk until it merges into P0.
+        self.mutate("delete", (0..563).map(|i| 20800 - i * 10));
+    }
+}
+
+fn has(view: &MaterializedView, i: u32) -> bool {
+    view.data()
+        .entries()
+        .iter()
+        .any(|e| e.row.get("id").and_then(Value::as_str) == Some(&id(i)))
+}
+
+#[test]
+fn a_row_after_a_merge_of_the_replica_tree_can_be_deleted() {
+    let (client, view) = Client::new();
+    client.lose_a_row();
+    assert!(has(&view, 10300));
+    client.mutate("delete", [10300]);
+    assert!(!has(&view, 10300), "the deleted row is still in the view");
+}
+
+#[test]
+fn a_row_after_a_merge_of_the_replica_tree_can_be_inserted_again() {
+    let (client, view) = Client::new();
+    client.lose_a_row();
+    client.mutate("delete", [10300]);
+    client.mutate("insert", [10300]);
+    assert!(has(&view, 10300));
+}
+
+/// The steps of [`Client::lose_a_row`] as pokes: puts and dels of the
+/// server's rows, then a del and a put of the row.
+#[test]
+fn a_row_after_a_merge_of_the_replica_tree_follows_the_server() {
+    let scheduler = ManualScheduler::new(1_700_000_000_000.0);
+    let factory = MockSocketFactory::new();
+    let zero = Zero::new(
+        schema(),
+        (),
+        Vec::new(),
+        ZeroOptions {
+            cache_url: Some("http://zero.test".into()),
+            user_id: Some("u1".into()),
+            log_sink: Some(Rc::new(RecordingLogSink::default())),
+            ..ZeroOptions::default()
+        },
+        scheduler.clone(),
+        Some(factory.clone() as Rc<dyn SocketFactory>),
+    )
+    .unwrap();
+    let view = zero
+        .materialize(
+            &AnyQuery::new(schema(), "issue", System::Client),
+            &Ttl::default(),
+        )
+        .unwrap();
+    scheduler.advance(0.0);
+    let socket = factory.sockets()[0].clone();
+    socket.receive_text(r#"["connected",{"wsid":"w1"}]"#);
+    let cookie = Cell::new(0);
+    let poke = |patch: Vec<String>| {
+        let n = cookie.get() + 1;
+        let base = if n == 1 {
+            "null".to_owned()
+        } else {
+            format!(r#""{:02}""#, n - 1)
+        };
+        cookie.set(n);
+        socket.receive_text(&format!(
+            r#"["pokeStart",{{"pokeID":"p{n}","baseCookie":{base}}}]"#
+        ));
+        socket.receive_text(&format!(
+            r#"["pokePart",{{"pokeID":"p{n}","rowsPatch":[{}]}}]"#,
+            patch.join(",")
+        ));
+        socket.receive_text(&format!(
+            r#"["pokeEnd",{{"pokeID":"p{n}","cookie":"{n:02}"}}]"#
+        ));
+        scheduler.advance(20.0);
+    };
+    let put = |ids: Vec<u32>| {
+        poke(
+            ids.into_iter()
+                .map(|i| {
+                    format!(
+                        r#"{{"op":"put","tableName":"issue","value":{{"id":"{}","title":"t"}}}}"#,
+                        id(i)
+                    )
+                })
+                .collect(),
+        )
+    };
+    let del = |ids: Vec<u32>| {
+        poke(
+            ids.into_iter()
+                .map(|i| {
+                    format!(
+                        r#"{{"op":"del","tableName":"issue","id":{{"id":"{}"}}}}"#,
+                        id(i)
+                    )
+                })
+                .collect(),
+        )
+    };
+    put((0..2081).map(|i| i * 10).collect());
+    put((0..32).map(|i| 10561 + i * 10).collect());
+    del((0..32).map(|i| 10240 + i * 10).collect());
+    put(vec![10300]);
+    del((0..563).map(|i| 20800 - i * 10).collect());
+    assert!(has(&view, 10300));
+    del(vec![10300]);
+    assert!(
+        !has(&view, 10300),
+        "the row the server deleted is still in the view"
+    );
+    put(vec![10300]);
+    assert!(has(&view, 10300));
+    assert_eq!(view.data().entries().len(), 2081 + 32 - 32 + 1 - 563);
+}