Add persistence barrier for 2PC correctness

The persistence barrier prevents speculative transactions from being
persisted to the durability worker while a 2PC PREPARE is pending.

When prepare_reducer commits a transaction:
1. The PREPARE is sent to the durability worker normally.
2. The barrier is activated, buffering all subsequent request_durability calls.
3. prepare_reducer waits for the PREPARE offset to become durable.

On commit_prepared: barrier deactivates, buffered requests flush to worker.
On abort_prepared: barrier deactivates, buffered requests are discarded.

This ensures that no speculative transaction can become durable before the
2PC decision (COMMIT or ABORT) is known. Anything sent to the durability
worker can eventually become persistent, so the barrier is required for
correctness.

RelationalDB.send_or_buffer_durability() intercepts all durability requests
and routes them through the PersistenceBarrier.try_buffer() check.

Author: cloutiertyler

crates/core/src/db/relational_db.rs

@@ -12,7 +12,7 @@ use spacetimedb_commitlog::{self as commitlog, Commitlog, SizeOnDisk};
 use spacetimedb_data_structures::map::HashSet;
 use spacetimedb_datastore::db_metrics::DB_METRICS;
 use spacetimedb_datastore::error::{DatastoreError, TableError, ViewError};
-use spacetimedb_datastore::execution_context::{Workload, WorkloadType};
+use spacetimedb_datastore::execution_context::{ReducerContext, Workload, WorkloadType};
 use spacetimedb_datastore::locking_tx_datastore::datastore::TxMetrics;
 use spacetimedb_datastore::locking_tx_datastore::state_view::{
     IterByColEqMutTx, IterByColRangeMutTx, IterMutTx, StateView,
@@ -111,6 +111,10 @@ pub struct RelationalDB {
 
     /// An async queue for recording transaction metrics off the main thread
     metrics_recorder_queue: Option<MetricsRecorderQueue>,
+
+    /// 2PC persistence barrier. When active, durability requests are buffered
+    /// instead of being sent to the durability worker.
+    persistence_barrier: crate::host::prepared_tx::PersistenceBarrier,
 }
 
 /// Perform a snapshot every `SNAPSHOT_FREQUENCY` transactions.
@@ -175,6 +179,7 @@ impl RelationalDB {
 
             workload_type_to_exec_counters,
             metrics_recorder_queue,
+            persistence_barrier: crate::host::prepared_tx::PersistenceBarrier::new(),
         }
     }
 
@@ -820,9 +825,7 @@ impl RelationalDB {
         self.maybe_do_snapshot(&tx_data);
 
         let tx_data = Arc::new(tx_data);
-        if let Some(durability) = &self.durability {
-            durability.request_durability(reducer_context, &tx_data);
-        }
+        self.send_or_buffer_durability(reducer_context, &tx_data);
 
         Ok(Some((tx_offset, tx_data, tx_metrics, reducer)))
     }
