replay: don't make ReplayCommittedState again and again

Author: Centril

crates/datastore/src/locking_tx_datastore/datastore.rs

@@ -24,7 +24,7 @@ use crate::{
     },
 };
 use anyhow::anyhow;
-use core::{cell::RefCell, ops::RangeBounds};
+use core::ops::RangeBounds;
 use parking_lot::{Mutex, RwLock};
 use spacetimedb_data_structures::map::{HashCollectionExt, HashMap};
 use spacetimedb_durability::TxOffset;
@@ -170,13 +170,9 @@ impl Locking {
     /// The provided closure will be called for each transaction found in the
     /// history, the parameter is the transaction's offset. The closure is called
     /// _before_ the transaction is applied to the database state.
-    pub fn replay<F: FnMut(u64)>(&self, progress: F, error_behavior: ErrorBehavior) -> Replay<F> {
-        Replay {
-            database_identity: self.database_identity,
-            committed_state: self.committed_state.clone(),
-            progress: RefCell::new(progress),
-            error_behavior,
-        }
+    pub fn replay<F: FnMut(u64)>(&self, progress: F, error_behavior: ErrorBehavior) -> Replay<'_, F> {
+        let committed_state = self.committed_state.write();
+        Replay::new(self.database_identity, committed_state, progress, error_behavior)
     }
 
     /// Construct a new [`Locking`] datastore containing the state stored in `snapshot`.

crates/datastore/src/locking_tx_datastore/replay.rs

@@ -7,8 +7,9 @@ use crate::system_tables::{
     is_built_in_meta_row, StColumnRow, StFields as _, StTableFields, StTableRow, ST_COLUMN_ID, ST_TABLE_ID,
 };
 use anyhow::{anyhow, Context};
+use core::cell::RefMut;
 use core::ops::{Deref, DerefMut, RangeBounds};
-use parking_lot::{RwLock, RwLockReadGuard};
+use parking_lot::RwLockWriteGuard;
 use prometheus::core::{AtomicF64, GenericGauge};
 use prometheus::IntGauge;
 use spacetimedb_commitlog::payload::txdata;
@@ -79,6 +80,7 @@ pub fn apply_history(
         .set((end_tx_offset - start_tx_offset) as _);
 
     log::info!("[{database_identity}] DATABASE: applied transaction history");
+    drop(replay); // Neccessary to avoid a deadlock.
     datastore.rebuild_state_after_replay()?;
     log::info!("[{database_identity}] DATABASE: rebuilt state after replay");
 
@@ -104,21 +106,32 @@ pub enum ReplayError {
 
 /// A [`spacetimedb_commitlog::Decoder`] suitable for replaying a transaction
 /// history into the database state.
-pub struct Replay<F> {
-    pub(super) database_identity: Identity,
-    pub(super) committed_state: Arc<RwLock<CommittedState>>,
-    pub(super) progress: RefCell<F>,
-    pub(super) error_behavior: ErrorBehavior,
+pub struct Replay<'a, F> {
+    database_identity: Identity,
+    committed_state: RefCell<ReplayCommittedState<'a>>,
+    progress: RefCell<F>,
+    error_behavior: ErrorBehavior,
 }
 
-impl<F> Replay<F> {
-    fn using_visitor<T>(&self, f: impl FnOnce(&mut ReplayVisitor<'_, F>) -> T) -> T {
-        let mut committed_state = self.committed_state.write();
-        let state = &mut *committed_state;
-        let committed_state = ReplayCommittedState::new(state);
+impl<'a, F> Replay<'a, F> {
+    pub fn new(
+        database_identity: Identity,
+        committed_state: RwLockWriteGuard<'a, CommittedState>,
+        progress: F,
+        error_behavior: ErrorBehavior,
+    ) -> Self {
+        Self {
+            database_identity,
+            committed_state: RefCell::new(ReplayCommittedState::new(committed_state)),
+            progress: RefCell::new(progress),
+            error_behavior,
+        }
+    }
+
+    fn using_visitor<T>(&self, f: impl FnOnce(&mut ReplayVisitor<'_, '_, F>) -> T) -> T {
         let mut visitor = ReplayVisitor {
             database_identity: &self.database_identity,
-            committed_state,
+            committed_state: &mut self.committed_state.borrow_mut(),
             progress: &mut *self.progress.borrow_mut(),
             dropped_table_names: IntMap::default(),
             error_behavior: self.error_behavior,
@@ -127,16 +140,16 @@ impl<F> Replay<F> {
     }
 
     pub fn next_tx_offset(&self) -> u64 {
-        self.committed_state.read_arc().next_tx_offset
+        self.committed_state.borrow().next_tx_offset
     }
 
     // NOTE: This is not unused.
-    pub fn committed_state(&self) -> RwLockReadGuard<'_, CommittedState> {
-        self.committed_state.read()
+    pub fn committed_state(&self) -> RefMut<'_, ReplayCommittedState<'a>> {
+        self.committed_state.borrow_mut()
     }
 }
 
-impl<F: FnMut(u64)> spacetimedb_commitlog::Decoder for &mut Replay<F> {
+impl<F: FnMut(u64)> spacetimedb_commitlog::Decoder for &mut Replay<'_, F> {
     type Record = txdata::Txdata<ProductValue>;
     type Error = txdata::DecoderError<ReplayError>;
 
@@ -217,9 +230,9 @@ pub enum ErrorBehavior {
     Warn,
 }
 
-struct ReplayVisitor<'a, F> {
+struct ReplayVisitor<'a, 'cs, F> {
     database_identity: &'a Identity,
-    committed_state: ReplayCommittedState<'a>,
+    committed_state: &'a mut ReplayCommittedState<'cs>,
     progress: &'a mut F,
     // Since deletes are handled before truncation / drop, sometimes the schema
     // info is gone. We save the name on the first delete of that table so metrics
@@ -228,7 +241,7 @@ struct ReplayVisitor<'a, F> {
     error_behavior: ErrorBehavior,
 }
 
-impl<F> ReplayVisitor<'_, F> {
+impl<F> ReplayVisitor<'_, '_, F> {
     /// Process `err` according to `self.error_behavior`,
     /// either warning about it or returning it.
     ///
@@ -244,7 +257,7 @@ impl<F> ReplayVisitor<'_, F> {
     }
 }
 
-impl<F: FnMut(u64)> spacetimedb_commitlog::payload::txdata::Visitor for ReplayVisitor<'_, F> {
+impl<F: FnMut(u64)> spacetimedb_commitlog::payload::txdata::Visitor for ReplayVisitor<'_, '_, F> {
     type Error = ReplayError;
     // NOTE: Technically, this could be `()` if and when we can extract the
     // row data without going through `ProductValue` (PV).
@@ -392,9 +405,9 @@ impl<F: FnMut(u64)> spacetimedb_commitlog::payload::txdata::Visitor for ReplayVi
 }
 
 /// A `CommittedState` under construction during replay.
-struct ReplayCommittedState<'cs> {
-    /// The committed state being contructed.
-    state: &'cs mut CommittedState,
+pub struct ReplayCommittedState<'cs> {
+    /// The committed state being constructed.
+    state: RwLockWriteGuard<'cs, CommittedState>,
 
     /// Whether the table was dropped within the current transaction during replay.
     ///
@@ -439,25 +452,25 @@ struct ReplayCommittedState<'cs> {
     ///
     /// [`RowPointer`]s from this set are passed to the `unsafe` [`Table::get_row_ref_unchecked`],
     /// so it's important to properly maintain only [`RowPointer`]s to valid, extant, non-deleted rows.
-    pub(super) replay_table_updated: IntMap<TableId, RowPointer>,
+    replay_table_updated: IntMap<TableId, RowPointer>,
 }
 
 impl Deref for ReplayCommittedState<'_> {
     type Target = CommittedState;
 
     fn deref(&self) -> &Self::Target {
-        self.state
+        &self.state
     }
 }
 
 impl DerefMut for ReplayCommittedState<'_> {
     fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut *self.state
+        &mut self.state
     }
 }
 
 impl<'cs> ReplayCommittedState<'cs> {
-    fn new(state: &'cs mut CommittedState) -> Self {
+    fn new(state: RwLockWriteGuard<'cs, CommittedState>) -> Self {
         Self {
             state,
             replay_table_dropped: <_>::default(),
@@ -844,7 +857,7 @@ mod tests {
         // Directly call replay_insert on committed state.
         let row = u32_str_u32(1, "Carol", 40);
         {
-            let state = &mut *datastore.committed_state.write();
+            let state = datastore.committed_state.write();
             let mut committed_state = ReplayCommittedState::new(state);
             committed_state.replay_insert(table_id, &schema, &row)?;
         }