move call_identity_disconnected + disconnect_clien to Instance

Author: Centril

crates/core/src/host/module_host.rs

@@ -424,6 +424,24 @@ impl Instance {
         }
     }
 
+    fn call_identity_disconnected(
+        &mut self,
+        caller_identity: Identity,
+        caller_connection_id: ConnectionId,
+    ) -> Result<(), ReducerCallError> {
+        match self {
+            Instance::Wasm(inst) => inst.call_identity_disconnected(caller_identity, caller_connection_id),
+            Instance::Js(inst) => inst.call_identity_disconnected(caller_identity, caller_connection_id),
+        }
+    }
+
+    fn disconnect_client(&mut self, client_id: ClientActorId) -> Result<(), ReducerCallError> {
+        match self {
+            Instance::Wasm(inst) => inst.disconnect_client(client_id),
+            Instance::Js(inst) => inst.disconnect_client(client_id),
+        }
+    }
+
     fn call_scheduled_reducer(&mut self, item: QueueItem) -> Result<ReducerCallResult, ReducerCallError> {
         match self {
             Instance::Wasm(inst) => inst.call_scheduled_reducer(item),
@@ -1103,20 +1121,32 @@ impl ModuleHost {
 
     pub async fn disconnect_client(&self, client_id: ClientActorId) {
         log::trace!("disconnecting client {client_id}");
-        let this = self.clone();
         if let Err(e) = self
-            .call("disconnect_client", move |inst| {
-                // Call the `client_disconnected` reducer, if it exists.
-                // This is a no-op if the module doesn't define such a reducer.
-                this.subscriptions().remove_subscriber(client_id);
-                this.call_identity_disconnected_inner(client_id.identity, client_id.connection_id, inst)
-            })
+            .call("disconnect_client", move |inst| inst.disconnect_client(client_id))
             .await
         {
             log::error!("Error from client_disconnected transaction: {e}");
         }
     }
 
+    pub fn disconnect_client_inner(
+        client_id: ClientActorId,
+        info: &ModuleInfo,
+        call_reducer: impl FnOnce(Option<MutTxId>, CallReducerParams) -> (ReducerCallResult, bool),
+        trapped_slot: &mut bool,
+    ) -> Result<(), ReducerCallError> {
+        // Call the `client_disconnected` reducer, if it exists.
+        // This is a no-op if the module doesn't define such a reducer.
+        info.subscriptions.remove_subscriber(client_id);
+        Self::call_identity_disconnected_inner(
+            client_id.identity,
+            client_id.connection_id,
+            info,
+            call_reducer,
+            trapped_slot,
+        )
+    }
+
     /// Invoke the module's `client_connected` reducer, if it has one,
     /// and insert a new row into `st_client` for `(caller_identity, caller_connection_id)`.
     ///
@@ -1147,12 +1177,15 @@ impl ModuleHost {
     /// If the reducer fails, the rows are still deleted.
     /// Calling this on an already-disconnected client is a no-op.
     pub fn call_identity_disconnected_inner(
-        &self,
         caller_identity: Identity,
         caller_connection_id: ConnectionId,
-        inst: &mut Instance,
+        info: &ModuleInfo,
+        call_reducer: impl FnOnce(Option<MutTxId>, CallReducerParams) -> (ReducerCallResult, bool),
+        trapped_slot: &mut bool,
     ) -> Result<(), ReducerCallError> {
-        let reducer_lookup = self.info.module_def.lifecycle_reducer(Lifecycle::OnDisconnect);
+        let stdb = info.relational_db();
+
+        let reducer_lookup = info.module_def.lifecycle_reducer(Lifecycle::OnDisconnect);
         let reducer_name = reducer_lookup
             .as_ref()
             .map(|(_, def)| &*def.name)
@@ -1162,12 +1195,9 @@ impl ModuleHost {
 
         let workload = || Workload::reducer_no_args(reducer_name, caller_identity, caller_connection_id);
 
-        let me = self.clone();
-        let stdb = me.module.replica_ctx().relational_db.clone();
-
         // A fallback transaction that deletes the client from `st_client`.
+        let database_identity = stdb.database_identity();
         let fallback = || {
-            let database_identity = me.info.database_identity;
             stdb.with_auto_commit(workload(), |mut_tx| {
                 if !is_client_exist(mut_tx) {
                     // The client is already gone. Nothing to do.
@@ -1207,8 +1237,9 @@ impl ModuleHost {
             // The module defined a lifecycle reducer to handle disconnects. Call it.
             // If it succeeds, `WasmModuleInstance::call_reducer_with_tx` has already ensured
             // that `st_client` is updated appropriately.
-            let result = me.call_reducer_inner_with_inst(
-                Some(mut_tx),
+            let tx = Some(mut_tx);
+            let result = Self::call_reducer_params(
+                info,
                 caller_identity,
                 Some(caller_connection_id),
                 None,
@@ -1217,8 +1248,12 @@ impl ModuleHost {
                 reducer_id,
                 reducer_def,
                 FunctionArgs::Nullary,
-                inst,
-            );
+            )
+            .map(|params| {
+                let (res, trapped) = call_reducer(tx, params);
+                *trapped_slot = trapped;
+                res
+            });
 
             // If it failed, we still need to update `st_client`: the client's not coming back.
             // Commit a separate transaction that just updates `st_client`.
@@ -1270,9 +1305,8 @@ impl ModuleHost {
         caller_identity: Identity,
         caller_connection_id: ConnectionId,
     ) -> Result<(), ReducerCallError> {
-        let me = self.clone();
         self.call("call_identity_disconnected", move |inst| {
-            me.call_identity_disconnected_inner(caller_identity, caller_connection_id, inst)
+            inst.call_identity_disconnected(caller_identity, caller_connection_id)
         })
         .await?
     }
@@ -1340,37 +1374,6 @@ impl ModuleHost {
             })
             .await?)
     }
-    fn call_reducer_inner_with_inst(
-        &self,
-        tx: Option<MutTxId>,
-        caller_identity: Identity,
-        caller_connection_id: Option<ConnectionId>,
-        client: Option<Arc<ClientConnectionSender>>,
-        request_id: Option<RequestId>,
-        timer: Option<Instant>,
-        reducer_id: ReducerId,
-        reducer_def: &ReducerDef,
-        args: FunctionArgs,
-        module_instance: &mut Instance,
-    ) -> Result<ReducerCallResult, ReducerCallError> {
-        let reducer_seed = ArgsSeed(self.info.module_def.typespace().with_type(reducer_def));
-        let args = args.into_tuple(reducer_seed).map_err(InvalidReducerArguments)?;
-        let caller_connection_id = caller_connection_id.unwrap_or(ConnectionId::ZERO);
-
-        Ok(module_instance.call_reducer(
-            tx,
-            CallReducerParams {
-                timestamp: Timestamp::now(),
-                caller_identity,
-                caller_connection_id,
-                client,
-                request_id,
-                timer,
-                reducer_id,
-                args,
-            },
-        ))
-    }
 
     pub async fn call_reducer(
         &self,

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

@@ -9,6 +9,7 @@ use self::syscall::{call_call_reducer, call_describe_module, call_reducer_fun, r
 use super::module_common::{build_common_module_from_raw, run_describer, ModuleCommon};
 use super::module_host::{CallProcedureParams, CallReducerParams, Module, ModuleInfo, ModuleRuntime};
 use super::UpdateDatabaseResult;
+use crate::client::ClientActorId;
 use crate::host::instance_env::{ChunkPool, InstanceEnv};
 use crate::host::module_host::{
     call_identity_connected, call_scheduled_reducer, init_database, ClientConnectedError, Instance,
@@ -17,7 +18,7 @@ use crate::host::scheduler::QueueItem;
 use crate::host::wasm_common::instrumentation::CallTimes;
 use crate::host::wasm_common::module_host_actor::{DescribeError, ExecuteResult, ExecutionTimings, InstanceCommon};
 use crate::host::wasm_common::{RowIters, TimingSpanSet};
-use crate::host::{ReducerCallError, ReducerCallResult, Scheduler};
+use crate::host::{ModuleHost, ReducerCallError, ReducerCallResult, Scheduler};
 use crate::module_host_context::{ModuleCreationContext, ModuleCreationContextLimited};
 use crate::replica_context::ReplicaContext;
 use crate::util::asyncify;
@@ -28,7 +29,7 @@ use itertools::Either;
 use spacetimedb_auth::identity::ConnectionAuthCtx;
 use spacetimedb_datastore::locking_tx_datastore::MutTxId;
 use spacetimedb_datastore::traits::Program;
-use spacetimedb_lib::{ConnectionId, RawModuleDef, Timestamp};
+use spacetimedb_lib::{ConnectionId, Identity, RawModuleDef, Timestamp};
 use spacetimedb_schema::auto_migrate::MigrationPolicy;
 use std::sync::mpsc::{Receiver, SyncSender};
 use std::sync::{mpsc, Arc, LazyLock};
@@ -333,6 +334,28 @@ impl JsInstance {
         })
     }
 
+    pub fn call_identity_disconnected(
+        &mut self,
+        caller_identity: Identity,
+        caller_connection_id: ConnectionId,
+    ) -> Result<(), ReducerCallError> {
+        self.can_trap(move |this| {
+            this.send_recv(
+                JsWorkerReply::into_call_identity_disconnected,
+                JsWorkerRequest::CallIdentityDisconnected(caller_identity, caller_connection_id),
+            )
+        })
+    }
+
+    pub fn disconnect_client(&mut self, client_id: ClientActorId) -> Result<(), ReducerCallError> {
+        self.can_trap(move |this| {
+            this.send_recv(
+                JsWorkerReply::into_disconnect_client,
+                JsWorkerRequest::DisconnectClient(client_id),
+            )
+        })
+    }
+
     pub(crate) fn call_scheduled_reducer(&mut self, item: QueueItem) -> Result<ReducerCallResult, ReducerCallError> {
         self.can_trap(move |this| {
             this.send_recv(
@@ -366,6 +389,8 @@ enum JsWorkerReply {
     CallReducer((ReducerCallResult, bool)),
     ClearAllClients(anyhow::Result<()>),
     CallIdentityConnected((Result<(), ClientConnectedError>, bool)),
+    CallIdentityDisconnected((Result<(), ReducerCallError>, bool)),
+    DisconnectClient((Result<(), ReducerCallError>, bool)),
     CallScheduledReducer((Result<ReducerCallResult, ReducerCallError>, bool)),
     InitDatabase((anyhow::Result<Option<ReducerCallResult>>, bool)),
 }
@@ -390,6 +415,10 @@ enum JsWorkerRequest {
     ClearAllClients,
     /// See [`JsInstance::call_identity_connected`].
     CallIdentityConnected(ConnectionAuthCtx, ConnectionId),
+    /// See [`JsInstance::call_identity_disconnected`].
+    CallIdentityDisconnected(Identity, ConnectionId),
+    /// See [`JsInstance::disconnect_client`].
+    DisconnectClient(ClientActorId),
     /// See [`JsInstance::call_scheduled_reducer`].
     CallScheduledReducer(QueueItem),
     /// See [`JsInstance::init_database`].
@@ -503,6 +532,9 @@ fn spawn_instance_worker(
             }
         };
         for request in request_rx.iter() {
+            let mut call_reducer =
+                |tx, params| call_reducer(&mut instance_common, replica_ctx, scope, call_reducer_fun, tx, params);
+
             match request {
                 JsWorkerRequest::UpdateDatabase {
                     program,
@@ -519,34 +551,46 @@ fn spawn_instance_worker(
                     // but rather let this happen by `return_instance` using `JsInstance::trapped`
                     // which will cause `JsInstance` to be dropped,
                     // which in turn results in the loop being terminated.
-                    let res = call_reducer(&mut instance_common, replica_ctx, scope, call_reducer_fun, tx, params);
+                    let res = call_reducer(tx, params);
                     reply("call_reducer", JsWorkerReply::CallReducer(res));
                 }
                 JsWorkerRequest::ClearAllClients => {
                     let res = instance_common.clear_all_clients();
                     reply("clear_all_clients", JsWorkerReply::ClearAllClients(res));
                 }
                 JsWorkerRequest::CallIdentityConnected(caller_auth, caller_connection_id) => {
-                    let call_reducer = |tx, params| {
-                        call_reducer(&mut instance_common, replica_ctx, scope, call_reducer_fun, tx, params)
-                    };
                     let mut trapped = false;
                     let res =
                         call_identity_connected(caller_auth, caller_connection_id, info, call_reducer, &mut trapped);
                     let res = (res, trapped);
                     reply("call_identity_connected", JsWorkerReply::CallIdentityConnected(res));
                 }
+                JsWorkerRequest::CallIdentityDisconnected(caller_identity, caller_connection_id) => {
+                    let mut trapped = false;
+                    let res = ModuleHost::call_identity_disconnected_inner(
+                        caller_identity,
+                        caller_connection_id,
+                        info,
+                        call_reducer,
+                        &mut trapped,
+                    );
+                    let res = (res, trapped);
+                    reply(
+                        "call_identity_disconnected",
+                        JsWorkerReply::CallIdentityDisconnected(res),
+                    );
+                }
+                JsWorkerRequest::DisconnectClient(client_id) => {
+                    let mut trapped = false;
+                    let res = ModuleHost::disconnect_client_inner(client_id, info, call_reducer, &mut trapped);
+                    let res = (res, trapped);
+                    reply("disconnect_client", JsWorkerReply::DisconnectClient(res));
+                }
                 JsWorkerRequest::CallScheduledReducer(queue_item) => {
-                    let call_reducer = |tx, params| {
-                        call_reducer(&mut instance_common, replica_ctx, scope, call_reducer_fun, tx, params)
-                    };
                     let res = call_scheduled_reducer(info, queue_item, call_reducer);
                     reply("call_scheduled_reducer", JsWorkerReply::CallScheduledReducer(res));
                 }
                 JsWorkerRequest::InitDatabase(program) => {
-                    let call_reducer = |tx, params| {
-                        call_reducer(&mut instance_common, replica_ctx, scope, call_reducer_fun, tx, params)
-                    };
                     let res = init_database(replica_ctx, &module_common.info().module_def, program, call_reducer);
                     reply("init_database", JsWorkerReply::InitDatabase(res));
                 }

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

@@ -11,7 +11,7 @@ use std::time::Duration;
 use tracing::span::EnteredSpan;
 
 use super::instrumentation::CallTimes;
-use crate::client::ClientConnectionSender;
+use crate::client::{ClientActorId, ClientConnectionSender};
 use crate::database_logger;
 use crate::energy::{EnergyMonitor, ReducerBudget, ReducerFingerprint};
 use crate::host::instance_env::InstanceEnv;
@@ -295,6 +295,34 @@ impl<T: WasmInstance> WasmModuleInstance<T> {
         res
     }
 
+    pub fn call_identity_disconnected(
+        &mut self,
+        caller_identity: Identity,
+        caller_connection_id: ConnectionId,
+    ) -> Result<(), ReducerCallError> {
+        let module = &self.common.info.clone();
+        let call_reducer = |tx, params| self.call_reducer_inner(tx, params);
+        let mut trapped = false;
+        let res = ModuleHost::call_identity_disconnected_inner(
+            caller_identity,
+            caller_connection_id,
+            module,
+            call_reducer,
+            &mut trapped,
+        );
+        self.trapped = trapped;
+        res
+    }
+
+    pub fn disconnect_client(&mut self, client_id: ClientActorId) -> Result<(), ReducerCallError> {
+        let module = &self.common.info.clone();
+        let call_reducer = |tx, params| self.call_reducer_inner(tx, params);
+        let mut trapped = false;
+        let res = ModuleHost::disconnect_client_inner(client_id, module, call_reducer, &mut trapped);
+        self.trapped = trapped;
+        res
+    }
+
     pub(crate) fn call_scheduled_reducer(&mut self, item: QueueItem) -> Result<ReducerCallResult, ReducerCallError> {
         let module = &self.common.info.clone();
         let call_reducer = |tx, params| self.call_reducer_inner(tx, params);