Fix disconnects breaking view updates for other connections (#4607)

# Description of Changes

Fixes a bug in client disconnect logic that would mark a client's views
as dropped(unsubscribed). However it was marking the identity's views as
dropped, not the connection. So if an identity had multiple connections
open, each subscribing to different views, and one of them disconnected,
the subscriptions for the other connections would break. The observed
behavior would be that they would stop receiving subscription updates.
This could potentially lead to their client cache getting into a
corrupted state.

Now, instead of dropping all of the views for a particular identity on
disconnect, we drop only the views for that particular connection. And
when I say drop, what I really mean is decrement. A view is not dropped
completely unless it no longer had any subscribers.

# API and ABI breaking changes

None

# Expected complexity level and risk

2

# Testing

Regression smoketest was added

Author: joshua-spacetime

crates/client-api/src/routes/database.rs

@@ -184,8 +184,7 @@ pub async fn call<S: ControlStateDelegate + NodeDelegate>(
     };
 
     module
-        // We don't clear views or procedures after reducer calls
-        .call_identity_disconnected(caller_identity, connection_id, false)
+        .call_identity_disconnected(caller_identity, connection_id)
         .await
         .map_err(client_disconnected_error_to_response)?;
 

crates/core/src/host/host_controller.rs

@@ -1012,7 +1012,7 @@ impl Host {
         // No need to clear view tables here since we do it in `clear_all_clients`.
         for (identity, connection_id) in connected_clients {
             module_host
-                .call_identity_disconnected(identity, connection_id, false)
+                .call_identity_disconnected(identity, connection_id)
                 .await
                 .with_context(|| {
                     format!(

crates/core/src/host/module_host.rs

@@ -1251,7 +1251,6 @@ impl ModuleHost {
             client_id.identity,
             client_id.connection_id,
             info,
-            true,
             call_reducer,
             trapped_slot,
         )
@@ -1293,7 +1292,6 @@ impl ModuleHost {
         caller_identity: Identity,
         caller_connection_id: ConnectionId,
         info: &ModuleInfo,
-        drop_view_subscribers: bool,
         call_reducer: impl FnOnce(Option<MutTxId>, CallReducerParams) -> (ReducerCallResult, bool),
         trapped_slot: &mut bool,
     ) -> Result<(), ReducerCallError> {
@@ -1309,20 +1307,10 @@ impl ModuleHost {
 
         let workload = || Workload::reducer_no_args(reducer_name.clone(), caller_identity, caller_connection_id);
 
-        // Decrement the number of subscribers for each view this caller is subscribed to
-        let dec_view_subscribers = |tx: &mut MutTxId| {
-            if drop_view_subscribers && let Err(err) = tx.unsubscribe_views(caller_identity) {
-                log::error!("`call_identity_disconnected`: failed to delete client view data: {err}");
-            }
-        };
-
         // A fallback transaction that deletes the client from `st_client`.
         let database_identity = stdb.database_identity();
         let fallback = || {
             stdb.with_auto_commit(workload(), |mut_tx| {
-
-                dec_view_subscribers(mut_tx);
-
                 if !is_client_exist(mut_tx) {
                     // The client is already gone. Nothing to do.
                     log::debug!(
@@ -1348,9 +1336,7 @@ impl ModuleHost {
         };
 
         if let Some((reducer_id, reducer_def)) = reducer_lookup {
-            let mut mut_tx = stdb.begin_mut_tx(IsolationLevel::Serializable, workload());
-
-            dec_view_subscribers(&mut mut_tx);
+            let mut_tx = stdb.begin_mut_tx(IsolationLevel::Serializable, workload());
 
             if !is_client_exist(&mut_tx) {
                 // The client is already gone. Nothing to do.
@@ -1430,13 +1416,12 @@ impl ModuleHost {
         &self,
         caller_identity: Identity,
         caller_connection_id: ConnectionId,
-        drop_view_subscribers: bool,
     ) -> Result<(), ReducerCallError> {
         self.call(
             "call_identity_disconnected",
-            (caller_identity, caller_connection_id, drop_view_subscribers),
-            async |(a, b, c), inst| inst.call_identity_disconnected(a, b, c),
-            async |(a, b, c), inst| inst.call_identity_disconnected(a, b, c).await,
+            (caller_identity, caller_connection_id),
+            async |(a, b), inst| inst.call_identity_disconnected(a, b),
+            async |(a, b), inst| inst.call_identity_disconnected(a, b).await,
         )
         .await?
     }

crates/core/src/host/v8/mod.rs

@@ -410,11 +410,10 @@ impl JsInstance {
         &mut self,
         caller_identity: Identity,
         caller_connection_id: ConnectionId,
-        drop_view_subscribers: bool,
     ) -> Result<(), ReducerCallError> {
         self.send_recv(
             JsWorkerReply::into_call_identity_disconnected,
-            JsWorkerRequest::CallIdentityDisconnected(caller_identity, caller_connection_id, drop_view_subscribers),
+            JsWorkerRequest::CallIdentityDisconnected(caller_identity, caller_connection_id),
         )
         .await
     }
@@ -505,7 +504,7 @@ enum JsWorkerRequest {
     /// See [`JsInstance::call_identity_connected`].
     CallIdentityConnected(ConnectionAuthCtx, ConnectionId),
     /// See [`JsInstance::call_identity_disconnected`].
-    CallIdentityDisconnected(Identity, ConnectionId, bool),
+    CallIdentityDisconnected(Identity, ConnectionId),
     /// See [`JsInstance::disconnect_client`].
     DisconnectClient(ClientActorId),
     /// See [`JsInstance::init_database`].
@@ -718,17 +717,12 @@ async fn spawn_instance_worker(
                         call_identity_connected(caller_auth, caller_connection_id, info, call_reducer, &mut trapped);
                     reply("call_identity_connected", CallIdentityConnected(res), trapped);
                 }
-                JsWorkerRequest::CallIdentityDisconnected(
-                    caller_identity,
-                    caller_connection_id,
-                    drop_view_subcribers,
-                ) => {
+                JsWorkerRequest::CallIdentityDisconnected(caller_identity, caller_connection_id) => {
                     let mut trapped = false;
                     let res = ModuleHost::call_identity_disconnected_inner(
                         caller_identity,
                         caller_connection_id,
                         info,
-                        drop_view_subcribers,
                         call_reducer,
                         &mut trapped,
                     );

crates/core/src/host/wasm_common/module_host_actor.rs

@@ -488,7 +488,6 @@ impl<T: WasmInstance> WasmModuleInstance<T> {
         &mut self,
         caller_identity: Identity,
         caller_connection_id: ConnectionId,
-        drop_view_subscribers: bool,
     ) -> Result<(), ReducerCallError> {
         let module = &self.common.info.clone();
         let call_reducer = |tx, params| self.call_reducer_with_tx(tx, params);
@@ -497,7 +496,6 @@ impl<T: WasmInstance> WasmModuleInstance<T> {
             caller_identity,
             caller_connection_id,
             module,
-            drop_view_subscribers,
             call_reducer,
             &mut trapped,
         );

crates/core/src/subscription/module_subscription_actor.rs

@@ -1560,8 +1560,24 @@ impl ModuleSubscriptions {
     }
 
     pub fn remove_subscriber(&self, client_id: ClientActorId) {
-        let mut subscriptions = self.subscriptions.write();
-        subscriptions.remove_all_subscriptions(&(client_id.identity, client_id.connection_id));
+        let removed_queries = {
+            let mut subscriptions = self.subscriptions.write();
+            subscriptions.remove_all_subscriptions(&(client_id.identity, client_id.connection_id))
+        };
+
+        if removed_queries.is_empty() {
+            return;
+        }
+
+        // TODO(perf): Removing a subscriber is currently O(subscribed_queries).
+        // Instead we should maintain an index to make this O(subscribed_views).
+        if let Err(err) = self.unsubscribe_views(&removed_queries, client_id.identity) {
+            log::error!(
+                "failed to unsubscribe views for disconnected client ({}, {}): {err}",
+                client_id.identity,
+                client_id.connection_id
+            );
+        }
     }
 
     /// Rolls back `tx` and returns the offset as it was before `tx`.

crates/core/src/subscription/module_subscription_manager.rs

@@ -1203,10 +1203,20 @@ impl SubscriptionManager {
     /// If a query no longer has any subscribers,
     /// it is removed from the index along with its table ids.
     #[tracing::instrument(level = "trace", skip_all)]
-    pub fn remove_all_subscriptions(&mut self, client: &ClientId) {
+    pub fn remove_all_subscriptions(&mut self, client: &ClientId) -> Vec<Query> {
+        let mut removed_query_hashes = HashSet::new();
+        if let Some(client_info) = self.clients.get(client) {
+            removed_query_hashes.extend(client_info.legacy_subscriptions.iter().copied());
+            removed_query_hashes.extend(client_info.subscription_ref_count.keys().copied());
+        }
+        let removed_queries = removed_query_hashes
+            .into_iter()
+            .filter_map(|hash| self.queries.get(&hash).map(|q| q.query.clone()))
+            .collect::<Vec<_>>();
+
         self.remove_legacy_subscriptions(client);
         let Some(client_info) = self.remove_client_and_inform_send_worker(*client) else {
-            return;
+            return removed_queries;
         };
 
         debug_assert!(client_info.legacy_subscriptions.is_empty());
@@ -1252,6 +1262,8 @@ impl SubscriptionManager {
         for query_hash in queries_to_remove {
             self.queries.remove(&query_hash);
         }
+
+        removed_queries
     }
 
     /// Find the queries that need to be evaluated for this table update.

crates/smoketests/modules/views-sql/src/lib.rs

@@ -24,6 +24,16 @@ pub fn add_player_level(ctx: &ReducerContext, id: u64, level: u64) {
     ctx.db.player_level().insert(PlayerState { id, level });
 }
 
+#[spacetimedb::reducer]
+pub fn set_player_state(ctx: &ReducerContext, id: u64, level: u64) {
+    if let Some(mut row) = ctx.db.player_state().id().find(id) {
+        row.level = level;
+        ctx.db.player_state().id().update(row);
+    } else {
+        ctx.db.player_state().insert(PlayerState { id, level });
+    }
+}
+
 #[spacetimedb::view(accessor = my_player_and_level, public)]
 pub fn my_player_and_level(ctx: &AnonymousViewContext) -> Option<PlayerState> {
     ctx.db.player_level().id().find(0)

crates/smoketests/tests/smoketests/views.rs

@@ -553,6 +553,75 @@ fn test_typescript_procedure_triggers_subscription_updates() {
     );
 }
 
+#[test]
+fn test_disconnect_does_not_break_sender_view() {
+    let test = Smoketest::builder().precompiled_module("views-sql").build();
+
+    test.call("set_player_state", &["42", "1"]).unwrap();
+
+    // Two connections subscribe to the same view.
+    let sub_keep = test.subscribe_background(&["SELECT * FROM player"], 2).unwrap();
+    let sub_drop = test.subscribe_background(&["SELECT * FROM player"], 1).unwrap();
+
+    // Both connections should receive the first update.
+    // After one connection disconnects, the other should still receive updates.
+    test.call("set_player_state", &["42", "2"]).unwrap();
+    let _ = sub_drop.collect().unwrap();
+    test.call("set_player_state", &["42", "3"]).unwrap();
+
+    let events = sub_keep.collect().unwrap();
+
+    assert_eq!(events.len(), 2, "Expected two updates for player, got: {events:?}");
+    let inserts = events[1]["player"]["inserts"]
+        .as_array()
+        .expect("Expected inserts array on player update");
+    assert!(
+        inserts
+            .iter()
+            .any(|row| row["id"] == json!(42) && row["level"] == json!(3)),
+        "Expected player id=42 level=3 insert after disconnect, got: {events:?}"
+    );
+}
+
+#[test]
+fn test_disconnect_does_not_break_anonymous_view() {
+    let test = Smoketest::builder().precompiled_module("views-sql").build();
+
+    // Seed a row in the anonymous-view source table.
+    test.call("add_player_level", &["0", "2"]).unwrap();
+
+    // Two connections subscribe to the same anonymous view.
+    let sub_keep = test
+        .subscribe_background(&["SELECT * FROM player_and_level"], 2)
+        .unwrap();
+    let sub_drop = test
+        .subscribe_background(&["SELECT * FROM player_and_level"], 1)
+        .unwrap();
+
+    // Both connections should receive the first update.
+    // After one connection disconnects, the other should still receive updates.
+    test.call("add_player_level", &["1", "2"]).unwrap();
+    let _ = sub_drop.collect().unwrap();
+    test.call("add_player_level", &["2", "2"]).unwrap();
+
+    let events = sub_keep.collect().unwrap();
+
+    assert_eq!(
+        events.len(),
+        2,
+        "Expected two updates for player_and_level, got: {events:?}"
+    );
+    let inserts = events[1]["player_and_level"]["inserts"]
+        .as_array()
+        .expect("Expected inserts array on player_and_level update");
+    assert!(
+        inserts
+            .iter()
+            .any(|row| row["id"] == json!(2) && row["level"] == json!(2)),
+        "Expected player id=2 level=2 insert after disconnect, got: {events:?}"
+    );
+}
+
 #[test]
 fn test_typescript_query_builder_view_query() {
     require_pnpm!();