[TS] Anonymous transactions

Author: coolreader18

crates/bindings-typescript/src/lib/procedures.ts

@@ -3,7 +3,7 @@ import type { ConnectionId } from '../lib/connection_id';
 import type { Identity } from '../lib/identity';
 import type { Timestamp } from '../lib/timestamp';
 import type { HttpClient } from '../server/http_internal';
-import type { ParamsObj } from './reducers';
+import type { ParamsObj, ReducerCtx } from './reducers';
 import {
   MODULE_DEF,
   registerTypesRecursively,
@@ -25,8 +25,12 @@ export interface ProcedureCtx<S extends UntypedSchemaDef> {
   readonly timestamp: Timestamp;
   readonly connectionId: ConnectionId | null;
   readonly http: HttpClient;
+  withTx<T>(body: (ctx: TransactionCtx<S>) => T): T;
 }
 
+export interface TransactionCtx<S extends UntypedSchemaDef>
+  extends ReducerCtx<S> {}
+
 export function procedure<
   S extends UntypedSchemaDef,
   Params extends ParamsObj,

crates/bindings-typescript/src/server/index.ts

@@ -5,6 +5,7 @@ export { reducers } from '../lib/reducers';
 export { SenderError, SpacetimeHostError, errors } from './errors';
 export { type Reducer, type ReducerCtx } from '../lib/reducers';
 export { type DbView } from './db_view';
+export type { ProcedureCtx, TransactionCtx } from '../lib/procedures';
 
 import './polyfills'; // Ensure polyfills are loaded
 import './register_hooks'; // Ensure module hooks are registered

crates/bindings-typescript/src/server/procedures.ts

@@ -3,11 +3,16 @@ import BinaryReader from '../lib/binary_reader';
 import BinaryWriter from '../lib/binary_writer';
 import type { ConnectionId } from '../lib/connection_id';
 import { Identity } from '../lib/identity';
-import { PROCEDURES, type ProcedureCtx } from '../lib/procedures';
+import {
+  PROCEDURES,
+  type ProcedureCtx,
+  type TransactionCtx,
+} from '../lib/procedures';
 import { MODULE_DEF, type UntypedSchemaDef } from '../lib/schema';
-import type { Timestamp } from '../lib/timestamp';
+import { Timestamp } from '../lib/timestamp';
+import type { DbView } from './db_view';
 import { httpClient } from './http_internal';
-import { sys } from './runtime';
+import { makeReducerCtx, sys } from './runtime';
 
 const { freeze } = Object;
 
@@ -25,15 +30,45 @@ export function callProcedure(
     MODULE_DEF.typespace
   );
 
-  const ctx: ProcedureCtx<UntypedSchemaDef> = freeze({
+  const ctx: ProcedureCtx<UntypedSchemaDef> = {
     sender,
     timestamp,
     connectionId,
     http: httpClient,
     get identity() {
       return new Identity(sys.identity().__identity__);
     },
-  });
+    withTx(body) {
+      const run = () => {
+        const timestamp = sys.procedure_start_mut_tx();
+
+        try {
+          const ctx: TransactionCtx<UntypedSchemaDef> = freeze(
+            makeReducerCtx(sender, new Timestamp(timestamp), connectionId)
+          );
+          return body(ctx);
+        } catch (e) {
+          sys.procedure_abort_mut_tx();
+          throw e;
+        }
+      };
+
+      let res = run();
+      try {
+        sys.procedure_commit_mut_tx();
+        return res;
+      } catch {}
+      console.warn('committing anonymous transaction failed');
+      res = run();
+      try {
+        sys.procedure_commit_mut_tx();
+        return res;
+      } catch (e) {
+        throw new Error('transaction retry failed again', { cause: e });
+      }
+    },
+  };
+  freeze(ctx);
 
   const ret = fn(ctx, args);
   const retBuf = new BinaryWriter(returnTypeBaseSize);

crates/bindings-typescript/src/server/runtime.ts

@@ -25,7 +25,7 @@ import {
   type JwtClaims,
   type ReducerCtx,
 } from '../lib/reducers';
-import { MODULE_DEF } from '../lib/schema';
+import { MODULE_DEF, type UntypedSchemaDef } from '../lib/schema';
 import { type RowType, type Table, type TableMethods } from '../lib/table';
 import { Timestamp } from '../lib/timestamp';
 import type { Infer } from '../lib/type_builders';
@@ -175,6 +175,21 @@ class AuthCtxImpl implements AuthCtx {
   }
 }
 
