Fix module hotswapping for connected clients (#3159)

The `Clone` impl for `ClientConnection` would create an independent
instance that could not observe module hotswapping. This would result in
methods called on a replaced `ModuleHost` to fail, because that host
exited already.

Fix by reading the `ModuleHost` from the watch channel directly, instead
of maintaining a redundant copy.

Also fix `watch_module_host` to properly mark the current module host as
seen.

# Expected complexity level and risk

2

# Testing

- [x] test suite passes
- [x] ran @bfops repro script

Author: kim

crates/bench/src/spacetime_module.rs

@@ -69,10 +69,11 @@ impl<L: ModuleLanguage> BenchDatabase for SpacetimeModule<L> {
             L::get_module().load_module(config, Some(&path)).await
         });
 
-        for table in module.client.module.info.module_def.tables() {
+        let module_info = module.client.module().info;
+        for table in module_info.module_def.tables() {
             log::trace!("SPACETIME_MODULE: LOADED TABLE: {table:?}");
         }
-        for reducer in module.client.module.info.module_def.reducers() {
+        for reducer in module_info.module_def.reducers() {
             log::trace!("SPACETIME_MODULE: LOADED REDUCER: {reducer:?}");
         }
         Ok(SpacetimeModule {
@@ -102,7 +103,7 @@ impl<L: ModuleLanguage> BenchDatabase for SpacetimeModule<L> {
             module.call_reducer_binary(&name, ProductValue::new(&[])).await?;
             */
             // workaround for now
-            module.client.module.clear_table(&table_id.pascal_case)?;
+            module.client.module().clear_table(&table_id.pascal_case)?;
             Ok(())
         })
     }

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

@@ -359,7 +359,7 @@ async fn ws_client_actor_inner(
     ws: WebSocketStream,
     sendrx: MeteredReceiver<SerializableMessage>,
 ) {
-    let database = client.module.info().database_identity;
+    let database = client.module().info().database_identity;
     let client_id = client.id;
     let client_closed_metric = WORKER_METRICS.ws_clients_closed_connection.with_label_values(&database);
     let state = Arc::new(ActorState::new(database, client_id, config));

crates/core/src/client/client_connection.rs

@@ -213,7 +213,6 @@ impl ClientConnectionSender {
 pub struct ClientConnection {
     sender: Arc<ClientConnectionSender>,
     pub replica_id: u64,
-    pub module: ModuleHost,
     module_rx: watch::Receiver<ModuleHost>,
 }
 
@@ -480,7 +479,6 @@ impl ClientConnection {
         let this = Self {
             sender,
             replica_id,
-            module,
             module_rx,
         };
 
@@ -495,13 +493,11 @@ impl ClientConnection {
         id: ClientActorId,
         config: ClientConfig,
         replica_id: u64,
-        mut module_rx: watch::Receiver<ModuleHost>,
+        module_rx: watch::Receiver<ModuleHost>,
     ) -> Self {
-        let module = module_rx.borrow_and_update().clone();
         Self {
             sender: Arc::new(ClientConnectionSender::dummy(id, config)),
             replica_id,
-            module,
             module_rx,
         }
     }
@@ -510,6 +506,20 @@ impl ClientConnection {
         self.sender.clone()
     }
 
+    /// Get the [`ModuleHost`] for this connection.
+    ///
+    /// Note that modules can be hotswapped, in which case the returned handle
+    /// becomes invalid (i.e. all calls on it will result in an error).
+    /// Callers should thus drop the value as soon as they are done, and obtain
+    /// a fresh one when needed.
+    ///
+    /// While this [`ClientConnection`] is active, [`Self::watch_module_host`]
+    /// should be polled in the background, and the connection closed if and
+    /// when it returns an error.
+    pub fn module(&self) -> ModuleHost {
+        self.module_rx.borrow().clone()
+    }
+
     #[inline]
     pub fn handle_message(
         &self,
@@ -519,13 +529,26 @@ impl ClientConnection {
         message_handlers::handle(self, message.into(), timer)
     }
 
+    /// Waits until the [`ModuleHost`] of this [`ClientConnection`] instance
+    /// exits, in which case `Err` containing [`NoSuchModule`] is returned.
+    ///
+    /// Should be polled while this [`ClientConnection`] is active, so as to be
+    /// able to shut down the connection gracefully if and when the module
+    /// exits.
+    ///
+    /// Note that this borrows `self` mutably, so may require cloning the
+    /// [`ClientConnection`] instance. The module is shared, however, so all
+    /// clones will observe a swapped module.
     pub async fn watch_module_host(&mut self) -> Result<(), NoSuchModule> {
-        match self.module_rx.changed().await {
-            Ok(()) => {
-                self.module = self.module_rx.borrow_and_update().clone();
-                Ok(())
+        loop {
+            // First check if the module exited between creating the client
+            // connection and calling `watch_module_host`...
+            if self.module_rx.changed().await.is_err() {
+                return Err(NoSuchModule);
             }
-            Err(_) => Err(NoSuchModule),
+            // ...then mark the current module as seen, so the next iteration
+            // of the loop waits until the module changes or exits.
+            self.module_rx.mark_unchanged();
         }
     }
 
@@ -544,7 +567,7 @@ impl ClientConnection {
             CallReducerFlags::NoSuccessNotify => None,
         };
 
-        self.module
+        self.module()
             .call_reducer(
                 self.id.identity,
                 Some(self.id.connection_id),
@@ -563,9 +586,9 @@ impl ClientConnection {
         timer: Instant,
     ) -> Result<Option<ExecutionMetrics>, DBError> {
         let me = self.clone();
-        self.module
+        self.module()
             .on_module_thread("subscribe_single", move || {
-                me.module
+                me.module()
                     .subscriptions()
                     .add_single_subscription(me.sender, subscription, timer, None)
             })
@@ -575,7 +598,7 @@ impl ClientConnection {
     pub async fn unsubscribe(&self, request: Unsubscribe, timer: Instant) -> Result<Option<ExecutionMetrics>, DBError> {
         let me = self.clone();
         asyncify(move || {
-            me.module
+            me.module()
                 .subscriptions()
                 .remove_single_subscription(me.sender, request, timer)
         })
@@ -588,9 +611,9 @@ impl ClientConnection {
         timer: Instant,
     ) -> Result<Option<ExecutionMetrics>, DBError> {
         let me = self.clone();
-        self.module
+        self.module()
             .on_module_thread("subscribe_multi", move || {
-                me.module
+                me.module()
                     .subscriptions()
                     .add_multi_subscription(me.sender, request, timer, None)
             })
@@ -603,9 +626,9 @@ impl ClientConnection {
         timer: Instant,
     ) -> Result<Option<ExecutionMetrics>, DBError> {
         let me = self.clone();
-        self.module
+        self.module()
             .on_module_thread("unsubscribe_multi", move || {
-                me.module
+                me.module()
                     .subscriptions()
                     .remove_multi_subscription(me.sender, request, timer)
             })
@@ -615,7 +638,7 @@ impl ClientConnection {
     pub async fn subscribe(&self, subscription: Subscribe, timer: Instant) -> Result<ExecutionMetrics, DBError> {
         let me = self.clone();
         asyncify(move || {
-            me.module
+            me.module()
                 .subscriptions()
                 .add_legacy_subscriber(me.sender, subscription, timer, None)
         })
@@ -628,7 +651,7 @@ impl ClientConnection {
         message_id: &[u8],
         timer: Instant,
     ) -> Result<(), anyhow::Error> {
-        self.module
+        self.module()
             .one_off_query::<JsonFormat>(
                 self.id.identity,
                 query.to_owned(),
@@ -646,7 +669,7 @@ impl ClientConnection {
         message_id: &[u8],
         timer: Instant,
     ) -> Result<(), anyhow::Error> {
-        self.module
+        self.module()
             .one_off_query::<BsatnFormat>(
                 self.id.identity,
                 query.to_owned(),
@@ -659,6 +682,6 @@ impl ClientConnection {
     }
 
     pub async fn disconnect(self) {
-        self.module.disconnect_client(self.id).await
+        self.module().disconnect_client(self.id).await
     }
 }

crates/core/src/client/message_handlers.rs

@@ -46,10 +46,11 @@ pub async fn handle(client: &ClientConnection, message: DataMessage, timer: Inst
             .map_args(|b| ReducerArgs::Bsatn(message_buf.slice_ref(b))),
     };
 
-    let mod_info = client.module.info();
+    let module = client.module();
+    let mod_info = module.info();
     let mod_metrics = &mod_info.metrics;
     let database_identity = mod_info.database_identity;
-    let db = &client.module.replica_ctx().relational_db;
+    let db = &module.replica_ctx().relational_db;
     let record_metrics = |wl| {
         move |metrics| {
             if let Some(metrics) = metrics {