add ProcedureContext::with_transaction

Author: Centril

crates/bindings-sys/src/lib.rs

@@ -674,14 +674,15 @@ pub mod raw {
         ///
         /// # Traps
         ///
-        /// This function does not trap.
+        /// Traps if:
+        /// - `out` is NULL or `out[..size_of::<i64>()]` is not in bounds of WASM memory.
         ///
         /// # Errors
         ///
         /// Returns an error:
         ///
         /// - `WOULD_BLOCK_TRANSACTION`, if there's already an ongoing transaction.
-        pub fn procedure_start_mut_transaction() -> u16;
+        pub fn procedure_start_mut_transaction(out: *mut i64) -> u16;
 
         /// Commits a mutable transaction,
         /// suspending execution of this WASM instance until
@@ -1322,7 +1323,7 @@ impl Drop for RowIter {
 pub mod procedure {
     //! Side-effecting or asynchronous operations which only procedures are allowed to perform.
 
-    use super::{call_no_ret, raw, Result};
+    use super::{call, call_no_ret, raw, Result};
 
     #[inline]
     pub fn sleep_until(wake_at_timestamp: i64) -> i64 {
@@ -1335,7 +1336,7 @@ pub mod procedure {
     /// suspending execution of this WASM instance until
     /// a mutable transaction lock is aquired.
     ///
-    /// Upon resuming, returns `Ok(())` on success,
+    /// Upon resuming, returns `Ok(timestamp)` on success,
     /// enabling further calls that require a pending transaction,
     /// or [`Errno`] otherwise.
     ///
@@ -1345,8 +1346,8 @@ pub mod procedure {
     ///
     /// - `WOULD_BLOCK_TRANSACTION`, if there's already an ongoing transaction.
     #[inline]
-    pub fn procedure_start_mut_transaction() -> Result<()> {
-        call_no_ret(|| unsafe { raw::procedure_start_mut_transaction() })
+    pub fn procedure_start_mut_transaction() -> Result<i64> {
+        unsafe { call(|out| raw::procedure_start_mut_transaction(out)) }
     }
 
     /// Commits a mutable transaction,

crates/bindings/src/lib.rs

@@ -4,6 +4,7 @@
 use core::cell::{LazyCell, OnceCell, RefCell};
 use core::ops::Deref;
 use spacetimedb_lib::bsatn;
+use std::rc::Rc;
 
 #[cfg(feature = "unstable")]
 mod client_visibility_filter;
@@ -888,8 +889,6 @@ pub struct ViewContext {
 /// number generation.
 ///
 /// Implements the `DbContext` trait for accessing views into a database.
-/// Currently, being this generic is only meaningful in clients,
-/// as `ReducerContext` is the only implementor of `DbContext` within modules.
 #[non_exhaustive]
 pub struct ReducerContext {
     /// The `Identity` of the client that invoked the reducer.
@@ -904,6 +903,8 @@ pub struct ReducerContext {
     /// including `init` and scheduled reducers.
     pub connection_id: Option<ConnectionId>,
 
+    sender_auth: AuthCtx,
+
     /// Allows accessing the local database attached to a module.
     ///
     /// This slightly strange type appears to have no methods, but that is misleading.
@@ -942,8 +943,6 @@ pub struct ReducerContext {
     /// See the [`#[table]`](macro@crate::table) macro for more information.
     pub db: Local,
 
-    sender_auth: AuthCtx,
-
     #[cfg(feature = "rand08")]
     rng: std::cell::OnceCell<StdbRng>,
 }
@@ -964,16 +963,12 @@ impl ReducerContext {
 
     #[doc(hidden)]
     fn new(db: Local, sender: Identity, connection_id: Option<ConnectionId>, timestamp: Timestamp) -> Self {
-        let sender_auth = match connection_id {
-            Some(cid) => AuthCtx::from_connection_id(cid),
-            None => AuthCtx::internal(),
-        };
         Self {
             db,
             sender,
             timestamp,
             connection_id,
-            sender_auth,
+            sender_auth: AuthCtx::from_connection_id_opt(connection_id),
             #[cfg(feature = "rand08")]
             rng: std::cell::OnceCell::new(),
         }
@@ -1011,6 +1006,57 @@ impl ReducerContext {
     }
 }
 
+/// The context that an anonymous transaction
+/// in [`ProcedureContext::with_transaction`] is provided with.
+///
+/// Includes information about the client starting the transaction
+/// and the time of the procedure/reducer,
+/// as well as a view into the module's database.
+///
+/// If the crate was compiled with the `rand` feature, also includes faculties for random
+/// number generation.
+///
+/// Implements the `DbContext` trait for accessing views into a database.
+pub struct TxContext(ReducerContext);
+
+impl AsRef<ReducerContext> for TxContext {
+    fn as_ref(&self) -> &ReducerContext {
+        &self.0
+    }
+}
+
+impl Deref for TxContext {
+    type Target = ReducerContext;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+/// Values which knows whether they signify an ok state as opposed to error.
+pub trait IsOk {
+    /// Returns whether the current state of `self` is "ok".
+    fn is_ok(&self) -> bool;
+}
+
+impl IsOk for () {
+    fn is_ok(&self) -> bool {
+        true
+    }
+}
+
+impl<T> IsOk for Option<T> {
+    fn is_ok(&self) -> bool {
+        self.is_some()
+    }
+}
+
+impl<T, E> IsOk for Result<T, E> {
+    fn is_ok(&self) -> bool {
+        self.is_ok()
+    }
+}
+
 /// The context that any procedure is provided with.
 ///
 /// Each procedure must accept `&mut ProcedureContext` as its first argument.
@@ -1074,6 +1120,65 @@ impl ProcedureContext {
         let new_time = Timestamp::from_micros_since_unix_epoch(new_time);
         self.timestamp = new_time;
     }
+
+    /// 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()`,
+    /// the transaction will be rolled back and its mutations discarded.
+    ///
+    /// 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.
+    /// This is done once. On the second failure, a panic will occur.
+    pub fn with_transaction<R: IsOk>(&mut self, body: impl Fn(&TxContext) -> R) -> R {
+        let run = || {
+            // Start the transaction.
+            let timestamp = sys::procedure::procedure_start_mut_transaction().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.
+            let tx = ReducerContext::new(Local {}, self.sender, self.connection_id, timestamp);
+            let tx = TxContext(tx);
+            body(&tx)
+        };
+
+        let mut res = run();
+        let abort = || {
+            sys::procedure::procedure_abort_mut_transaction()
+                .expect("should have a pending mutable anon tx as `procedure_start_mut_transaction` preceded")
+        };
+
+        // Commit or roll back?
+        if res.is_ok() {
+            if sys::procedure::procedure_commit_mut_transaction().is_err() {
+                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_transaction().expect("transaction retry failed again")
+                } else {
+                    abort();
+                }
+            }
+        } else {
+            abort();
+        }
+
+        res
+    }
 }
 
 /// A handle on a database with a particular table schema.
@@ -1102,9 +1207,16 @@ impl DbContext for ReducerContext {
     }
 }
 
-// `ProcedureContext` is *not* a `DbContext`. We will add a `TxContext`
-// which can be obtained from `ProcedureContext::start_tx`,
-// and that will be a `DbContext`.
+impl DbContext for TxContext {
+    type DbView = Local;
+
+    fn db(&self) -> &Self::DbView {
+        &self.db
+    }
+}
+
+// `ProcedureContext` is *not* a `DbContext`
+// but a `TxContext` derived from it is.
 
 /// Allows accessing the local database attached to the module.
 ///
@@ -1125,18 +1237,25 @@ pub struct JwtClaims {
 }
 
 /// Authentication information for the caller of a reducer.
+#[derive(Clone)]
 pub struct AuthCtx {
     is_internal: bool,
     // NOTE(jsdt): cannot directly use a `LazyCell` without making this struct generic,
     // which would cause `ReducerContext` to become generic as well.
-    jwt: Box<dyn Deref<Target = Option<JwtClaims>>>,
+    jwt: Rc<dyn Deref<Target = Option<JwtClaims>>>,
 }
 
 impl AuthCtx {
+    /// Creates an [`AuthCtx`] both for cases where there's a [`ConnectionId`]
+    /// and for when there isn't.
+    fn from_connection_id_opt(conn_id: Option<ConnectionId>) -> Self {
+        conn_id.map(Self::from_connection_id).unwrap_or_else(Self::internal)
+    }
+
     fn new(is_internal: bool, jwt_fn: impl FnOnce() -> Option<JwtClaims> + 'static) -> Self {
         AuthCtx {
             is_internal,
-            jwt: Box::new(LazyCell::new(jwt_fn)),
+            jwt: Rc::new(LazyCell::new(jwt_fn)),
         }
     }
 

crates/bindings/src/rng.rs

@@ -1,13 +1,10 @@
-use std::cell::UnsafeCell;
-use std::marker::PhantomData;
-
-use crate::rand;
-
+use crate::{rand, ReducerContext};
+use core::cell::UnsafeCell;
+use core::marker::PhantomData;
 use rand::distributions::{Distribution, Standard};
 use rand::rngs::StdRng;
 use rand::{RngCore, SeedableRng};
-
-use crate::ReducerContext;
+use spacetimedb_lib::Timestamp;
 
 impl ReducerContext {
     /// Generates a random value.
@@ -48,10 +45,7 @@ impl ReducerContext {
     ///
     /// For more information, see [`StdbRng`] and [`rand::Rng`].
     pub fn rng(&self) -> &StdbRng {
-        self.rng.get_or_init(|| StdbRng {
-            rng: StdRng::seed_from_u64(self.timestamp.to_micros_since_unix_epoch() as u64).into(),
-            _marker: PhantomData,
-        })
+        self.rng.get_or_init(|| StdbRng::seed_from_ts(self.timestamp))
     }
 }
 
@@ -82,6 +76,16 @@ pub struct StdbRng {
     _marker: PhantomData<*mut ()>,
 }
 
+impl StdbRng {
+    /// Seeds a [`StdbRng`] from a timestamp.
+    fn seed_from_ts(timestamp: Timestamp) -> Self {
+        Self {
+            rng: StdRng::seed_from_u64(timestamp.to_micros_since_unix_epoch() as u64).into(),
+            _marker: PhantomData,
+        }
+    }
+}
+
 impl RngCore for StdbRng {
     fn next_u32(&mut self) -> u32 {
         (&*self).next_u32()

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

@@ -1454,7 +1454,8 @@ impl WasmInstanceEnv {
     ///
     /// # Traps
     ///
-    /// This function does not trap.
+    /// Traps if:
+    /// - `out` is NULL or `out[..size_of::<i64>()]` is not in bounds of WASM memory.
     ///
     /// # Errors
     ///
@@ -1463,21 +1464,27 @@ impl WasmInstanceEnv {
     /// - `WOULD_BLOCK_TRANSACTION`, if there's already an ongoing transaction.
     pub fn procedure_start_mut_transaction<'caller>(
         caller: Caller<'caller, Self>,
-        (): (),
+        (out,): (WasmPtr<u64>,),
     ) -> Fut<'caller, RtResult<u32>> {
-        Self::async_with_span(caller, AbiCall::ProcedureStartMutTransaction, |mut caller| async {
-            let (_, env) = Self::mem_env(&mut caller);
-            let res = env.instance_env.start_mutable_tx().await;
-
-            let result = res
-                .map(|()| {
-                    env.in_anon_tx = true;
-                    0u16.into()
-                })
-                .or_else(|err| Self::convert_wasm_result(AbiCall::ProcedureStartMutTransaction, err.into()));
-
-            (caller, result)
-        })
+        Self::async_with_span(
+            caller,
+            AbiCall::ProcedureStartMutTransaction,
+            move |mut caller| async move {
+                let (mem, env) = Self::mem_env(&mut caller);
+                let res = env.instance_env.start_mutable_tx().await.map_err(Into::into);
+                let timestamp = Timestamp::now().to_micros_since_unix_epoch() as u64;
+                let res = res.and_then(|()| Ok(timestamp.write_to(mem, out)?));
+
+                let result = res
+                    .map(|()| {
+                        env.in_anon_tx = true;
+                        0u16.into()
+                    })
+                    .or_else(|err| Self::convert_wasm_result(AbiCall::ProcedureStartMutTransaction, err));
+
+                (caller, result)
+            },
+        )
     }
 
     /// Commits a mutable transaction,