WIP on wound-wait for 2pc

Author: jsdt

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

@@ -20,6 +20,7 @@ use axum::response::{ErrorResponse, IntoResponse};
 use axum::routing::MethodRouter;
 use axum::Extension;
 use axum_extra::TypedHeader;
+use http::HeaderMap;
 use futures::TryStreamExt;
 use http::StatusCode;
 use log::{info, warn};
@@ -41,7 +42,7 @@ use spacetimedb_lib::bsatn;
 use spacetimedb_lib::db::raw_def::v10::RawModuleDefV10;
 use spacetimedb_lib::db::raw_def::v9::RawModuleDefV9;
 use spacetimedb_lib::de::DeserializeSeed;
-use spacetimedb_lib::{sats, AlgebraicValue, Hash, ProductValue, Timestamp};
+use spacetimedb_lib::{sats, AlgebraicValue, GlobalTxId, Hash, ProductValue, Timestamp, TX_ID_HEADER};
 use spacetimedb_schema::auto_migrate::{
     MigrationPolicy as SchemaMigrationPolicy, MigrationToken, PrettyPrintStyle as AutoMigratePrettyPrintStyle,
 };
@@ -133,6 +134,7 @@ fn map_procedure_error(e: ProcedureCallError, procedure: &str) -> (StatusCode, S
 pub async fn call<S: ControlStateDelegate + NodeDelegate>(
     State(worker_ctx): State<S>,
     Extension(auth): Extension<SpacetimeAuth>,
+    headers: HeaderMap,
     Path(CallParams {
         name_or_identity,
         reducer,
@@ -141,6 +143,10 @@ pub async fn call<S: ControlStateDelegate + NodeDelegate>(
     body: Bytes,
 ) -> axum::response::Result<impl IntoResponse> {
     let caller_identity = auth.claims.identity;
+    let tx_id = headers
+        .get(TX_ID_HEADER)
+        .and_then(|v| v.to_str().ok())
+        .and_then(|s| s.parse::<GlobalTxId>().ok());
 
     let args = parse_call_args(content_type, body)?;
 
@@ -162,6 +168,7 @@ pub async fn call<S: ControlStateDelegate + NodeDelegate>(
         .call_reducer_with_return(
             caller_identity,
             Some(connection_id),
+            tx_id,
             None,
             None,
             None,
@@ -251,6 +258,7 @@ fn parse_call_args(content_type: headers::ContentType, body: Bytes) -> axum::res
 pub async fn prepare<S: ControlStateDelegate + NodeDelegate>(
     State(worker_ctx): State<S>,
     Extension(auth): Extension<SpacetimeAuth>,
+    headers: HeaderMap,
     Path(CallParams {
         name_or_identity,
         reducer,
@@ -260,14 +268,18 @@ pub async fn prepare<S: ControlStateDelegate + NodeDelegate>(
 ) -> axum::response::Result<impl IntoResponse> {
     let args = parse_call_args(content_type, body)?;
     let caller_identity = auth.claims.identity;
+    let tx_id = headers
+        .get(TX_ID_HEADER)
+        .and_then(|v| v.to_str().ok())
+        .and_then(|s| s.parse::<GlobalTxId>().ok());
 
     let (module, Database { owner_identity, .. }) = find_module_and_database(&worker_ctx, name_or_identity).await?;
 
     // 2PC prepare is a server-to-server call; no client lifecycle management needed.
     // call_identity_connected/disconnected submit jobs to the module's executor, which
     // will be blocked holding the 2PC write lock after prepare_reducer returns — deadlock.
     let result = module
-        .prepare_reducer(caller_identity, None, &reducer, args)
+        .prepare_reducer(caller_identity, None, tx_id, &reducer, args)
         .await;
 
     match result {
@@ -298,6 +310,12 @@ pub struct TwoPcParams {
     prepare_id: String,
 }
 
+#[derive(Deserialize)]
+pub struct GlobalTxParams {
+    name_or_identity: NameOrIdentity,
+    global_tx_id: String,
+}
+
 /// 2PC commit endpoint: finalize a prepared transaction.
 ///
 /// `POST /v1/database/:name_or_identity/2pc/commit/:prepare_id`
@@ -389,6 +407,30 @@ pub async fn ack_commit_2pc<S: ControlStateDelegate + NodeDelegate>(
     Ok(StatusCode::OK)
 }
 
+/// 2PC wound endpoint.
+///
+/// `POST /v1/database/:name_or_identity/2pc/wound/:global_tx_id`
+pub async fn wound_2pc<S: ControlStateDelegate + NodeDelegate>(
+    State(worker_ctx): State<S>,
+    Extension(_auth): Extension<SpacetimeAuth>,
+    Path(GlobalTxParams {
+        name_or_identity,
+        global_tx_id,
+    }): Path<GlobalTxParams>,
+) -> axum::response::Result<impl IntoResponse> {
+    let tx_id = global_tx_id
+        .parse::<GlobalTxId>()
+        .map_err(|e| (StatusCode::BAD_REQUEST, e).into_response())?;
+    let (module, _database) = find_module_and_database(&worker_ctx, name_or_identity).await?;
+
+    module.wound_global_tx(tx_id).await.map_err(|e| {
+        log::warn!("2PC wound failed for {tx_id}: {e}");
+        (StatusCode::NOT_FOUND, e).into_response()
+    })?;
+
+    Ok(StatusCode::OK)
+}
+
 fn reducer_outcome_response(
     module: &ModuleHost,
     owner_identity: &Identity,
@@ -426,6 +468,7 @@ fn reducer_outcome_response(
             // TODO: different status code? this is what cloudflare uses, sorta
             Ok((StatusCode::from_u16(530).unwrap(), (*errmsg).into_response()))
         }
+        ReducerOutcome::Wounded(errmsg) => Ok((StatusCode::CONFLICT, (*errmsg).into_response())),
         ReducerOutcome::BudgetExceeded => {
             log::warn!("Node's energy budget exceeded for identity: {owner_identity} while executing {reducer}");
             Ok((
@@ -1401,6 +1444,8 @@ pub struct DatabaseRoutes<S> {
     pub commit_2pc_post: MethodRouter<S>,
     /// POST: /database/:name_or_identity/2pc/abort/:prepare_id
     pub abort_2pc_post: MethodRouter<S>,
+    /// POST: /database/:name_or_identity/2pc/wound/:global_tx_id
+    pub wound_2pc_post: MethodRouter<S>,
     /// GET: /database/:name_or_identity/2pc/status/:prepare_id
     pub status_2pc_get: MethodRouter<S>,
     /// POST: /database/:name_or_identity/2pc/ack-commit/:prepare_id
@@ -1433,6 +1478,7 @@ where
             prepare_post: post(prepare::<S>),
             commit_2pc_post: post(commit_2pc::<S>),
             abort_2pc_post: post(abort_2pc::<S>),
+            wound_2pc_post: post(wound_2pc::<S>),
             status_2pc_get: get(status_2pc::<S>),
             ack_commit_2pc_post: post(ack_commit_2pc::<S>),
         }
@@ -1463,6 +1509,7 @@ where
             .route("/prepare/:reducer", self.prepare_post)
             .route("/2pc/commit/:prepare_id", self.commit_2pc_post)
             .route("/2pc/abort/:prepare_id", self.abort_2pc_post)
+            .route("/2pc/wound/:global_tx_id", self.wound_2pc_post)
             .route("/2pc/status/:prepare_id", self.status_2pc_get)
             .route("/2pc/ack-commit/:prepare_id", self.ack_commit_2pc_post);
 

crates/core/DISTRIBUTED-WOUND-WAIT.md

@@ -0,0 +1,143 @@
+# Distributed Wound-Wait for 2PC Deadlock Breaking
+
+## Problem
+
+Distributed reducers can deadlock when one distributed transaction holds a participant lock on one database and waits on another database locked by a younger distributed transaction. Existing 2PC ensures atomic commit/abort, but it does not resolve distributed lock cycles.
+
+## Chosen Model
+
+- Use wound-wait.
+- Transaction identity is `GlobalTxId`.
+- `GlobalTxId.creator_db` is the authoritative coordinator database.
+- Participant and coordinator runtime state are keyed by `GlobalTxId`.
+- `prepare_id` remains a participant-local 2PC phase handle only.
+- Older transactions wound younger transactions.
+- Younger transactions wait behind older transactions.
+- Lock acquisition is managed by an owner-aware scheduler that tracks running and pending `GlobalTxId`s.
+- A wound RPC is required because a younger lock holder may belong to a distributed transaction coordinated on a different database.
+
+## Runtime Model
+
+- Add a distributed session registry keyed by `GlobalTxId`.
+- Session tracks role, state, local `prepare_id`, coordinator identity, participants, and a local wounded/abort signal.
+- Add a per-database async lock scheduler for distributed reducer write acquisition.
+- Scheduler state includes current running owner, pending queue/set, and wounded marker for the current owner.
+- Requesters await scheduler admission before reducer execution starts a local mutable transaction.
+
+## Wound Protocol
+
+Participant detecting conflict compares requester and owner by `GlobalTxId` ordering.
+
+- If requester is older:
+  - mark local owner as wounded
+  - send wound RPC to coordinator `GlobalTxId.creator_db` if needed
+  - keep requester pending until local owner releases
+- If requester is younger:
+  - requester stays pending
+
+Wound RPC is idempotent and targets the distributed session at the coordinator.
+
+Coordinator receiving wound:
+
+- transitions session to `Aborting`
+- sets wounded flag
+- aborts local execution cooperatively
+- fans out abort to known prepared participants
+
+## Safe Points
+
+- Before remote reducer calls
+- Before PREPARE / COMMIT path work
+- After reducer body returns, before expensive post-processing
+
+On safe-point wound detection:
+
+- rollback local tx
+- unregister scheduler ownership
+- wake waiters
+- surface retryable `wounded` error
+
+## Compatibility
+
+- Keep existing `/2pc/commit`, `/2pc/abort`, `/2pc/status`, and ack-commit flows.
+- Add a new wound endpoint.
+- `/prepare` must propagate `GlobalTxId` the same way `/call` already does.
+- No durable format change unless recovery work later proves it necessary.
+
+## Implementation Sequence
+
+### 1. Propagate `GlobalTxId` through 2PC prepare path
+
+- Update outgoing 2PC prepare requests to send `X-Spacetime-Tx-Id`.
+- Update incoming `/prepare/:reducer` to parse `X-Spacetime-Tx-Id`.
+- Thread `GlobalTxId` through `prepare_reducer` and any participant execution params.
+- Ensure recovered/replayed participant work can recover or reconstruct the same session identity.
+
+### 2. Replace minimal prepared registry with `GlobalTxId` session registry
+
+- Extend the current prepared transaction registry into a session manager keyed by `GlobalTxId`.
+- Track:
+  - role
+  - state
+  - local `prepare_id`
+  - participants
+  - coordinator identity
+  - wounded/abort signal
+- Provide lookup by both `GlobalTxId` and `prepare_id`.
+
+### 3. Add distributed lock scheduler
+
+- Add an async scheduler in `core`, adjacent to reducer tx startup, not inside raw datastore locking.
+- Track running owner and pending `GlobalTxId`s.
+- Require distributed reducer write acquisition to await scheduler admission before blocking datastore acquisition.
+- Implement wound-wait ordering and wakeup behavior there.
+
+### 4. Add wound RPC endpoint and coordinator handler
+
+- Add `POST /v1/database/:name_or_identity/2pc/wound/:global_tx_id`.
+- Parse and route by `GlobalTxId.creator_db`.
+- Coordinator session handler must:
+  - mark session `Aborting`
+  - set wounded flag
+  - begin participant abort fanout
+  - behave idempotently
+
+### 5. Add cooperative abort checks in reducer execution
+
+- Add wound checks at required safe points.
+- On wound:
+  - rollback
+  - unregister running owner
+  - notify scheduler waiters
+  - surface retryable wounded error
+
+### 6. Integrate scheduler + wound with 2PC transitions
+
+- Ensure PREPARE, COMMIT, ABORT, and recovery all keep scheduler and session registry consistent.
+- Make local owner release happen on all terminal paths.
+- Keep participant recovery compatible with session state.
+
+### 7. Add tests
+
+- Scheduler ordering and wakeup tests
+- Local same-database wound tests
+- Distributed deadlock cycle tests
+- Wound RPC idempotency tests
+- Recovery tests for wounded prepared transactions
+- Regression tests for existing 2PC success/failure flows
+
+## Acceptance Criteria
+
+- Distributed deadlock cycles are broken deterministically by wound-wait.
+- Older distributed transactions eventually proceed without manual intervention.
+- Younger distributed transactions abort globally, not just locally.
+- `/prepare` and `/call` both carry `GlobalTxId`.
+- Existing 2PC happy paths continue to pass.
+- Repeated wound or abort requests are safe and idempotent.
+
+## Assumptions
+
+- `GlobalTxId.creator_db` is always the coordinator database.
+- `GlobalTxId` ordering is the authoritative age/tie-break rule.
+- Cooperative abort at safe points is sufficient for v1; no preemptive interruption is required.
+- Lock scheduler state is in-memory runtime state, not durable state.

crates/core/src/client/client_connection.rs

@@ -852,6 +852,7 @@ impl ClientConnection {
             .call_reducer(
                 self.id.identity,
                 Some(self.id.connection_id),
+                None,
                 caller,
                 Some(request_id),
                 Some(timer),
@@ -873,6 +874,7 @@ impl ClientConnection {
             .call_reducer(
                 self.id.identity,
                 Some(self.id.connection_id),
+                None,
                 Some(self.sender()),
                 Some(request_id),
                 Some(timer),

crates/core/src/client/messages.rs

@@ -379,7 +379,9 @@ impl ToProtocol for TransactionUpdateMessage {
 
             let status = match &event.status {
                 EventStatus::Committed(_) => ws_v1::UpdateStatus::Committed(update),
-                EventStatus::FailedUser(errmsg) | EventStatus::FailedInternal(errmsg) => {
+                EventStatus::FailedUser(errmsg)
+                | EventStatus::FailedInternal(errmsg)
+                | EventStatus::Wounded(errmsg) => {
                     ws_v1::UpdateStatus::Failed(errmsg.clone().into())
                 }
                 EventStatus::OutOfEnergy => ws_v1::UpdateStatus::OutOfEnergy,

crates/core/src/error.rs

@@ -275,6 +275,8 @@ pub enum NodesError {
     BadIndexType(u8),
     #[error("Failed to scheduled timer: {0}")]
     ScheduleError(#[source] ScheduleError),
+    #[error("Distributed transaction wounded: {0}")]
+    Wounded(String),
     #[error("HTTP request failed: {0}")]
     HttpError(String),
 }