Fix SDK bug with overlapping subscriptions

crates/codegen/src/rust.rs

@@ -1000,7 +1000,7 @@ impl TryFrom<__ws::DatabaseUpdate<__ws::BsatnFormat>> for DbUpdate {
             for table in iter_tables(module) {
                 writeln!(
                     out,
-                    "{:?} => db_update.{} = {}::parse_table_update(table_update)?,",
+                    "{:?} => db_update.{}.append({}::parse_table_update(table_update)?),",
                     table.name.deref(),
                     table_method_name(&table.name),
                     table_module_name(&table.name),

crates/codegen/tests/snapshots/codegen__codegen_rust.snap

@@ -1,5 +1,5 @@
 ---
-source: crates/cli/tests/codegen.rs
+source: crates/codegen/tests/codegen.rs
 expression: outfiles
 ---
 "add_player_reducer.rs" = '''
@@ -1611,18 +1611,18 @@ impl TryFrom<__ws::DatabaseUpdate<__ws::BsatnFormat>> for DbUpdate {
         for table_update in raw.tables {
             match &table_update.table_name[..] {
 
-        "has_special_stuff" => db_update.has_special_stuff = has_special_stuff_table::parse_table_update(table_update)?,
-    "logged_out_player" => db_update.logged_out_player = logged_out_player_table::parse_table_update(table_update)?,
-    "person" => db_update.person = person_table::parse_table_update(table_update)?,
-    "pk_multi_identity" => db_update.pk_multi_identity = pk_multi_identity_table::parse_table_update(table_update)?,
-    "player" => db_update.player = player_table::parse_table_update(table_update)?,
-    "points" => db_update.points = points_table::parse_table_update(table_update)?,
-    "private_table" => db_update.private_table = private_table_table::parse_table_update(table_update)?,
-    "repeating_test_arg" => db_update.repeating_test_arg = repeating_test_arg_table::parse_table_update(table_update)?,
-    "test_a" => db_update.test_a = test_a_table::parse_table_update(table_update)?,
-    "test_d" => db_update.test_d = test_d_table::parse_table_update(table_update)?,
-    "test_e" => db_update.test_e = test_e_table::parse_table_update(table_update)?,
-    "test_f" => db_update.test_f = test_f_table::parse_table_update(table_update)?,
+        "has_special_stuff" => db_update.has_special_stuff.append(has_special_stuff_table::parse_table_update(table_update)?),
+    "logged_out_player" => db_update.logged_out_player.append(logged_out_player_table::parse_table_update(table_update)?),
+    "person" => db_update.person.append(person_table::parse_table_update(table_update)?),
+    "pk_multi_identity" => db_update.pk_multi_identity.append(pk_multi_identity_table::parse_table_update(table_update)?),
+    "player" => db_update.player.append(player_table::parse_table_update(table_update)?),
+    "points" => db_update.points.append(points_table::parse_table_update(table_update)?),
+    "private_table" => db_update.private_table.append(private_table_table::parse_table_update(table_update)?),
+    "repeating_test_arg" => db_update.repeating_test_arg.append(repeating_test_arg_table::parse_table_update(table_update)?),
+    "test_a" => db_update.test_a.append(test_a_table::parse_table_update(table_update)?),
+    "test_d" => db_update.test_d.append(test_d_table::parse_table_update(table_update)?),
+    "test_e" => db_update.test_e.append(test_e_table::parse_table_update(table_update)?),
+    "test_f" => db_update.test_f.append(test_f_table::parse_table_update(table_update)?),
 
                 unknown => {
                     return Err(__sdk::InternalError::unknown_name(

crates/sdk/src/spacetime_module.rs

@@ -212,6 +212,11 @@ impl<Row> Default for TableUpdate<Row> {
 }
 
 impl<Row> TableUpdate<Row> {
+    pub fn append(&mut self, mut other: TableUpdate<Row>) {
+        self.inserts.append(&mut other.inserts);
+        self.deletes.append(&mut other.deletes);
+    }
+
     pub(crate) fn is_empty(&self) -> bool {
         self.inserts.is_empty() && self.deletes.is_empty()
     }

crates/sdk/tests/test-client/src/main.rs

@@ -131,6 +131,9 @@ fn main() {
         "test-rls-subscription" => test_rls_subscription(),
         "pk-simple-enum" => exec_pk_simple_enum(),
         "indexed-simple-enum" => exec_indexed_simple_enum(),
+
+        "overlapping-subscriptions" => exec_overlapping_subscriptions(),
+
         _ => panic!("Unknown test: {}", test),
     }
 }
@@ -2523,3 +2526,83 @@ fn exec_indexed_simple_enum() {
     });
     test_counter.wait_for_all();
 }
+
+/// This tests for a bug we once had where the Rust client SDK would
+/// drop all but the last `TableUpdate` for a particular table within a `DatabaseUpdate`.
+///
+/// This manifested as a panic when applying incremental updates
+/// after an initial subscription to two or more overlapping queries each of which matched the same row.
+/// That row would incorrectly have multiplicity 1 in the client cache,
+/// since the SDK would ignore all but the first query's initial responses,
+/// but would then see incremental updates from all of the queries.
+///
+/// A simple reproducer is available at [https://github.com/lavirlifiliol/spacetime-repro].
+fn exec_overlapping_subscriptions() {
+    // First, a bit of setup: insert the row `{ n: 1, data: 0 }`,
+    // and wait for it to be present.
+    let setup_counter = TestCounter::new();
+
+    let call_insert_result = setup_counter.add_test("call_insert_reducer");
+    let mut row_inserted = Some(setup_counter.add_test("insert_reducer_done"));
+
+    let conn = connect_then(&setup_counter, move |ctx| {
+        ctx.reducers.on_insert_pk_u_8(move |ctx, _n, _data| {
+            (row_inserted.take().unwrap())(match &ctx.event.status {
+                Status::Committed => Ok(()),
+                s @ (Status::Failed(_) | Status::OutOfEnergy) => {
+                    Err(anyhow::anyhow!("insert_pk_u_8 during setup failed: {s:?}"))
+                }
+            });
+        });
+
+        call_insert_result(ctx.reducers.insert_pk_u_8(1, 0).map_err(|e| e.into()));
+    });
+
+    setup_counter.wait_for_all();
+
+    let test_counter = TestCounter::new();
+
+    let subscribe_result = test_counter.add_test("subscribe_with_row_present");
+
+    let call_update_result = test_counter.add_test("call_update_reducer");
+
+    let mut update_result = Some(test_counter.add_test("update_row"));
+
+    // Now, subscribe to two queries which each match that row.
+    subscribe_these_then(
+        &conn,
+        &["select * from pk_u8 where n < 100", "select * from pk_u8 where n > 0"],
+        move |ctx| {
+            // It's not exposed to users of the SDK, so we won't assert on it,
+            // but we expect the row to have multiplicity 2.
+            subscribe_result(if ctx.db.pk_u_8().count() == 1 {
+                Ok(())
+            } else {
+                Err(anyhow::anyhow!(
+                    "Expected one row for PkU8 but found {}",
+                    ctx.db.pk_u_8().count()
+                ))
+            });
+        },
+    );
+
+    // Once the row is in the cache, update it by replacing it with `{ n: 1, data: 1 }`.
+    conn.db.pk_u_8().on_update(move |ctx, old, new| {
+        // It's not exposed, so no assert,
+        // but we expect to have received two deletes for the old row,
+        // and two inserts for the new row.
+        // The SDK will combine all of these into a single update event.
+        (update_result.take().unwrap())((|| {
+            anyhow::ensure!(old.n == new.n);
+            anyhow::ensure!(old.n == 1);
+            anyhow::ensure!(old.data == 0);
+            anyhow::ensure!(new.data == 1);
+            anyhow::ensure!(ctx.db.pk_u_8().count() == 1);
+            Ok(())
+        })())
+    });
+
+    call_update_result(conn.reducers.update_pk_u_8(1, 1).map_err(|e| e.into()));
+
+    test_counter.wait_for_all();
+}