move call_identity_connected to Instance

Author: Centril

crates/core/src/host/module_host.rs

@@ -413,10 +413,14 @@ impl Instance {
         }
     }
 
-    async fn call_procedure(&mut self, params: CallProcedureParams) -> Result<ProcedureCallResult, ProcedureCallError> {
+    fn call_identity_connected(
+        &mut self,
+        caller_auth: ConnectionAuthCtx,
+        caller_connection_id: ConnectionId,
+    ) -> Result<(), ClientConnectedError> {
         match self {
-            Instance::Wasm(inst) => inst.call_procedure(params).await,
-            Instance::Js(inst) => inst.call_procedure(params).await,
+            Instance::Wasm(inst) => inst.call_identity_connected(caller_auth, caller_connection_id),
+            Instance::Js(inst) => inst.call_identity_connected(caller_auth, caller_connection_id),
         }
     }
 
@@ -433,6 +437,13 @@ impl Instance {
             Instance::Js(inst) => inst.init_database(program),
         }
     }
+
+    async fn call_procedure(&mut self, params: CallProcedureParams) -> Result<ProcedureCallResult, ProcedureCallError> {
+        match self {
+            Instance::Wasm(inst) => inst.call_procedure(params).await,
+            Instance::Js(inst) => inst.call_procedure(params).await,
+        }
+    }
 }
 
 /// Creates the table for `table_def` in `stdb`.
@@ -550,6 +561,95 @@ fn init_database_inner(
     Ok(rcr)
 }
 
+pub fn call_identity_connected(
+    caller_auth: ConnectionAuthCtx,
+    caller_connection_id: ConnectionId,
+    module: &ModuleInfo,
+    call_reducer: impl FnOnce(Option<MutTxId>, CallReducerParams) -> (ReducerCallResult, bool),
+    trapped_slot: &mut bool,
+) -> Result<(), ClientConnectedError> {
+    let reducer_lookup = module.module_def.lifecycle_reducer(Lifecycle::OnConnect);
+    let stdb = module.relational_db();
+    let workload = Workload::reducer_no_args(
+        "call_identity_connected",
+        caller_auth.claims.identity,
+        caller_connection_id,
+    );
+    let mut_tx = stdb.begin_mut_tx(IsolationLevel::Serializable, workload);
+    let mut mut_tx = scopeguard::guard(mut_tx, |mut_tx| {
+        // If we crash before committing, we need to ensure that the transaction is rolled back.
+        // This is necessary to avoid leaving the database in an inconsistent state.
+        log::debug!("call_identity_connected: rolling back transaction");
+        let (_, metrics, reducer_name) = mut_tx.rollback();
+        stdb.report_mut_tx_metrics(reducer_name, metrics, None);
+    });
+
+    mut_tx
+        .insert_st_client(
+            caller_auth.claims.identity,
+            caller_connection_id,
+            &caller_auth.jwt_payload,
+        )
+        .map_err(DBError::from)?;
+
+    if let Some((reducer_id, reducer_def)) = reducer_lookup {
+        // The module defined a lifecycle reducer to handle new connections.
+        // Call this reducer.
+        // If the call fails (as in, something unexpectedly goes wrong with guest execution),
+        // abort the connection: we can't really recover.
+        let tx = Some(ScopeGuard::into_inner(mut_tx));
+        let params = ModuleHost::call_reducer_params(
+            module,
+            caller_auth.claims.identity,
+            Some(caller_connection_id),
+            None,
+            None,
+            None,
+            reducer_id,
+            reducer_def,
+            FunctionArgs::Nullary,
+        )
+        .map_err(ReducerCallError::from)?;
+        let (reducer_outcome, trapped) = call_reducer(tx, params);
+        *trapped_slot = trapped;
+
+        match reducer_outcome.outcome {
+            // If the reducer committed successfully, we're done.
+            // `WasmModuleInstance::call_reducer_with_tx` has already ensured
+            // that `st_client` is updated appropriately.
+            //
+            // It's necessary to spread out the responsibility for updating `st_client` in this way
+            // because it's important that `call_identity_connected` commit at most one transaction.
+            // A naive implementation of this method would just run the reducer first,
+            // then insert into `st_client`,
+            // but if we crashed in between, we'd be left in an inconsistent state
+            // where the reducer had run but `st_client` was not yet updated.
+            ReducerOutcome::Committed => Ok(()),
+
+            // If the reducer returned an error or couldn't run due to insufficient energy,
+            // abort the connection: the module code has decided it doesn't want this client.
+            ReducerOutcome::Failed(message) => Err(ClientConnectedError::Rejected(message)),
+            ReducerOutcome::BudgetExceeded => Err(ClientConnectedError::OutOfEnergy),
+        }
+    } else {
+        // The module doesn't define a client_connected reducer.
+        // We need to commit the transaction to update st_clients and st_connection_credentials.
+        //
+        // This is necessary to be able to disconnect clients after a server crash.
+
+        // TODO: Is this being broadcast? Does it need to be, or are st_client table subscriptions
+        // not allowed?
+        // I (jsdt) don't think it was being broadcast previously. See:
+        // https://github.com/clockworklabs/SpacetimeDB/issues/3130
+        stdb.finish_tx(ScopeGuard::into_inner(mut_tx), Ok(()))
+            .map_err(|e: DBError| {
+                log::error!("`call_identity_connected`: finish transaction failed: {e:#?}");
+                ClientConnectedError::DBError(e)
+            })?;
+        Ok(())
+    }
+}
+
 // Only for logging purposes.
 const SCHEDULED_REDUCER: &str = "scheduled_reducer";
 
