split with_tx into self and try_with_tx

Author: Centril

crates/bindings/src/lib.rs

@@ -1121,12 +1121,44 @@ impl ProcedureContext {
         self.timestamp = new_time;
     }
 
+    /// Acquire a mutable transaction
+    /// and execute `body` with read-write access to the database.
+    ///
+    /// When a panic occurs,
+    /// the transaction will be rolled back and its mutations discarded.
+    /// Otherwise, the transaction will be committed and its mutations persisted.
+    ///
+    /// Regardless of the transaction's success or failure,
+    /// the return value of `body` is not persisted to the commitlog
+    /// or broadcast to subscribed clients.
+    /// Clients attribute mutations performed by this transaction to `Event::UnknownTransaction`.
+    ///
+    /// If the transaction fails to commit after `body` returns,
+    /// e.g., due to a conflict with a concurrent transaction,
+    /// this method will re-invoke `body` with a new transaction in order to retry.
+    /// The transaction will be retried at most once.
+    /// If it fails to commit a second time, this method will panic.
+    ///
+    /// Because `body` may be run multiple times,
+    /// and is expected to perform the same set of database operations
+    /// and return the same result on each invocation,
+    /// callers should avoid writing to any captured mutable state within `body`,
+    /// This includes interior mutability through types like [`std::cell::Cell`].
+    #[cfg(feature = "unstable")]
+    pub fn with_tx<T>(&mut self, body: impl Fn(&TxContext) -> T) -> T {
+        use core::convert::Infallible;
+        match self.try_with_tx::<T, Infallible>(|tx| Ok(body(tx))) {
+            Ok(v) => v,
+            Err(e) => match e {},
+        }
+    }
+
     /// Acquire a mutable transaction
     /// and execute `body` with read-write access to the database.
     ///
     /// When `body().is_ok()`,
     /// the transaction will be committed and its mutations persisted.
-    /// When `!body().is_ok()`,
+    /// When `!body().is_ok()` or a panic occurs in `body`,
     /// the transaction will be rolled back and its mutations discarded.
     ///
     /// Regardless of the transaction's success or failure,
@@ -1146,43 +1178,57 @@ impl ProcedureContext {
     /// callers should avoid writing to any captured mutable state within `body`,
     /// This includes interior mutability through types like [`std::cell::Cell`].
     #[cfg(feature = "unstable")]
-    pub fn with_tx<R: IsOk>(&mut self, body: impl Fn(&TxContext) -> R) -> R {
+    pub fn try_with_tx<T, E>(&mut self, body: impl Fn(&TxContext) -> Result<T, E>) -> Result<T, E> {
+        let abort = || {
+            sys::procedure::procedure_abort_mut_tx()
+                .expect("should have a pending mutable anon tx as `procedure_start_mut_tx` preceded")
+        };
+
         let run = || {
             // Start the transaction.
             let timestamp = sys::procedure::procedure_start_mut_tx().expect(
                 "holding `&mut ProcedureContext`, so should not be in a tx already; called manually elsewhere?",
             );
             let timestamp = Timestamp::from_micros_since_unix_epoch(timestamp);
 
-            // We've resumed, so do the work.
+            // We've resumed, so let's do the work, but first prepare the context.
             let tx = ReducerContext::new(Local {}, self.sender, self.connection_id, timestamp);
             let tx = TxContext(tx);
-            body(&tx)
+
+            // Guard the execution of `body` with a scope-guard that `abort`s on panic.
+            // Wasmtime now supports unwinding, so we need to protect against that.
+            // We're not using `scopeguard::guard` here to avoid an extra dependency.
+            struct DoOnDrop<F: Fn()>(F);
+            impl<F: Fn()> Drop for DoOnDrop<F> {
+                fn drop(&mut self) {
+                    (self.0)();
+                }
+            }
+            let abort_guard = DoOnDrop(abort);
+            let res = body(&tx);
+            // Defuse the bomb.
+            let DoOnDrop(_) = abort_guard;
+            res
         };
 
         let mut res = run();
-        let abort = || {
-            sys::procedure::procedure_abort_mut_tx()
-                .expect("should have a pending mutable anon tx as `procedure_start_mut_tx` preceded")
-        };
 
         // Commit or roll back?
-        if res.is_ok() {
-            if sys::procedure::procedure_commit_mut_tx().is_err() {
+        match res {
+            Ok(_) if sys::procedure::procedure_commit_mut_tx().is_err() => {
+                // Tried to commit, but couldn't. Retry once.
                 log::warn!("committing anonymous transaction failed");
-
                 // NOTE(procedure,centril): there's no actual guarantee that `body`
                 // does the exact same as the time before, as the timestamps differ
                 // and due to interior mutability.
                 res = run();
-                if res.is_ok() {
-                    sys::procedure::procedure_commit_mut_tx().expect("transaction retry failed again")
-                } else {
-                    abort();
+                match res {
+                    Ok(_) => sys::procedure::procedure_commit_mut_tx().expect("transaction retry failed again"),
+                    Err(_) => abort(),
                 }
             }
-        } else {
-            abort();
+            Ok(_) => {}
+            Err(_) => abort(),
         }
 
         res

crates/core/src/host/wasmtime/wasm_instance_env.rs

@@ -5,7 +5,6 @@ use crate::database_logger::{BacktraceFrame, BacktraceProvider, ModuleBacktrace,
 use crate::error::NodesError;
 use crate::host::instance_env::{ChunkPool, InstanceEnv};
 use crate::host::module_host::{DatabaseUpdate, EventStatus, ModuleEvent, ModuleFunctionCall};
-use crate::host::wasm_common::instrumentation::noop::CallSpanStart;
 use crate::host::wasm_common::instrumentation::{span, CallTimes};
 use crate::host::wasm_common::module_host_actor::ExecutionTimings;
 use crate::host::wasm_common::{err_to_errno_and_log, RowIterIdx, RowIters, TimingSpan, TimingSpanIdx, TimingSpanSet};
@@ -295,7 +294,7 @@ impl WasmInstanceEnv {
     }
 
     /// Record a span with `start`.
-    fn end_span(mut caller: Caller<'_, Self>, start: CallSpanStart) {
+    fn end_span(mut caller: Caller<'_, Self>, start: span::CallSpanStart) {
         let span = start.end();
         span::record_span(&mut caller.data_mut().call_times, span);
     }