Close scheduler on module exit

crates/core/src/host/module_host.rs

@@ -1,4 +1,4 @@
-use super::{ArgsTuple, InvalidReducerArguments, ReducerArgs, ReducerCallResult, ReducerId};
+use super::{ArgsTuple, InvalidReducerArguments, ReducerArgs, ReducerCallResult, ReducerId, Scheduler};
 use crate::client::{ClientActorId, ClientConnectionSender};
 use crate::database_logger::{LogLevel, Record};
 use crate::db::datastore::locking_tx_datastore::MutTxId;
@@ -16,7 +16,7 @@ use crate::sql::ast::SchemaViewer;
 use crate::subscription::module_subscription_actor::ModuleSubscriptions;
 use crate::subscription::tx::DeltaTx;
 use crate::subscription::{execute_plan, record_exec_metrics};
-use crate::util::lending_pool::{Closed, LendingPool, LentResource, PoolClosed};
+use crate::util::lending_pool::{LendingPool, LentResource, PoolClosed};
 use crate::vm::check_row_limit;
 use crate::worker_metrics::WORKER_METRICS;
 use anyhow::Context;
@@ -253,7 +253,7 @@ pub trait Module: Send + Sync + 'static {
     fn info(&self) -> Arc<ModuleInfo>;
     fn create_instance(&self) -> Self::Instance;
     fn replica_ctx(&self) -> &ReplicaContext;
-    fn close(self);
+    fn scheduler(&self) -> &Scheduler;
 }
 
 pub trait ModuleInstance: Send + 'static {
@@ -344,8 +344,8 @@ impl fmt::Debug for ModuleHost {
 trait DynModuleHost: Send + Sync + 'static {
     async fn get_instance(&self, db: Identity) -> Result<Box<dyn ModuleInstance>, NoSuchModule>;
     fn replica_ctx(&self) -> &ReplicaContext;
-    fn exit(&self) -> Closed<'_>;
-    fn exited(&self) -> Closed<'_>;
+    async fn exit(&self);
+    async fn exited(&self);
 }
 
 struct HostControllerActor<T: Module> {
@@ -406,12 +406,14 @@ impl<T: Module> DynModuleHost for HostControllerActor<T> {
         self.module.replica_ctx()
     }
 
-    fn exit(&self) -> Closed<'_> {
-        self.instance_pool.close()
+    async fn exit(&self) {
+        self.module.scheduler().close();
+        self.instance_pool.close();
+        self.exited().await
     }
 
-    fn exited(&self) -> Closed<'_> {
-        self.instance_pool.closed()
+    async fn exited(&self) {
+        tokio::join!(self.module.scheduler().closed(), self.instance_pool.closed());
     }
 }
 

crates/core/src/host/scheduler.rs

@@ -233,6 +233,10 @@ impl Scheduler {
     pub fn close(&self) {
         let _ = self.tx.send(MsgOrExit::Exit);
     }
+
+    pub async fn closed(&self) {
+        self.tx.closed().await
+    }
 }
 
 struct SchedulerActor {

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

@@ -230,8 +230,8 @@ impl<T: WasmModule> Module for WasmModuleHostActor<T> {
         &self.replica_context
     }
 
-    fn close(self) {
-        self.scheduler.close()
+    fn scheduler(&self) -> &Scheduler {
+        &self.scheduler
     }
 }
 

crates/core/src/util/lending_pool.rs

@@ -129,7 +129,7 @@ impl<T> LendingPool<T> {
         self.sem.available_permits()
     }
 
-    pub fn close(&self) -> Closed<'_> {
+    pub fn close(&self) {
         let mut vec = self.inner.vec.lock();
         self.sem.close();
         if let Some(deque) = vec.deque.take() {
@@ -138,7 +138,6 @@ impl<T> LendingPool<T> {
         if vec.total_count == 0 {
             self.inner.closed_notify.notify();
         }
-        self.closed()
     }
 
     pub fn closed(&self) -> Closed<'_> {