send tx offset for procedures & abort dangling anon txes

Author: Centril

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

@@ -281,6 +281,7 @@ async fn procedure<S: ControlStateDelegate + NodeDelegate>(
     let result = match module
         .call_procedure(caller_identity, Some(connection_id), None, &procedure, args)
         .await
+        .result
     {
         Ok(res) => Ok(res),
         Err(e) => {

crates/core/src/client/client_connection.rs

@@ -12,7 +12,7 @@ use super::{message_handlers, ClientActorId, MessageHandleError};
 use crate::db::relational_db::RelationalDB;
 use crate::error::DBError;
 use crate::host::module_host::ClientConnectedError;
-use crate::host::{FunctionArgs, ModuleHost, NoSuchModule, ReducerCallError, ReducerCallResult};
+use crate::host::{CallProcedureReturn, FunctionArgs, ModuleHost, NoSuchModule, ReducerCallError, ReducerCallResult};
 use crate::messages::websocket::Subscribe;
 use crate::subscription::module_subscription_manager::BroadcastError;
 use crate::util::asyncify;
@@ -848,7 +848,7 @@ impl ClientConnection {
         request_id: RequestId,
         timer: Instant,
     ) -> Result<(), BroadcastError> {
-        let res = self
+        let CallProcedureReturn { result, tx_offset } = self
             .module()
             .call_procedure(
                 self.id.identity,
@@ -859,9 +859,11 @@ impl ClientConnection {
             )
             .await;
 
+        let message = ProcedureResultMessage::from_result(&result, request_id);
+
         self.module()
             .subscriptions()
-            .send_procedure_message(self.sender(), ProcedureResultMessage::from_result(&res, request_id))
+            .send_procedure_message(self.sender(), message, tx_offset)
     }
 
     pub async fn subscribe_single(

crates/core/src/host/host_controller.rs

@@ -10,11 +10,12 @@ use crate::db::{self, spawn_tx_metrics_recorder};
 use crate::energy::{EnergyMonitor, EnergyQuanta, NullEnergyMonitor};
 use crate::host::module_host::ModuleRuntime as _;
 use crate::host::v8::V8Runtime;
+use crate::host::ProcedureCallError;
 use crate::messages::control_db::{Database, HostType};
 use crate::module_host_context::ModuleCreationContext;
 use crate::replica_context::ReplicaContext;
 use crate::subscription::module_subscription_actor::ModuleSubscriptions;
-use crate::subscription::module_subscription_manager::{spawn_send_worker, SubscriptionManager};
+use crate::subscription::module_subscription_manager::{spawn_send_worker, SubscriptionManager, TransactionOffset};
 use crate::util::asyncify;
 use crate::util::jobs::{JobCores, SingleCoreExecutor};
 use crate::worker_metrics::WORKER_METRICS;
@@ -178,6 +179,11 @@ pub struct ProcedureCallResult {
     pub start_timestamp: Timestamp,
 }
 
+pub struct CallProcedureReturn {
+    pub result: Result<ProcedureCallResult, ProcedureCallError>,
+    pub tx_offset: Option<TransactionOffset>,
+}
+
 impl HostController {
     pub fn new(
         data_dir: Arc<ServerDataDir>,

crates/core/src/host/instance_env.rs

@@ -230,11 +230,8 @@ impl InstanceEnv {
         );
     }
 
-    /// End a console timer by logging the span at INFO level.
-    pub(crate) fn console_timer_end(&self, span: &TimingSpan, function: Option<&str>) {
-        let elapsed = span.start.elapsed();
-        let message = format!("Timing span {:?}: {:?}", &span.name, elapsed);
-
+    /// Logs a simple `message` at `level`.
+    pub(crate) fn console_log_simple_message(&self, level: LogLevel, function: Option<&str>, message: &str) {
         /// A backtrace provider that provides nothing.
         struct Noop;
         impl BacktraceProvider for Noop {
@@ -254,9 +251,17 @@ impl InstanceEnv {
             filename: None,
             line_number: None,
             function,
-            message: &message,
+            message,
         };
-        self.console_log(LogLevel::Info, &record, &Noop);
+        self.console_log(level, &record, &Noop);
+    }
+
+    /// End a console timer by logging the span at INFO level.
+    pub(crate) fn console_timer_end(&self, span: &TimingSpan, function: Option<&str>) {
+        let elapsed = span.start.elapsed();
+        let message = format!("Timing span {:?}: {:?}", &span.name, elapsed);
+
+        self.console_log_simple_message(LogLevel::Info, function, &message);
     }
 
     /// Returns the current time suitable for logging.
@@ -622,7 +627,7 @@ impl TxSlot {
         MutexGuard::try_map(self.inner.lock(), |map| map.as_mut()).map_err(|_| GetTxError)
     }
 
-    /// Steals th tx from the slot.
+    /// Steals the tx from the slot.
     pub fn take(&self) -> Result<MutTxId, GetTxError> {
         self.inner.lock().take().ok_or(GetTxError)
     }

crates/core/src/host/mod.rs

@@ -24,8 +24,8 @@ mod wasm_common;
 
 pub use disk_storage::DiskStorage;
 pub use host_controller::{
-    extract_schema, ExternalDurability, ExternalStorage, HostController, MigratePlanResult, ProcedureCallResult,
-    ProgramStorage, ReducerCallResult, ReducerOutcome,
+    extract_schema, CallProcedureReturn, ExternalDurability, ExternalStorage, HostController, MigratePlanResult,
+    ProcedureCallResult, ProgramStorage, ReducerCallResult, ReducerOutcome,
 };
 pub use module_host::{ModuleHost, NoSuchModule, ProcedureCallError, ReducerCallError, UpdateDatabaseResult};
 pub use scheduler::Scheduler;

crates/core/src/host/module_host.rs

@@ -1,6 +1,6 @@
 use super::{
-    ArgsTuple, FunctionArgs, InvalidProcedureArguments, InvalidReducerArguments, ProcedureCallResult,
-    ReducerCallResult, ReducerId, ReducerOutcome, Scheduler,
+    ArgsTuple, FunctionArgs, InvalidProcedureArguments, InvalidReducerArguments, ReducerCallResult, ReducerId,
+    ReducerOutcome, Scheduler,
 };
 use crate::client::messages::{OneOffQueryResponseMessage, SerializableMessage};
 use crate::client::{ClientActorId, ClientConnectionSender};
@@ -10,6 +10,7 @@ use crate::energy::EnergyQuanta;
 use crate::error::DBError;
 use crate::estimation::estimate_rows_scanned;
 use crate::hash::Hash;
+use crate::host::host_controller::CallProcedureReturn;
 use crate::host::scheduler::{handle_queued_call_reducer_params, QueueItem};
 use crate::host::v8::JsInstance;
 use crate::host::wasmtime::ModuleInstance;
@@ -401,7 +402,7 @@ impl Instance {
         }
     }
 
-    async fn call_procedure(&mut self, params: CallProcedureParams) -> Result<ProcedureCallResult, ProcedureCallError> {
+    async fn call_procedure(&mut self, params: CallProcedureParams) -> CallProcedureReturn {
         match self {
             Instance::Wasm(inst) => inst.call_procedure(params).await,
             Instance::Js(inst) => inst.call_procedure(params).await,
@@ -1550,7 +1551,7 @@ impl ModuleHost {
         timer: Option<Instant>,
         procedure_name: &str,
         args: FunctionArgs,
-    ) -> Result<ProcedureCallResult, ProcedureCallError> {
+    ) -> CallProcedureReturn {
         let res = async {
             let (procedure_id, procedure_def) = self
                 .info
@@ -1569,7 +1570,15 @@ impl ModuleHost {
         }
         .await;
 
-        let log_message = match &res {
+        let ret = match res {
+            Ok(ret) => ret,
+            Err(err) => CallProcedureReturn {
+                result: Err(err),
+                tx_offset: None,
+            },
+        };
+
+        let log_message = match &ret.result {
             Err(ProcedureCallError::NoSuchProcedure) => Some(no_such_function_log_message("procedure", procedure_name)),
             Err(ProcedureCallError::Args(_)) => Some(args_error_log_message("procedure", procedure_name)),
             _ => None,
@@ -1579,7 +1588,7 @@ impl ModuleHost {
             self.inject_logs(LogLevel::Error, procedure_name, &log_message)
         }
 
-        res
+        ret
     }
 
     async fn call_procedure_inner(
@@ -1590,10 +1599,10 @@ impl ModuleHost {
         procedure_id: ProcedureId,
         procedure_def: &ProcedureDef,
         args: FunctionArgs,
-    ) -> Result<ProcedureCallResult, ProcedureCallError> {
+    ) -> Result<CallProcedureReturn, ProcedureCallError> {
         let procedure_seed = ArgsSeed(self.info.module_def.typespace().with_type(procedure_def));
-        let args = args.into_tuple(procedure_seed).map_err(InvalidProcedureArguments)?;
         let caller_connection_id = caller_connection_id.unwrap_or(ConnectionId::ZERO);
+        let args = args.into_tuple(procedure_seed).map_err(InvalidProcedureArguments)?;
 
         self.call_async_with_instance(&procedure_def.name, async move |mut inst| {
             let res = inst
@@ -1608,7 +1617,8 @@ impl ModuleHost {
                 .await;
             (res, inst)
         })
-        .await?
+        .await
+        .map_err(Into::into)
     }
 
     // Scheduled reducers require a different function here to call their reducer

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

@@ -13,6 +13,7 @@ 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::client::ClientActorId;
+use crate::host::host_controller::CallProcedureReturn;
 use crate::host::instance_env::{ChunkPool, InstanceEnv, TxSlot};
 use crate::host::module_host::{
     call_identity_connected, call_scheduled_reducer, init_database, CallViewParams, ClientConnectedError, Instance,
@@ -29,6 +30,7 @@ use crate::host::wasm_common::{RowIters, TimingSpanSet};
 use crate::host::{ModuleHost, ReducerCallError, ReducerCallResult, Scheduler};
 use crate::module_host_context::{ModuleCreationContext, ModuleCreationContextLimited};
 use crate::replica_context::ReplicaContext;
+use crate::subscription::module_subscription_manager::TransactionOffset;
 use crate::util::asyncify;
 use anyhow::Context as _;
 use core::any::type_name;
@@ -380,10 +382,7 @@ impl JsInstance {
         .await
     }
 
-    pub async fn call_procedure(
-        &mut self,
-        _params: CallProcedureParams,
-    ) -> Result<super::ProcedureCallResult, super::ProcedureCallError> {
+    pub async fn call_procedure(&mut self, _params: CallProcedureParams) -> CallProcedureReturn {
         todo!("JS/TS module procedure support")
     }
 
@@ -781,7 +780,11 @@ impl WasmInstance for V8Instance<'_, '_, '_> {
         log_traceback(self.replica_ctx, func_type, func, trap)
     }
 
-    async fn call_procedure(&mut self, _op: ProcedureOp, _budget: FunctionBudget) -> ProcedureExecuteResult {
+    async fn call_procedure(
+        &mut self,
+        _op: ProcedureOp,
+        _budget: FunctionBudget,
+    ) -> (ProcedureExecuteResult, Option<TransactionOffset>) {
         todo!("JS/TS module procedure support")
     }
 }

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

@@ -3,6 +3,7 @@ use super::*;
 use crate::client::ClientActorId;
 use crate::database_logger;
 use crate::energy::{EnergyMonitor, FunctionBudget, FunctionFingerprint};
+use crate::host::host_controller::CallProcedureReturn;
 use crate::host::instance_env::{InstanceEnv, TxSlot};
 use crate::host::module_common::{build_common_module_from_raw, ModuleCommon};
 use crate::host::module_host::{
@@ -20,6 +21,7 @@ use crate::messages::control_db::HostType;
 use crate::module_host_context::ModuleCreationContextLimited;
 use crate::replica_context::ReplicaContext;
 use crate::subscription::module_subscription_actor::commit_and_broadcast_event;
+use crate::subscription::module_subscription_manager::TransactionOffset;
 use crate::util::prometheus_handle::{HistogramExt, TimerGuard};
 use crate::worker_metrics::WORKER_METRICS;
 use bytes::Bytes;
@@ -79,7 +81,7 @@ pub trait WasmInstance {
         &mut self,
         op: ProcedureOp,
         budget: FunctionBudget,
-    ) -> impl Future<Output = ProcedureExecuteResult>;
+    ) -> impl Future<Output = (ProcedureExecuteResult, Option<TransactionOffset>)>;
 }
 
 pub struct EnergyStats {
@@ -378,15 +380,12 @@ impl<T: WasmInstance> WasmModuleInstance<T> {
         res
     }
 
-    pub async fn call_procedure(
-        &mut self,
-        params: CallProcedureParams,
-    ) -> Result<ProcedureCallResult, ProcedureCallError> {
-        let res = self.common.call_procedure(params, &mut self.instance).await;
-        if res.is_err() {
+    pub async fn call_procedure(&mut self, params: CallProcedureParams) -> CallProcedureReturn {
+        let ret = self.common.call_procedure(params, &mut self.instance).await;
+        if ret.result.is_err() {
             self.trapped = true;
         }
-        res
+        ret
     }
 }
 
@@ -575,7 +574,7 @@ impl InstanceCommon {
         &mut self,
         params: CallProcedureParams,
         inst: &mut I,
-    ) -> Result<ProcedureCallResult, ProcedureCallError> {
+    ) -> CallProcedureReturn {
         let CallProcedureParams {
             timestamp,
             caller_identity,
@@ -612,7 +611,7 @@ impl InstanceCommon {
         // TODO(procedure-energy): replace with call to separate function `procedure_budget`.
         let budget = self.energy_monitor.reducer_budget(&energy_fingerprint);
 
-        let result = inst.call_procedure(op, budget).await;
+        let (result, tx_offset) = inst.call_procedure(op, budget).await;
 
         let ProcedureExecuteResult {
             stats:
@@ -630,7 +629,7 @@ impl InstanceCommon {
             self.allocated_memory = memory_allocation;
         }
 
-        match call_result {
+        let result = match call_result {
             Err(err) => {
                 inst.log_traceback("procedure", &procedure_def.name, &err);
 
@@ -655,16 +654,17 @@ impl InstanceCommon {
             Ok(return_val) => {
                 let return_type = &procedure_def.return_type;
                 let seed = spacetimedb_sats::WithTypespace::new(self.info.module_def.typespace(), return_type);
-                let return_val = seed
-                    .deserialize(bsatn::Deserializer::new(&mut &return_val[..]))
-                    .map_err(|err| ProcedureCallError::InternalError(format!("{err}")))?;
-                Ok(ProcedureCallResult {
-                    return_val,
-                    execution_duration: timer.map(|timer| timer.elapsed()).unwrap_or_default(),
-                    start_timestamp: timestamp,
-                })
+                seed.deserialize(bsatn::Deserializer::new(&mut &return_val[..]))
+                    .map_err(|err| ProcedureCallError::InternalError(format!("{err}")))
+                    .map(|return_val| ProcedureCallResult {
+                        return_val,
+                        execution_duration: timer.map(|timer| timer.elapsed()).unwrap_or_default(),
+                        start_timestamp: timestamp,
+                    })
             }
-        }
+        };
+
+        CallProcedureReturn { result, tx_offset }
     }
 
     /// Execute a reducer.