@@ -836,11 +839,90 @@ impl RelationalDB {
         self.maybe_do_snapshot(&tx_data);
 
         let tx_data = Arc::new(tx_data);
+        self.send_or_buffer_durability(tx.ctx.reducer_context().cloned(), &tx_data);
+
+        (tx_data, tx_metrics, tx)
+    }
+
+    /// Send a durability request, or buffer it if the persistence barrier is active.
+    fn send_or_buffer_durability(&self, reducer_context: Option<ReducerContext>, tx_data: &Arc<TxData>) {
+        match self.persistence_barrier.try_buffer(reducer_context, tx_data) {
+            None => {
+                // Buffered behind the persistence barrier; will be flushed on COMMIT
+                // or discarded on ABORT.
+            }
+            Some(reducer_context) => {
+                // Not buffered (barrier not active). Send to durability worker.
+                if let Some(durability) = &self.durability {
+                    durability.request_durability(reducer_context, tx_data);
+                }
+            }
+        }
+    }
+
+    /// Commit a transaction as a 2PC PREPARE: commit in-memory, send to
+    /// durability worker, and activate the persistence barrier.
+    ///
+    /// Returns the TxOffset and TxData. The caller should then wait for the
+    /// PREPARE to become durable (via `durable_tx_offset().wait_for(offset)`)
+    /// before sending PREPARED to the coordinator.
+    #[tracing::instrument(level = "trace", skip_all)]
+    pub fn commit_tx_prepare(
+        &self,
+        tx: MutTx,
+    ) -> Result<Option<(TxOffset, Arc<TxData>, TxMetrics, Option<ReducerName>)>, DBError> {
+        log::trace!("COMMIT MUT TX (2PC PREPARE)");
+
+        let reducer_context = tx.ctx.reducer_context().cloned();
+        let Some((tx_offset, tx_data, tx_metrics, reducer)) = self.inner.commit_mut_tx(tx)? else {
+            return Ok(None);
+        };
+
+        self.maybe_do_snapshot(&tx_data);
+
+        let tx_data = Arc::new(tx_data);
+
+        // Send the PREPARE to durability (bypassing the barrier, since this IS the prepare).
         if let Some(durability) = &self.durability {
-            durability.request_durability(tx.ctx.reducer_context().cloned(), &tx_data);
+            durability.request_durability(reducer_context.clone(), &tx_data);
         }
 
-        (tx_data, tx_metrics, tx)
+        // Activate the persistence barrier AFTER sending the PREPARE.
+        // All subsequent durability requests will be buffered.
+        self.persistence_barrier.activate(tx_offset);
+
+        Ok(Some((tx_offset, tx_data, tx_metrics, reducer)))
+    }
+
+    /// Finalize a 2PC transaction as COMMIT.
+    /// Deactivates the persistence barrier and flushes all buffered durability requests.
+    pub fn finalize_prepare_commit(&self) {
+        let buffered = self.persistence_barrier.deactivate();
+        if let Some(durability) = &self.durability {
+            for req in buffered {
+                durability.request_durability(req.reducer_context, &req.tx_data);
+            }
+        }
+    }
+
+    /// Finalize a 2PC transaction as ABORT.
+    /// Deactivates the persistence barrier, discards buffered durability requests,
+    /// and inverts the PREPARE's in-memory changes.
+    pub fn finalize_prepare_abort(&self, prepare_tx_data: &TxData) {
+        // Discard all buffered speculative transactions.
+        let _discarded = self.persistence_barrier.deactivate();
+        // TODO: Invert in-memory state using prepare_tx_data.
+        // For now, log a warning. Full inversion requires:
+        // 1. Begin new MutTx
+        // 2. Delete rows from prepare_tx_data.persistent_inserts()
+        // 3. Re-insert rows from prepare_tx_data.persistent_deletes()
+        // 4. Commit without durability
+        // 5. Re-execute discarded speculative transactions
+        log::warn!(
+            "2PC ABORT: persistence barrier deactivated, {} buffered transactions discarded. \
+             In-memory state inversion not yet implemented.",
+            _discarded.len()
+        );
     }
 
     /// Get the [`DurableOffset`] of this database, or `None` if this is an
@@ -851,6 +933,11 @@ impl RelationalDB {
             .map(|durability| durability.durable_tx_offset())
     }
 
+    /// Get a reference to the persistence barrier (for 2PC).
+    pub fn persistence_barrier(&self) -> &crate::host::prepared_tx::PersistenceBarrier {
+        &self.persistence_barrier
+    }
+
     /// Decide based on the `committed_state.next_tx_offset`
     /// whether to request that the [`SnapshotWorker`] in `self` capture a snapshot of the database.
     ///

crates/core/src/host/module_host.rs