@@ -1035,85 +1135,8 @@ impl ModuleHost {
         caller_auth: ConnectionAuthCtx,
         caller_connection_id: ConnectionId,
     ) -> Result<(), ClientConnectedError> {
-        let me = self.clone();
         self.call("call_identity_connected", move |inst| {
-            let reducer_lookup = me.info.module_def.lifecycle_reducer(Lifecycle::OnConnect);
-            let stdb = &me.module.replica_ctx().relational_db;
-            let workload = Workload::reducer_no_args(
-                "call_identity_connected",
-                caller_auth.claims.identity,
-                caller_connection_id,
-            );
-            let mut_tx = stdb.begin_mut_tx(IsolationLevel::Serializable, workload);
-            let mut mut_tx = scopeguard::guard(mut_tx, |mut_tx| {
-                // If we crash before committing, we need to ensure that the transaction is rolled back.
-                // This is necessary to avoid leaving the database in an inconsistent state.
-                log::debug!("call_identity_connected: rolling back transaction");
-                let (_, metrics, reducer_name) = mut_tx.rollback();
-                stdb.report_mut_tx_metrics(reducer_name, metrics, None);
-            });
-
-            mut_tx
-                .insert_st_client(
-                    caller_auth.claims.identity,
-                    caller_connection_id,
-                    &caller_auth.jwt_payload,
-                )
-                .map_err(DBError::from)?;
-
-            if let Some((reducer_id, reducer_def)) = reducer_lookup {
-                // The module defined a lifecycle reducer to handle new connections.
-                // Call this reducer.
-                // If the call fails (as in, something unexpectedly goes wrong with guest execution),
-                // abort the connection: we can't really recover.
-                let reducer_outcome = me.call_reducer_inner_with_inst(
-                    Some(ScopeGuard::into_inner(mut_tx)),
-                    caller_auth.claims.identity,
-                    Some(caller_connection_id),
-                    None,
-                    None,
-                    None,
-                    reducer_id,
-                    reducer_def,
-                    FunctionArgs::Nullary,
-                    inst,
-                )?;
-
-                match reducer_outcome.outcome {
-                    // If the reducer committed successfully, we're done.
-                    // `WasmModuleInstance::call_reducer_with_tx` has already ensured
-                    // that `st_client` is updated appropriately.
-                    //
-                    // It's necessary to spread out the responsibility for updating `st_client` in this way
-                    // because it's important that `call_identity_connected` commit at most one transaction.
-                    // A naive implementation of this method would just run the reducer first,
-                    // then insert into `st_client`,
-                    // but if we crashed in between, we'd be left in an inconsistent state
-                    // where the reducer had run but `st_client` was not yet updated.
-                    ReducerOutcome::Committed => Ok(()),
-
-                    // If the reducer returned an error or couldn't run due to insufficient energy,
-                    // abort the connection: the module code has decided it doesn't want this client.
-                    ReducerOutcome::Failed(message) => Err(ClientConnectedError::Rejected(message)),
-                    ReducerOutcome::BudgetExceeded => Err(ClientConnectedError::OutOfEnergy),
-                }
-            } else {
-                // The module doesn't define a client_connected reducer.
-                // We need to commit the transaction to update st_clients and st_connection_credentials.
-                //
-                // This is necessary to be able to disconnect clients after a server crash.
-
-                // TODO: Is this being broadcast? Does it need to be, or are st_client table subscriptions
-                // not allowed?
-                // I (jsdt) don't think it was being broadcast previously. See:
-                // https://github.com/clockworklabs/SpacetimeDB/issues/3130
-                stdb.finish_tx(ScopeGuard::into_inner(mut_tx), Ok(()))
-                    .map_err(|e: DBError| {
-                        log::error!("`call_identity_connected`: finish transaction failed: {e:#?}");
-                        ClientConnectedError::DBError(e)
-                    })?;
-                Ok(())
-            }
+            inst.call_identity_connected(caller_auth, caller_connection_id)
         })
         .await
         .map_err(ReducerCallError::from)?
@@ -1260,6 +1283,32 @@ impl ModuleHost {
             .await?
     }
 
+    fn call_reducer_params(
+        module: &ModuleInfo,
+        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,
+    ) -> Result<CallReducerParams, InvalidReducerArguments> {
+        let reducer_seed = ArgsSeed(module.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(CallReducerParams {
+            timestamp: Timestamp::now(),
+            caller_identity,
+            caller_connection_id,
+            client,
+            request_id,
+            timer,
+            reducer_id,
+            args,
+        })
+    }
+
     async fn call_reducer_inner(
         &self,
         caller_identity: Identity,

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

@@ -10,7 +10,9 @@ use super::module_common::{build_common_module_from_raw, run_describer, ModuleCo
 use super::module_host::{CallProcedureParams, CallReducerParams, Module, ModuleInfo, ModuleRuntime};
 use super::UpdateDatabaseResult;
 use crate::host::instance_env::{ChunkPool, InstanceEnv};
-use crate::host::module_host::{call_scheduled_reducer, init_database, Instance};
+use crate::host::module_host::{
+    call_identity_connected, call_scheduled_reducer, init_database, ClientConnectedError, Instance,
+};
 use crate::host::scheduler::QueueItem;
 use crate::host::wasm_common::instrumentation::CallTimes;
 use crate::host::wasm_common::module_host_actor::{DescribeError, ExecuteResult, ExecutionTimings, InstanceCommon};
@@ -23,9 +25,10 @@ use core::any::type_name;
 use core::str;
 use enum_as_inner::EnumAsInner;
 use itertools::Either;
+use spacetimedb_auth::identity::ConnectionAuthCtx;
 use spacetimedb_datastore::locking_tx_datastore::MutTxId;
 use spacetimedb_datastore::traits::Program;
-use spacetimedb_lib::{RawModuleDef, Timestamp};
+use spacetimedb_lib::{ConnectionId, RawModuleDef, Timestamp};
 use spacetimedb_schema::auto_migrate::MigrationPolicy;
 use std::sync::mpsc::{Receiver, SyncSender};
 use std::sync::{mpsc, Arc, LazyLock};
@@ -317,6 +320,19 @@ impl JsInstance {
         self.send_recv(JsWorkerReply::into_clear_all_clients, JsWorkerRequest::ClearAllClients)
     }
 
+    pub fn call_identity_connected(
+        &mut self,
+        caller_auth: ConnectionAuthCtx,
+        caller_connection_id: ConnectionId,
+    ) -> Result<(), ClientConnectedError> {
+        self.can_trap(move |this| {
+            this.send_recv(
+                JsWorkerReply::into_call_identity_connected,
+                JsWorkerRequest::CallIdentityConnected(caller_auth, caller_connection_id),
+            )
+        })
+    }
+
     pub(crate) fn call_scheduled_reducer(&mut self, item: QueueItem) -> Result<ReducerCallResult, ReducerCallError> {
         self.can_trap(move |this| {
             this.send_recv(
@@ -349,6 +365,7 @@ enum JsWorkerReply {
     UpdateDatabase(anyhow::Result<UpdateDatabaseResult>),
     CallReducer((ReducerCallResult, bool)),
     ClearAllClients(anyhow::Result<()>),
+    CallIdentityConnected((Result<(), ClientConnectedError>, bool)),
     CallScheduledReducer((Result<ReducerCallResult, ReducerCallError>, bool)),
     InitDatabase((anyhow::Result<Option<ReducerCallResult>>, bool)),
 }
@@ -371,6 +388,8 @@ enum JsWorkerRequest {
     },
     /// See [`JsInstance::clear_all_clients`].
     ClearAllClients,
+    /// See [`JsInstance::call_identity_connected`].
+    CallIdentityConnected(ConnectionAuthCtx, ConnectionId),
     /// See [`JsInstance::call_scheduled_reducer`].
     CallScheduledReducer(QueueItem),
     /// See [`JsInstance::init_database`].
@@ -465,6 +484,7 @@ fn spawn_instance_worker(
         };
 
         // Setup the instance common and environment.
+        let info = &module_common.info();
         let mut instance_common = InstanceCommon::new(&module_common);
         let replica_ctx: &Arc<ReplicaContext> = module_common.replica_ctx();
         let scheduler = module_common.scheduler().clone();
@@ -506,12 +526,21 @@ fn spawn_instance_worker(
                     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::CallScheduledReducer(queue_item) => {
                     let call_reducer = |tx, params| {
                         call_reducer(&mut instance_common, replica_ctx, scope, call_reducer_fun, tx, params)
                     };
-                    let module = module_common.info();
-                    let res = call_scheduled_reducer(&module, queue_item, call_reducer);
+                    let res = call_scheduled_reducer(info, queue_item, call_reducer);
                     reply("call_scheduled_reducer", JsWorkerReply::CallScheduledReducer(res));
                 }
                 JsWorkerRequest::InitDatabase(program) => {

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

@@ -1,5 +1,6 @@
 use bytes::Bytes;
 use prometheus::{Histogram, IntCounter, IntGauge};
+use spacetimedb_auth::identity::ConnectionAuthCtx;
 use spacetimedb_lib::db::raw_def::v9::Lifecycle;
 use spacetimedb_lib::de::DeserializeSeed;
 use spacetimedb_primitives::ProcedureId;
@@ -20,8 +21,8 @@ use crate::host::module_host::{
 };
 use crate::host::scheduler::QueueItem;
 use crate::host::{
-    call_scheduled_reducer, init_database, ArgsTuple, ProcedureCallError, ProcedureCallResult, ReducerCallResult,
-    ReducerId, ReducerOutcome, Scheduler,
+    call_identity_connected, call_scheduled_reducer, init_database, ArgsTuple, ProcedureCallError, ProcedureCallResult,
+    ReducerCallResult, ReducerId, ReducerOutcome, Scheduler,
 };
 use crate::identity::Identity;
 use crate::messages::control_db::HostType;
@@ -281,6 +282,19 @@ impl<T: WasmInstance> WasmModuleInstance<T> {
         self.common.clear_all_clients()
     }
 
+    pub fn call_identity_connected(
+        &mut self,
+        caller_auth: ConnectionAuthCtx,
+        caller_connection_id: ConnectionId,
+    ) -> Result<(), ClientConnectedError> {
+        let module = &self.common.info.clone();
+        let call_reducer = |tx, params| self.call_reducer_inner(tx, params);
+        let mut trapped = false;
+        let res = call_identity_connected(caller_auth, caller_connection_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);