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;
@@ -104,21 +105,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 +139,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 +229,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 +240,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 +256,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 +404,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.
     ///
@@ -446,18 +458,18 @@ 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(),
@@ -590,7 +602,7 @@ impl<'cs> ReplayCommittedState<'cs> {
         let target_col_id = ColId::deserialize(ValueDeserializer::from_ref(&st_column_row.elements[1]))
             .expect("second field in `st_column` should decode to a `ColId`");
 
-        let outdated_st_column_rows = iter_st_column_for_table(self.state, &target_table_id.into())?
+        let outdated_st_column_rows = iter_st_column_for_table(self, &target_table_id.into())?
             .filter_map(|row_ref| {
                 StColumnRow::try_from(row_ref)
                     .map(|c| (c.col_pos == target_col_id && row_ref.pointer() != row_ptr).then(|| row_ref.pointer()))
@@ -620,7 +632,7 @@ impl<'cs> ReplayCommittedState<'cs> {
         // and not the other one, as it is being replaced.
         // `Self::ignore_previous_version_of_column` has marked the old version as ignored,
         // so filter only the non-ignored columns.
-        let mut columns = iter_st_column_for_table(self.state, &table_id.into())?
+        let mut columns = iter_st_column_for_table(self, &table_id.into())?
             .filter(|row_ref| !self.replay_columns_to_ignore.contains(&row_ref.pointer()))
             .map(|row_ref| {
                 let row = StColumnRow::try_from(row_ref)?;
@@ -844,7 +856,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)?;
         }