+export const makeReducerCtx = (
+  sender: Identity,
+  timestamp: Timestamp,
+  connectionId: ConnectionId | null
+): ReducerCtx<UntypedSchemaDef> => ({
+  sender,
+  get identity() {
+    return new Identity(sys.identity().__identity__);
+  },
+  timestamp,
+  connectionId,
+  db: getDbView(),
+  senderAuth: AuthCtxImpl.fromSystemTables(connectionId, sender),
+});
+
 export const hooks: ModuleHooks = {
   __describe_module__() {
     const writer = new BinaryWriter(128);
@@ -195,19 +210,13 @@ export const hooks: ModuleHooks = {
       MODULE_DEF.typespace
     );
     const senderIdentity = new Identity(sender);
-    const ctx: ReducerCtx<any> = freeze({
-      sender: senderIdentity,
-      get identity() {
-        return new Identity(sys.identity().__identity__);
-      },
-      timestamp: new Timestamp(timestamp),
-      connectionId: ConnectionId.nullIfZero(new ConnectionId(connId)),
-      db: getDbView(),
-      senderAuth: AuthCtxImpl.fromSystemTables(
-        ConnectionId.nullIfZero(new ConnectionId(connId)),
-        senderIdentity
-      ),
-    });
+    const ctx: ReducerCtx<any> = freeze(
+      makeReducerCtx(
+        senderIdentity,
+        new Timestamp(timestamp),
+        ConnectionId.nullIfZero(new ConnectionId(connId))
+      )
+    );
     try {
       return REDUCERS[reducerId](ctx, args) ?? { tag: 'ok' };
     } catch (e) {

crates/bindings-typescript/src/server/sys.d.ts

@@ -93,4 +93,10 @@ declare module 'spacetime:sys@1.2' {
     request: Uint8Array,
     body: Uint8Array | string
   ): [response: Uint8Array, body: Uint8Array];
+
+  export function procedure_start_mut_tx(): bigint;
+
+  export function procedure_commit_mut_tx();
+
+  export function procedure_abort_mut_tx();
 }

crates/core/src/host/instance_env.rs

@@ -601,7 +601,11 @@ impl InstanceEnv {
         written
     }
 
-    pub async fn start_mutable_tx(&mut self) -> Result<(), NodesError> {
+    // Async procedure syscalls return a `Result<impl Future>`, so that we can check `get_tx()`
+    // *before* requiring an async runtime. Otherwise, the v8 module host would have to call
+    // on `tokio::runtime::Handle::try_current()` before being able to run the `get_tx()` check.
+
+    pub fn start_mutable_tx(&mut self) -> Result<impl Future<Output = ()> + use<'_>, NodesError> {
         if self.get_tx().is_ok() {
             return Err(NodesError::WouldBlockTransaction(
                 super::AbiCall::ProcedureStartMutTransaction,
@@ -610,11 +614,13 @@ impl InstanceEnv {
 
         let stdb = self.replica_ctx.relational_db.clone();
         // TODO(procedure-tx): should we add a new workload, e.g., `AnonTx`?
-        let tx = asyncify(move || stdb.begin_mut_tx(IsolationLevel::Serializable, Workload::Internal)).await;
-        self.tx.set_raw(tx);
-        self.in_anon_tx = true;
+        let fut = async move {
+            let tx = asyncify(move || stdb.begin_mut_tx(IsolationLevel::Serializable, Workload::Internal)).await;
+            self.tx.set_raw(tx);
+            self.in_anon_tx = true;
+        };
 
-        Ok(())
+        Ok(fut)
     }
 
     /// Finishes an anonymous transaction,
@@ -637,7 +643,7 @@ impl InstanceEnv {
     // *before* requiring an async runtime. Otherwise, the v8 module host would have to call
     // on `tokio::runtime::Handle::try_current()` before being able to run the `get_tx()` check.
 
-    pub async fn commit_mutable_tx(&mut self) -> Result<(), NodesError> {
+    pub fn commit_mutable_tx(&mut self) -> Result<impl Future<Output = ()> + use<'_>, NodesError> {
         self.finish_anon_tx()?;
 
         let stdb = self.relational_db().clone();
@@ -660,10 +666,12 @@ impl InstanceEnv {
         // This is somewhat expensive,
         // and can block for a while,
         // so we need to asyncify it.
-        let event = asyncify(move || commit_and_broadcast_event(&subs, None, event, tx)).await;
-        self.procedure_last_tx_offset = Some(event.tx_offset);
+        let fut = async move {
+            let event = asyncify(move || commit_and_broadcast_event(&subs, None, event, tx)).await;
+            self.procedure_last_tx_offset = Some(event.tx_offset);
+        };
 
-        Ok(())
+        Ok(fut)
     }
 
     pub fn abort_mutable_tx(&mut self) -> Result<(), NodesError> {

crates/core/src/host/v8/syscall/v1.rs

@@ -20,7 +20,7 @@ use crate::host::wasm_common::{err_to_errno_and_log, RowIterIdx, TimingSpan, Tim
 use crate::host::AbiCall;
 use anyhow::Context;
 use bytes::Bytes;
-use spacetimedb_lib::{bsatn, ConnectionId, Identity, RawModuleDef};
+use spacetimedb_lib::{bsatn, ConnectionId, Identity, RawModuleDef, Timestamp};
 use spacetimedb_primitives::{errno, ColId, IndexId, ProcedureId, ReducerId, TableId, ViewFnPtr};
 use spacetimedb_sats::Serialize;
 use v8::{
@@ -129,7 +129,22 @@ pub(super) fn sys_v1_2<'scope>(scope: &mut PinScope<'scope, '_>) -> Local<'scope
             with_sys_result_value,
             AbiCall::ProcedureHttpRequest,
             procedure_http_request
-        )
+        ),
+        (
+            with_sys_result_value,
+            AbiCall::ProcedureStartMutTransaction,
+            procedure_start_mut_tx
+        ),
+        (
+            with_sys_result_noret,
+            AbiCall::ProcedureAbortMutTransaction,
+            procedure_abort_mut_tx
+        ),
+        (
+            with_sys_result_noret,
+            AbiCall::ProcedureCommitMutTransaction,
+            procedure_commit_mut_tx
+        ),
     )
 }
 
@@ -1585,3 +1600,37 @@ fn procedure_http_request<'scope>(
         &[response.into(), response_body.into()],
     ))
 }
+
+fn procedure_start_mut_tx<'scope>(
+    scope: &mut PinScope<'scope, '_>,
+    _args: FunctionCallbackArguments<'_>,
+) -> SysCallResult<Local<'scope, v8::BigInt>> {
+    let env = get_env(scope)?;
+
+    let fut = env.instance_env.start_mutable_tx()?;
+
+    let rt = tokio::runtime::Handle::current();
+    rt.block_on(fut);
+
+    let timestamp = Timestamp::now().to_micros_since_unix_epoch() as u64;
+
+    Ok(v8::BigInt::new_from_u64(scope, timestamp))
+}
+
+fn procedure_abort_mut_tx(scope: &mut PinScope<'_, '_>, _args: FunctionCallbackArguments<'_>) -> SysCallResult<()> {
+    let env = get_env(scope)?;
+
+    env.instance_env.abort_mutable_tx()?;
+    Ok(())
+}
+
+fn procedure_commit_mut_tx(scope: &mut PinScope<'_, '_>, _args: FunctionCallbackArguments<'_>) -> SysCallResult<()> {
+    let env = get_env(scope)?;
+
+    let fut = env.instance_env.commit_mutable_tx()?;
+
+    let rt = tokio::runtime::Handle::current();
+    rt.block_on(fut);
+
+    Ok(())
+}

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

@@ -1465,7 +1465,7 @@ impl WasmInstanceEnv {
             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 res = async { Ok(env.instance_env.start_mutable_tx()?.await) }.await;
                 let timestamp = Timestamp::now().to_micros_since_unix_epoch() as u64;
                 let res = res.and_then(|()| Ok(timestamp.write_to(mem, out)?));
 
@@ -1508,12 +1508,12 @@ impl WasmInstanceEnv {
             |mut caller| async move {
                 let (_, env) = Self::mem_env(&mut caller);
 
-                let res = env
-                    .instance_env
-                    .commit_mutable_tx()
-                    .await
-                    .map(|()| 0u16.into())
-                    .or_else(|err| Self::convert_wasm_result(AbiCall::ProcedureCommitMutTransaction, err.into()));
+                let res = async {
+                    env.instance_env.commit_mutable_tx()?.await;
+                    Ok(0u16.into())
+                }
+                .await
+                .or_else(|err| Self::convert_wasm_result(AbiCall::ProcedureCommitMutTransaction, err));
 
                 (caller, res)
             },

modules/sdk-test-procedure-ts/src/index.ts

@@ -1,7 +1,14 @@
 // ─────────────────────────────────────────────────────────────────────────────
 // IMPORTS
 // ─────────────────────────────────────────────────────────────────────────────
-import { errors, schema, t, table } from 'spacetimedb/server';
+import {
+  errors,
+  schema,
+  t,
+  table,
+  type ProcedureCtx,
+  type TransactionCtx,
+} from 'spacetimedb/server';
 
 const ReturnStruct = t.object('ReturnStruct', {
   a: t.u32(),
@@ -13,7 +20,12 @@ const ReturnEnum = t.enum('ReturnEnum', {
   B: t.string(),
 });
 
-const spacetimedb = schema();
+const MyTable = table(
+  { name: 'my_table', public: true },
+  { field: ReturnStruct }
+);
+
+const spacetimedb = schema(MyTable);
 
 spacetimedb.procedure(
   'return_primitive',
@@ -68,3 +80,42 @@ spacetimedb.procedure('invalid_request', t.string(), ctx => {
     throw e;
   }
 });
+
+function insertMyTable(ctx: TransactionCtx<typeof spacetimedb.schemaType>) {
+  ctx.db.myTable.insert({ field: { a: 42, b: 'magic' } });
+}
+
+function assertRowCount(
+  ctx: ProcedureCtx<typeof spacetimedb.schemaType>,
+  count: number
+) {
+  ctx.withTx(ctx => {
+    assertEqual(ctx.db.myTable.count(), BigInt(count));
+  });
+}
+
+function assertEqual<T>(a: T, b: T) {
+  if (a !== b) {
+    throw new Error(`assertion failed: ${a} != ${b}`);
+  }
+}
+
+spacetimedb.procedure('insert_with_tx_commit', t.unit(), ctx => {
+  ctx.withTx(insertMyTable);
+  assertRowCount(ctx, 1);
+  return {};
+});
+
+spacetimedb.procedure('insert_with_tx_rollback', t.unit(), ctx => {
+  const error = {};
+  try {
+    ctx.withTx(ctx => {
+      insertMyTable(ctx);
+      throw error;
+    });
+  } catch (e) {
+    if (e !== error) throw e;
+  }
+  assertRowCount(ctx, 0);
+  return {};
+});