update: disconnect clients

Author: Shubham8287

crates/core/src/db/update.rs

@@ -1,6 +1,7 @@
 use super::relational_db::RelationalDB;
 use crate::database_logger::SystemLogger;
 use crate::sql::parser::RowLevelExpr;
+use crate::subscription::module_subscription_actor::ModuleSubscriptions;
 use spacetimedb_data_structures::map::HashMap;
 use spacetimedb_datastore::locking_tx_datastore::MutTxId;
 use spacetimedb_lib::db::auth::StTableType;
@@ -35,6 +36,7 @@ impl UpdateLogger for SystemLogger {
 // drop_* become transactional.
 pub fn update_database(
     stdb: &RelationalDB,
+    subscriptions: &ModuleSubscriptions,
     tx: &mut MutTxId,
     auth_ctx: AuthCtx,
     plan: MigratePlan,
@@ -57,7 +59,9 @@ pub fn update_database(
 
     match plan {
         MigratePlan::Manual(plan) => manual_migrate_database(stdb, tx, plan, logger, existing_tables),
-        MigratePlan::Auto(plan) => auto_migrate_database(stdb, tx, auth_ctx, plan, logger, existing_tables),
+        MigratePlan::Auto(plan) => {
+            auto_migrate_database(stdb, subscriptions, tx, auth_ctx, plan, logger, existing_tables)
+        }
     }
 }
 
@@ -83,6 +87,7 @@ macro_rules! log {
 /// Automatically migrate a database.
 fn auto_migrate_database(
     stdb: &RelationalDB,
+    subscriptions: &ModuleSubscriptions,
     tx: &mut MutTxId,
     auth_ctx: AuthCtx,
     plan: AutoMigratePlan,
@@ -264,7 +269,12 @@ fn auto_migrate_database(
                     .collect();
                 stdb.add_columns_to_table(tx, table_id, column_schemas, default_values)?;
             }
-            _ => anyhow::bail!("migration step not implemented: {step:?}"),
+            spacetimedb_schema::auto_migrate::AutoMigrateStep::DisconnectAllUsers => {
+                // Disconnect all clients from subscriptions.
+                // Any dangling clients will be handled during the launch of module hosts,
+                // which invokes `ModuleHost::call_identity_disconnected`.
+                subscriptions.remove_all_subscribers();
+            }
         }
     }
 
@@ -349,7 +359,14 @@ mod test {
         // Try to update the db.
         let mut tx = begin_mut_tx(&stdb);
         let plan = ponder_migrate(&old, &new)?;
-        update_database(&stdb, &mut tx, auth_ctx, plan, &TestLogger)?;
+        update_database(
+            &stdb,
+            &ModuleSubscriptions::for_test_new_runtime(Arc::new(stdb.db.clone())).0,
+            &mut tx,
+            auth_ctx,
+            plan,
+            &TestLogger,
+        )?;
 
         // Expect the schema change.
         let idx_b_id = stdb

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

@@ -306,7 +306,14 @@ impl InstanceCommon {
         system_logger.info(&format!("Updated program to {program_hash}"));
 
         let auth_ctx = AuthCtx::for_current(replica_ctx.database.owner_identity);
-        let res = crate::db::update::update_database(stdb, &mut tx, auth_ctx, plan, system_logger);
+        let res = crate::db::update::update_database(
+            stdb,
+            &replica_ctx.subscriptions,
+            &mut tx,
+            auth_ctx,
+            plan,
+            system_logger,
+        );
 
         match res {
             Err(e) => {

crates/core/src/subscription/module_subscription_actor.rs

@@ -853,6 +853,11 @@ impl ModuleSubscriptions {
         subscriptions.remove_all_subscriptions(&(client_id.identity, client_id.connection_id));
     }
 
+    pub fn remove_all_subscribers(&self) {
+        let mut subscriptions = self.subscriptions.write();
+        subscriptions.remove_all_clients();
+    }
+
     /// Commit a transaction and broadcast its ModuleEvent to all interested subscribers.
     ///
     /// The returned [`ExecutionMetrics`] are reported in this method via `report_tx_metrics`.

crates/core/src/subscription/module_subscription_manager.rs

@@ -779,6 +779,12 @@ impl SubscriptionManager {
         }
     }
 
+    pub fn remove_all_clients(&mut self) {
+        for id in self.clients.keys().copied().collect::<Vec<_>>() {
+            self.remove_all_subscriptions(&id);
+        }
+    }
+
     /// Remove a single subscription for a client.
     /// This will return an error if the client does not have a subscription with the given query id.
     pub fn remove_subscription(&mut self, client_id: ClientId, query_id: ClientQueryId) -> Result<Vec<Query>, DBError> {