disconnect client by dropping module

Author: Shubham8287

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

@@ -664,7 +664,9 @@ pub async fn publish<S: NodeDelegate + ControlStateDelegate>(
                 )
                     .into());
             }
-            UpdateDatabaseResult::NoUpdateNeeded | UpdateDatabaseResult::UpdatePerformed => {}
+            UpdateDatabaseResult::NoUpdateNeeded
+            | UpdateDatabaseResult::UpdatePerformed
+            | UpdateDatabaseResult::UpdatePerformedWithClientDisconnect => {}
         }
     }
 

crates/core/src/db/update.rs

@@ -1,7 +1,6 @@
 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;
@@ -24,6 +23,13 @@ impl UpdateLogger for SystemLogger {
     }
 }
 
+/// The result of a database update.
+/// Indicates whether clients should be disconnected when the update is complete.
+pub enum UpdateResult {
+    Success,
+    RequiresClientDisconnect,
+}
+
 /// Update the database according to the migration plan.
 ///
 /// The update is performed within the transactional context `tx`.
@@ -36,12 +42,11 @@ impl UpdateLogger for SystemLogger {
 // drop_* become transactional.
 pub fn update_database(
     stdb: &RelationalDB,
-    subscriptions: &ModuleSubscriptions,
     tx: &mut MutTxId,
     auth_ctx: AuthCtx,
     plan: MigratePlan,
     logger: &dyn UpdateLogger,
-) -> anyhow::Result<()> {
+) -> anyhow::Result<UpdateResult> {
     let existing_tables = stdb.get_all_tables_mut(tx)?;
 
     // TODO: consider using `ErrorStream` here.
@@ -59,9 +64,7 @@ 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, subscriptions, tx, auth_ctx, plan, logger, existing_tables)
-        }
+        MigratePlan::Auto(plan) => auto_migrate_database(stdb, tx, auth_ctx, plan, logger, existing_tables),
     }
 }
 
@@ -72,7 +75,7 @@ fn manual_migrate_database(
     _plan: ManualMigratePlan,
     _logger: &dyn UpdateLogger,
     _existing_tables: Vec<Arc<TableSchema>>,
-) -> anyhow::Result<()> {
+) -> anyhow::Result<UpdateResult> {
     unimplemented!("Manual database migrations are not yet implemented")
 }
 
@@ -87,13 +90,12 @@ macro_rules! log {
 /// Automatically migrate a database.
 fn auto_migrate_database(
     stdb: &RelationalDB,
-    subscriptions: &ModuleSubscriptions,
     tx: &mut MutTxId,
     auth_ctx: AuthCtx,
     plan: AutoMigratePlan,
     logger: &dyn UpdateLogger,
     existing_tables: Vec<Arc<TableSchema>>,
-) -> anyhow::Result<()> {
+) -> anyhow::Result<UpdateResult> {
     // We have already checked in `migrate_database` that `existing_tables` are compatible with the `old` definition in `plan`.
     // So we can look up tables in there using unwrap.
 
@@ -132,6 +134,7 @@ fn auto_migrate_database(
     }
 
     log::info!("Running database update steps: {}", stdb.database_identity());
+    let mut res = UpdateResult::Success;
 
     for step in plan.steps {
         match step {
@@ -270,16 +273,16 @@ fn auto_migrate_database(
                 stdb.add_columns_to_table(tx, table_id, column_schemas, default_values)?;
             }
             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();
+                log!(logger, "Disconnecting all users");
+                // It does disconnect clients right away,
+                // but send response indicated that caller should drop clients
+                res = UpdateResult::RequiresClientDisconnect;
             }
         }
     }
 
     log::info!("Database update complete");
-    Ok(())
+    Ok(res)
 }
 
 #[cfg(test)]
