Fix v2 client disconnects dropping subscriptions for other v2 clients (#4648)

# Description of Changes

This bug was introduced when we added the v2 subscription protocol.

Before this change, when a v2 client disconnected or unsubscribed, other
clients could stop receiving updates for sender-scoped views they were
subscribed to.

This happened because `SubscriptionManager::remove_all_subscriptions()`
was cleaning up v2 subscription state correctly, but we would also run
the legacy v1 cleanup loop which mixed v1 and v2 queries. This meant the
same query hash could be processed again in the v1 cleanup loop after
the v2 removal had already happened. In the sender-view case, that
second pass could remove metadata (`indexes` / `search_args`) for active
subscribers.

Specifically, we could have the following sequence:
1. Client B disconnects.
2. The v2 loop removes B’s v2 subscription for `query_hash`.
3. `query_hash` now has no subscribers, so
`remove_query_from_tables(query_hash)` runs once.
4. That compatibility loop then visits `query_hash` again via
`subscription_ref_count.keys()`.
5. `query_hash` still has no subscribers, so
`remove_query_from_tables(query_hash)` runs a second time.

and when `remove_query_from_tables(query_hash)` runs a 2nd time, it
calls `index_ids.delete_index_ids_for_query()` which is refcounted, not
idempotent. So calling it twice for the same query decrements the shared
index count twice, and the 2nd decrement drops an index that client A’s
query still needs.

The fix was to:
1. Collect the disconnected client’s v1 query hashes from
`client_info.v1_subscriptions`
2. Run the v2 cleanup loop over `client_info.v2_subscriptions` (as we
did before this change)
3. Run the old compatibility cleanup loop only over the deduplicated v1
hashes

# API and ABI breaking changes

None

# Expected complexity level and risk

2

Small fix to a rather intricate bug. Most of the code in this patch is
testing code.

# Testing

This bug was not reproducible with the current smoketest harness,
because `spacetime subscribe` still uses the `v1` subscription protocol.
I added a Rust SDK test, which did reproduce the bug before the fix,
because the Rust SDK uses the `v2` subscription protocol.

Author: joshua-spacetime

crates/core/src/subscription/module_subscription_actor.rs

@@ -1880,7 +1880,8 @@ mod tests {
         Protocol, WsVersion,
     };
     use crate::db::relational_db::tests_utils::{
-        begin_mut_tx, begin_tx, insert, with_auto_commit, with_read_only, TestDB,
+        begin_mut_tx, begin_tx, create_view_for_test, insert, insert_into_view, with_auto_commit, with_read_only,
+        TestDB,
     };
     use crate::db::relational_db::{Persistence, RelationalDB, Txdata};
     use crate::error::DBError;
@@ -1919,6 +1920,8 @@ mod tests {
     use tokio::sync::mpsc::{self};
     use tokio::sync::watch;
 
+    const TEST_MESSAGE_TIMEOUT: Duration = Duration::from_millis(20);
+
     fn add_subscriber(db: Arc<RelationalDB>, sql: &str, assert: Option<AssertTxFn>) -> Result<(), DBError> {
         // Create and enter a Tokio runtime to run the `ModuleSubscriptions`' background workers in parallel.
         let runtime = tokio::runtime::Runtime::new().unwrap();
@@ -2180,6 +2183,24 @@ mod tests {
         )
     }
 
+    /// Instantiate a v2 client connection with the default test settings.
+    fn v2_client_connection(
+        client_id: ClientActorId,
+        db: &Arc<RelationalDB>,
+    ) -> (Arc<ClientConnectionSender>, ClientConnectionReceiver) {
+        client_connection_with_config(
+            client_id,
+            db,
+            ClientConfig {
+                protocol: Protocol::Binary,
+                version: WsVersion::V2,
+                compression: ws_v1::Compression::None,
+                tx_update_full: true,
+                confirmed_reads: false,
+            },
+        )
+    }
+
     /// Insert rules into the RLS system table
     fn insert_rls_rules(
         db: &RelationalDB,
@@ -2263,17 +2284,7 @@ mod tests {
         let db = relational_db()?;
 
         let client_id = client_id_from_u8(1);
-        let (sender, mut rx) = client_connection_with_config(
-            client_id,
-            &db,
-            ClientConfig {
-                protocol: Protocol::Binary,
-                version: WsVersion::V2,
-                compression: ws_v1::Compression::None,
-                tx_update_full: true,
-                confirmed_reads: false,
-            },
-        );
+        let (sender, mut rx) = v2_client_connection(client_id, &db);
 
         let auth = AuthCtx::new(db.owner_identity(), client_id.identity);
         let subs = ModuleSubscriptions::for_test_enclosing_runtime(db.clone());
@@ -2333,17 +2344,7 @@ mod tests {
         let db = relational_db()?;
 
         let client_id = client_id_from_u8(1);
-        let (sender, mut rx) = client_connection_with_config(
-            client_id,
-            &db,
-            ClientConfig {
-                protocol: Protocol::Binary,
-                version: WsVersion::V2,
-                compression: ws_v1::Compression::None,
-                tx_update_full: true,
-                confirmed_reads: false,
-            },
-        );
+        let (sender, mut rx) = v2_client_connection(client_id, &db);
 
         let auth = AuthCtx::new(db.owner_identity(), client_id.identity);
         let subs = ModuleSubscriptions::for_test_enclosing_runtime(db.clone());
@@ -2405,17 +2406,7 @@ mod tests {
         let db = relational_db()?;
 
         let client_id = client_id_from_u8(1);
-        let (sender, mut rx) = client_connection_with_config(
-            client_id,
-            &db,
-            ClientConfig {
-                protocol: Protocol::Binary,
-                version: WsVersion::V2,
-                compression: ws_v1::Compression::None,
-                tx_update_full: true,
-                confirmed_reads: false,
-            },
-        );
+        let (sender, mut rx) = v2_client_connection(client_id, &db);
 
         let auth = AuthCtx::new(db.owner_identity(), client_id.identity);
         let subs = ModuleSubscriptions::for_test_enclosing_runtime(db.clone());
@@ -2460,8 +2451,84 @@ mod tests {
 
         let _ = commit_tx(&db, &subs, [], [(table_id, product![2_u8])])?;
 
-        let recv = tokio::time::timeout(Duration::from_millis(20), rx.recv()).await;
-        assert!(recv.is_err(), "expected no updates after unsubscribe");
+        assert_no_outbound_message(rx.recv()).await;
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn unsubscribe_v2_other_clients_receive_sender_view_updates() -> anyhow::Result<()> {
+        let db = relational_db()?;
+
+        let id_for_a = identity_from_u8(1);
+        let client_id_for_a = client_id_from_u8(1);
+        let client_id_for_b = client_id_from_u8(2);
+
+        let (tx_for_a, mut rx_for_a) = v2_client_connection(client_id_for_a, &db);
+        let (tx_for_b, mut rx_for_b) = v2_client_connection(client_id_for_b, &db);
+
+        let auth_for_a = AuthCtx::new(db.owner_identity(), client_id_for_a.identity);
+        let auth_for_b = AuthCtx::new(db.owner_identity(), client_id_for_b.identity);
+        let subs = ModuleSubscriptions::for_test_enclosing_runtime(db.clone());
+
+        let (_, view_table_id) = create_view_for_test(&db, "my_view", &[("counter", AlgebraicType::U8)], false)?;
+
+        // Seed a sender-scoped row that only client A should observe through the view.
+        with_auto_commit(&db, |tx| -> anyhow::Result<_> {
+            insert_into_view(&db, tx, view_table_id, Some(id_for_a), product![7_u8])?;
+            Ok(())
+        })?;
+
+        subs.add_v2_subscription_inner::<crate::host::wasmtime::WasmtimeInstance>(
+            None,
+            tx_for_a.clone(),
+            auth_for_a,
+            ws_v2::Subscribe {
+                request_id: 1,
+                query_set_id: ws_v2::QuerySetId::new(1),
+                query_strings: ["select * from my_view".into()].into(),
+            },
+            Instant::now(),
+            None,
+        )?;
+        subs.add_v2_subscription_inner::<crate::host::wasmtime::WasmtimeInstance>(
+            None,
+            tx_for_b.clone(),
+            auth_for_b,
+            ws_v2::Subscribe {
+                request_id: 2,
+                query_set_id: ws_v2::QuerySetId::new(2),
+                query_strings: ["select * from my_view".into()].into(),
+            },
+            Instant::now(),
+            None,
+        )?;
+
+        assert!(matches!(
+            rx_for_a.recv().await,
+            Some(OutboundMessage::V2(ws_v2::ServerMessage::SubscribeApplied(_)))
+        ));
+        assert!(matches!(
+            rx_for_b.recv().await,
+            Some(OutboundMessage::V2(ws_v2::ServerMessage::SubscribeApplied(_)))
+        ));
+
+        // Dropping client B must not break client A's sender-view bookkeeping.
+        subs.remove_subscriber(client_id_for_b);
+
+        // Delete the backing row and verify the surviving subscriber still receives the view delta.
+        let _ = commit_tx(&db, &subs, [(view_table_id, product![id_for_a, 7_u8])], [])?;
+
+        let schema = ProductType::from([AlgebraicType::U8]);
+        assert_v2_tx_update_for_table(
+            rx_for_a.recv(),
+            ws_v2::QuerySetId::new(1),
+            "my_view",
+            &schema,
+            [],
+            [product![7_u8]],
+        )
+        .await;
 
         Ok(())
     }
@@ -2477,6 +2544,59 @@ mod tests {
         Ok(())
     }
 
+    fn update_row_counts<I, D, BI, BD>(
+        rows_received: &mut HashMap<ProductValue, i32>,
+        schema: &ProductType,
+        inserts: I,
+        deletes: D,
+    ) where
+        I: IntoIterator<Item = BI>,
+        D: IntoIterator<Item = BD>,
+        BI: AsRef<[u8]>,
+        BD: AsRef<[u8]>,
+    {
+        for row in inserts.into_iter().map(|bytes| {
+            let mut bytes = bytes.as_ref();
+            ProductValue::decode(schema, &mut bytes).unwrap()
+        }) {
+            *rows_received.entry(row).or_insert(0) += 1;
+        }
+
+        for row in deletes.into_iter().map(|bytes| {
+            let mut bytes = bytes.as_ref();
+            ProductValue::decode(schema, &mut bytes).unwrap()
+        }) {
+            *rows_received.entry(row).or_insert(0) -= 1;
+        }
+    }
+
+    fn assert_received_rows(
+        rows_received: HashMap<ProductValue, i32>,
+        inserts: impl IntoIterator<Item = ProductValue>,
+        deletes: impl IntoIterator<Item = ProductValue>,
+    ) {
+        assert_eq!(
+            rows_received
+                .iter()
+                .filter(|(_, n)| n > &&0)
+                .map(|(row, _)| row)
+                .cloned()
+                .sorted()
+                .collect::<Vec<_>>(),
+            inserts.into_iter().sorted().collect::<Vec<_>>()
+        );
+        assert_eq!(
+            rows_received
+                .iter()
+                .filter(|(_, n)| n < &&0)
+                .map(|(row, _)| row)
+                .cloned()
+                .sorted()
+                .collect::<Vec<_>>(),
+            deletes.into_iter().sorted().collect::<Vec<_>>()
+        );
+    }
+
     /// Pull a message from receiver and assert that it is a `TxUpdate` with the expected rows
     async fn assert_tx_update_for_table(
         rx: impl Future<Output = Option<OutboundMessage>>,
@@ -2485,7 +2605,7 @@ mod tests {
         inserts: impl IntoIterator<Item = ProductValue>,
         deletes: impl IntoIterator<Item = ProductValue>,
     ) {
-        match rx.await {
+        match recv_outbound_message(rx, "TxUpdate").await {
             Some(OutboundMessage::V1(SerializableMessage::TxUpdate(TransactionUpdateMessage {
                 database_update:
                     SubscriptionUpdateMessage {
@@ -2512,49 +2632,70 @@ mod tests {
                         panic!("expected an uncompressed table update")
                     };
 
-                    for row in table_update
-                        .inserts
-                        .into_iter()
-                        .map(|bytes| ProductValue::decode(schema, &mut &*bytes).unwrap())
-                    {
-                        *rows_received.entry(row).or_insert(0) += 1;
-                    }
-
-                    for row in table_update
-                        .deletes
-                        .into_iter()
-                        .map(|bytes| ProductValue::decode(schema, &mut &*bytes).unwrap())
-                    {
-                        *rows_received.entry(row).or_insert(0) -= 1;
-                    }
+                    update_row_counts(&mut rows_received, schema, &table_update.inserts, &table_update.deletes);
                 }
 
-                assert_eq!(
-                    rows_received
-                        .iter()
-                        .filter(|(_, n)| n > &&0)
-                        .map(|(row, _)| row)
-                        .cloned()
-                        .sorted()
-                        .collect::<Vec<_>>(),
-                    inserts.into_iter().sorted().collect::<Vec<_>>()
-                );
-                assert_eq!(
-                    rows_received
-                        .iter()
-                        .filter(|(_, n)| n < &&0)
-                        .map(|(row, _)| row)
-                        .cloned()
-                        .sorted()
-                        .collect::<Vec<_>>(),
-                    deletes.into_iter().sorted().collect::<Vec<_>>()
-                );
+                assert_received_rows(rows_received, inserts, deletes);
             }
             Some(msg) => panic!("expected a TxUpdate, but got {msg:#?}"),
             None => panic!("The receiver closed due to an error"),
         }
     }
 
+    /// Pull a message from receiver and assert that it is a v2 `TransactionUpdate`
+    /// with the expected rows for a single table in a single query set.
+    async fn assert_v2_tx_update_for_table(
+        rx: impl Future<Output = Option<OutboundMessage>>,
+        query_set_id: ws_v2::QuerySetId,
+        table_name: &str,
+        schema: &ProductType,
+        inserts: impl IntoIterator<Item = ProductValue>,
+        deletes: impl IntoIterator<Item = ProductValue>,
+    ) {
+        match recv_outbound_message(rx, "v2 TransactionUpdate").await {
+            Some(OutboundMessage::V2(ws_v2::ServerMessage::TransactionUpdate(update))) => {
+                assert_eq!(update.query_sets.len(), 1);
+                let query_set = &update.query_sets[0];
+                assert_eq!(query_set.query_set_id, query_set_id);
+                assert_eq!(query_set.tables.len(), 1);
+
+                let table_update = &query_set.tables[0];
+                assert_eq!(table_update.table_name.as_ref(), table_name);
+
+                let mut rows_received: HashMap<ProductValue, i32> = HashMap::new();
+
+                for rows in table_update.rows.iter() {
+                    let ws_v2::TableUpdateRows::PersistentTable(rows) = rows else {
+                        panic!("expected a persistent-table update")
+                    };
+
+                    update_row_counts(&mut rows_received, schema, &rows.inserts, &rows.deletes);
+                }
+
+                assert_received_rows(rows_received, inserts, deletes);
+            }
+            Some(msg) => panic!("expected a v2 TransactionUpdate, but got {msg:#?}"),
+            None => panic!("The receiver closed due to an error"),
+        }
+    }
+
+    async fn recv_outbound_message(
+        rx: impl Future<Output = Option<OutboundMessage>>,
+        expected: &str,
+    ) -> Option<OutboundMessage> {
+        tokio::time::timeout(TEST_MESSAGE_TIMEOUT, rx)
+            .await
+            .unwrap_or_else(|_| panic!("timed out waiting for {expected}"))
+    }
+
+    async fn assert_no_outbound_message(rx: impl Future<Output = Option<OutboundMessage>>) {
+        match tokio::time::timeout(TEST_MESSAGE_TIMEOUT, rx).await {
+            Err(_) => {}
+            Ok(Some(msg)) => panic!("expected no message, got {msg:#?}"),
+            Ok(None) => panic!("the receiver closed due to an error"),
+        }
+    }
+
     /// Assert that the future `f` completes only after `durability` is marked
     /// durable.
     ///