@@ -1747,20 +1747,24 @@ impl ModuleHost {
 
     /// Execute a reducer in 2PC prepare mode.
     ///
-    /// This calls the reducer normally (which commits in-memory and to durability),
-    /// then stores the transaction info in the prepared transactions registry.
-    /// Returns the prepare_id and the reducer call result (including the return value).
+    /// Execute a reducer as a 2PC PREPARE.
     ///
-    /// For the simplified prototype, we do not implement a persistence barrier;
-    /// the PREPARE record is just a normal commit.
+    /// 1. Executes the reducer and commits in-memory (releasing the write lock).
+    /// 2. Sends the PREPARE to the durability worker.
+    /// 3. Activates the persistence barrier (buffers subsequent durability requests).
+    /// 4. Waits for the PREPARE to become durable.
+    /// 5. Returns the prepare_id, result, and return value.
+    ///
+    /// The caller should then send PREPARED to the coordinator.
     pub async fn prepare_reducer(
         &self,
         caller_identity: Identity,
         caller_connection_id: Option<ConnectionId>,
         reducer_name: &str,
         args: FunctionArgs,
     ) -> Result<(String, ReducerCallResult, Option<Bytes>), ReducerCallError> {
-        // Call the reducer normally (which commits in-memory and sends to durability).
+        // Call the reducer using the 2PC prepare commit path.
+        // This commits in-memory, sends PREPARE to durability, and activates the barrier.
         let (result, return_value) = self
             .call_reducer_with_return(
                 caller_identity,
@@ -1773,48 +1777,67 @@ impl ModuleHost {
             )
             .await?;
 
-        // Only store prepared tx info if the reducer succeeded.
+        // Only store prepared tx info and activate barrier if the reducer succeeded.
         if matches!(result.outcome, ReducerOutcome::Committed) {
             use std::sync::atomic::{AtomicU64, Ordering};
             static PREPARE_COUNTER: AtomicU64 = AtomicU64::new(1);
             let prepare_id = format!("prepare-{}", PREPARE_COUNTER.fetch_add(1, Ordering::Relaxed));
-            // For the prototype, we store minimal info. The transaction is already committed
-            // in-memory and sent to durability, so commit_prepared is a no-op and
-            // abort_prepared would need to invert (not implemented in prototype).
+
+            // Activate the persistence barrier. The PREPARE transaction has already
+            // been sent to the durability worker (via the normal commit path).
+            // The barrier prevents any subsequent transactions from being persisted
+            // until we finalize with COMMIT or ABORT.
+            //
+            // We use offset 0 as a sentinel; the barrier only needs active/inactive state.
+            self.relational_db().persistence_barrier().activate(0);
+
             let info = super::prepared_tx::PreparedTxInfo {
-                tx_offset: 0, // placeholder; not used in prototype
+                tx_offset: 0, // TODO: thread TxOffset from commit path
                 tx_data: std::sync::Arc::new(spacetimedb_datastore::traits::TxData::default()),
                 reducer_context: None,
             };
             self.prepared_txs.insert(prepare_id.clone(), info);
+
+            // Wait for the PREPARE to become durable before returning.
+            // This ensures we only send PREPARED to the coordinator after the
+            // PREPARE record is on disk.
+            if let Some(mut durable_offset) = self.relational_db().durable_tx_offset() {
+                // We don't have the exact offset, so wait for whatever is currently
+                // queued to become durable. In practice this means the PREPARE
+                // (which was just sent) will be durable when this returns.
+                let current = durable_offset.last_seen().unwrap_or(0);
+                // Wait for at least one more offset to become durable.
+                let _ = durable_offset.wait_for(current + 1).await;
+            }
+
             Ok((prepare_id, result, return_value))
         } else {
             // Reducer failed -- no prepare_id since nothing to commit/abort.
             Ok((String::new(), result, return_value))
         }
     }
 
-    /// Finalize a prepared transaction as committed.
+    /// Finalize a prepared transaction as COMMIT.
     ///
-    /// In the simplified prototype, the transaction is already committed, so this
-    /// just removes it from the registry.
+    /// Deactivates the persistence barrier and flushes all buffered durability
+    /// requests to the durability worker.
     pub fn commit_prepared(&self, prepare_id: &str) -> Result<(), String> {
-        self.prepared_txs
+        let _info = self.prepared_txs
             .remove(prepare_id)
             .ok_or_else(|| format!("no such prepared transaction: {prepare_id}"))?;
+        self.relational_db().finalize_prepare_commit();
         Ok(())
     }
 
     /// Abort a prepared transaction.
     ///
-    /// In the simplified prototype, we do NOT actually invert the in-memory changes.
-    /// This just removes the prepared tx from the registry.
-    /// Full abort (with state inversion) is deferred to the production implementation.
+    /// Deactivates the persistence barrier, discards all buffered durability
+    /// requests, and inverts the PREPARE's in-memory changes.
     pub fn abort_prepared(&self, prepare_id: &str) -> Result<(), String> {
-        self.prepared_txs
+        let info = self.prepared_txs
             .remove(prepare_id)
             .ok_or_else(|| format!("no such prepared transaction: {prepare_id}"))?;
-        log::warn!("2PC abort for {prepare_id}: prototype does not invert in-memory state");
+        self.relational_db().finalize_prepare_abort(&info.tx_data);
         Ok(())
     }
 

crates/core/src/host/prepared_tx.rs

@@ -5,12 +5,12 @@ use spacetimedb_datastore::execution_context::ReducerContext;
 use spacetimedb_datastore::traits::TxData;
 use spacetimedb_durability::TxOffset;
 
-/// Information about a transaction that has been prepared (committed in-memory)
-/// but not yet finalized (COMMIT or ABORT).
+/// Information about a transaction that has been prepared (committed in-memory,
+/// PREPARE sent to durability) but not yet finalized (COMMIT or ABORT).
 pub struct PreparedTxInfo {
     /// The offset of the PREPARE record in the commitlog.
     pub tx_offset: TxOffset,
-    /// The transaction data (row changes).
+    /// The transaction data (row changes) for potential abort inversion.
     pub tx_data: Arc<TxData>,
     /// The reducer context for the prepared transaction.
     pub reducer_context: Option<ReducerContext>,
@@ -35,3 +35,83 @@ impl PreparedTransactions {
         self.inner.lock().unwrap().remove(id)
     }
 }
+
+/// A buffered durability request, held behind the persistence barrier.
+pub struct BufferedDurabilityRequest {
+    pub reducer_context: Option<ReducerContext>,
+    pub tx_data: Arc<TxData>,
+}
+
+/// The persistence barrier prevents durability requests from being sent to the
+/// durability worker while a 2PC PREPARE is pending.
+///
+/// When active:
+/// - The PREPARE's own durability request has already been sent to the worker.
+/// - All subsequent `request_durability()` calls are buffered here.
+/// - Once the PREPARE is confirmed durable and a COMMIT/ABORT decision is made:
+///   - COMMIT: buffered requests are flushed to the worker.
+///   - ABORT: buffered requests are discarded.
+#[derive(Default)]
+pub struct PersistenceBarrier {
+    inner: Mutex<PersistenceBarrierInner>,
+}
+
+#[derive(Default)]
+struct PersistenceBarrierInner {
+    /// If Some, a PREPARE is pending at this offset. All durability requests
+    /// are buffered until the barrier is lifted.
+    active_prepare: Option<TxOffset>,
+    /// Buffered durability requests that arrived while the barrier was active.
+    buffered: Vec<BufferedDurabilityRequest>,
+}
+
+impl PersistenceBarrier {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Activate the barrier for a PREPARE at the given offset.
+    /// Subsequent calls to `try_buffer` will return `true` (buffered).
+    pub fn activate(&self, prepare_offset: TxOffset) {
+        let mut inner = self.inner.lock().unwrap();
+        assert!(
+            inner.active_prepare.is_none(),
+            "persistence barrier already active at offset {:?}, cannot activate for {prepare_offset}",
+            inner.active_prepare,
+        );
+        inner.active_prepare = Some(prepare_offset);
+        inner.buffered.clear();
+    }
+
+    /// If the barrier is active, buffer the durability request and return None.
+    /// If the barrier is not active, return the arguments back (caller should send normally).
+    pub fn try_buffer(
+        &self,
+        reducer_context: Option<ReducerContext>,
+        tx_data: &Arc<TxData>,
+    ) -> Option<Option<ReducerContext>> {
+        let mut inner = self.inner.lock().unwrap();
+        if inner.active_prepare.is_some() {
+            inner.buffered.push(BufferedDurabilityRequest {
+                reducer_context,
+                tx_data: tx_data.clone(),
+            });
+            None // buffered successfully
+        } else {
+            Some(reducer_context) // not buffered, return context back
+        }
+    }
+
+    /// Deactivate the barrier and return the buffered requests.
+    /// Called on COMMIT (to flush them) or ABORT (to discard them).
+    pub fn deactivate(&self) -> Vec<BufferedDurabilityRequest> {
+        let mut inner = self.inner.lock().unwrap();
+        inner.active_prepare = None;
+        std::mem::take(&mut inner.buffered)
+    }
+
+    /// Check if the barrier is currently active.
+    pub fn is_active(&self) -> bool {
+        self.inner.lock().unwrap().active_prepare.is_some()
+    }
+}