@@ -359,14 +362,7 @@ mod test {
         // Try to update the db.
         let mut tx = begin_mut_tx(&stdb);
         let plan = ponder_migrate(&old, &new)?;
-        update_database(
-            &stdb,
-            &ModuleSubscriptions::for_test_new_runtime(Arc::new(stdb.db.clone())).0,
-            &mut tx,
-            auth_ctx,
-            plan,
-            &TestLogger,
-        )?;
+        update_database(&stdb, &mut tx, auth_ctx, plan, &TestLogger)?;
 
         // Expect the schema change.
         let idx_b_id = stdb

crates/core/src/host/host_controller.rs

@@ -19,7 +19,7 @@ use crate::worker_metrics::WORKER_METRICS;
 use anyhow::{anyhow, Context};
 use async_trait::async_trait;
 use durability::{Durability, EmptyHistory};
-use log::{info, trace, warn};
+use log::{debug, info, trace, warn};
 use parking_lot::Mutex;
 use spacetimedb_data_structures::error_stream::ErrorStream;
 use spacetimedb_data_structures::map::IntMap;
@@ -414,7 +414,14 @@ impl HostController {
                 )
                 .await?;
 
-            *guard = Some(host);
+            // If hotswap is disabled, we drop the host after the update.
+            // which will drop all connected clients.
+            if !update_result.hotswap_disabled() {
+                *guard = Some(host);
+            } else {
+                debug!("dropping host after update with hotswap disabled");
+            }
+
             Ok::<_, anyhow::Error>(update_result)
         })
         .await??;
@@ -977,7 +984,7 @@ impl Host {
         trace!("update result: {update_result:?}");
         // Only replace the module + scheduler if the update succeeded.
         // Otherwise, we want the database to continue running with the old state.
-        if update_result.was_successful() {
+        if update_result.was_successful() && !update_result.hotswap_disabled() {
             self.scheduler = scheduler;
             scheduler_starter.start(&module)?;
             let old_module = self.module.send_replace(module);

crates/core/src/host/module_host.rs

@@ -506,6 +506,7 @@ pub struct WeakModuleHost {
 pub enum UpdateDatabaseResult {
     NoUpdateNeeded,
     UpdatePerformed,
+    UpdatePerformedWithClientDisconnect,
     AutoMigrateError(ErrorStream<AutoMigrateError>),
     ErrorExecutingMigration(anyhow::Error),
 }
@@ -514,9 +515,16 @@ impl UpdateDatabaseResult {
     pub fn was_successful(&self) -> bool {
         matches!(
             self,
-            UpdateDatabaseResult::UpdatePerformed | UpdateDatabaseResult::NoUpdateNeeded
+            UpdateDatabaseResult::UpdatePerformed
+                | UpdateDatabaseResult::NoUpdateNeeded
+                | UpdateDatabaseResult::UpdatePerformedWithClientDisconnect
         )
     }
+
+    /// Check if hotswap was disabled due to the update.
+    pub fn hotswap_disabled(&self) -> bool {
+        matches!(self, UpdateDatabaseResult::UpdatePerformedWithClientDisconnect)
+    }
 }
 
 #[derive(thiserror::Error, Debug)]

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

@@ -306,14 +306,7 @@ 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,
-            &replica_ctx.subscriptions,
-            &mut tx,
-            auth_ctx,
-            plan,
-            system_logger,
-        );
+        let res = crate::db::update::update_database(stdb, &mut tx, auth_ctx, plan, system_logger);
 
         match res {
             Err(e) => {
@@ -323,13 +316,18 @@ impl InstanceCommon {
                 stdb.report_mut_tx_metrics(reducer, tx_metrics, None);
                 Ok(UpdateDatabaseResult::ErrorExecutingMigration(e))
             }
-            Ok(()) => {
+            Ok(res) => {
                 if let Some((_tx_offset, tx_data, tx_metrics, reducer)) = stdb.commit_tx(tx)? {
                     stdb.report_mut_tx_metrics(reducer, tx_metrics, Some(tx_data));
                 }
                 system_logger.info("Database updated");
                 log::info!("Database updated, {}", stdb.database_identity());
-                Ok(UpdateDatabaseResult::UpdatePerformed)
+                match res {
+                    crate::db::update::UpdateResult::Success => Ok(UpdateDatabaseResult::UpdatePerformed),
+                    crate::db::update::UpdateResult::RequiresClientDisconnect => {
+                        Ok(UpdateDatabaseResult::UpdatePerformedWithClientDisconnect)
+                    }
+                }
             }
         }
     }

crates/core/src/subscription/module_subscription_actor.rs

@@ -853,11 +853,6 @@ 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,12 +779,6 @@ 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> {