Move payload conversion into Durability trait

And remove intermediate worker from generic durability path entirely.

Author: joshua-spacetime

crates/core/src/db/durability.rs

@@ -9,7 +9,7 @@ use spacetimedb_snapshot::SnapshotRepository;
 use crate::{messages::control_db::Database, util::asyncify};
 
 use super::{
-    relational_db::{self, LocalDurability, Txdata},
+    relational_db::{self, Txdata},
     snapshot::{self, SnapshotDatabaseState, SnapshotWorker},
 };
 
@@ -30,11 +30,6 @@ pub type DiskSizeFn = Arc<dyn Fn() -> io::Result<SizeOnDisk> + Send + Sync>;
 pub struct Persistence {
     /// The [Durability] to use, for persisting transactions.
     pub durability: Arc<Durability>,
-    /// TODO: Merge this local durability handle with the generic trait object above.
-    /// This allows us to bypass an actor whose only responsibility is to generate a
-    /// commitlog payload from `TxData`. Ultimately though this should just be a part
-    /// of the [Durability] implementation itself.
-    pub local_durability: Option<LocalDurability>,
     /// The [DiskSizeFn].
     ///
     /// Currently the expectation is that the reported size is the commitlog
@@ -60,7 +55,6 @@ impl Persistence {
     ) -> Self {
         Self {
             durability: Arc::new(durability),
-            local_durability: None,
             disk_size: Arc::new(disk_size),
             snapshots,
             runtime,
@@ -89,34 +83,23 @@ impl Persistence {
 
     /// Convenience to deconstruct an [Option<Self>] into parts.
     ///
-    /// Returns `(Some(durability), local_durability, Some(disk_size), Option<SnapshotWorker>, Some(runtime))`
-    /// if `this` is `Some`, and `(None, None, None, None, None)` if `this` is `None`.
-    #[allow(clippy::type_complexity)]
+    /// Returns `(Some(durability), Some(disk_size), Option<SnapshotWorker>, Some(runtime))`
+    /// if `this` is `Some`, and `(None, None, None, None)` if `this` is `None`.
     pub(super) fn unzip(
         this: Option<Self>,
     ) -> (
         Option<Arc<Durability>>,
-        Option<LocalDurability>,
         Option<DiskSizeFn>,
         Option<SnapshotWorker>,
         Option<tokio::runtime::Handle>,
     ) {
         this.map(
             |Self {
                  durability,
-                 local_durability,
                  disk_size,
                  snapshots,
                  runtime,
-             }| {
-                (
-                    Some(durability),
-                    local_durability,
-                    Some(disk_size),
-                    snapshots,
-                    Some(runtime),
-                )
-            },
+             }| (Some(durability), Some(disk_size), snapshots, Some(runtime)),
         )
         .unwrap_or_default()
     }
@@ -176,8 +159,7 @@ impl PersistenceProvider for LocalPersistenceProvider {
         ));
 
         Ok(Persistence {
-            durability: durability.clone(),
-            local_durability: Some(durability),
+            durability,
             disk_size,
             snapshots: Some(snapshot_worker),
             runtime: tokio::runtime::Handle::current(),

crates/core/src/db/persistence.rs

@@ -1,4 +1,4 @@
-use crate::db::durability::DurabilityWorker;
+use crate::db::durability::{request_durability, spawn_close as spawn_durability_close};
 use crate::db::MetricsRecorderQueue;
 use crate::error::{DBError, RestoreSnapshotError};
 use crate::subscription::ExecutionCounters;
@@ -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,
@@ -98,7 +98,8 @@ pub struct RelationalDB {
     owner_identity: Identity,
 
     inner: Locking,
-    durability: Option<DurabilityWorker>,
+    durability: Option<Arc<Durability>>,
+    durability_runtime: Option<tokio::runtime::Handle>,
     snapshot_worker: Option<SnapshotWorker>,
 
     row_count_fn: RowCountFn,
@@ -133,8 +134,8 @@ impl std::fmt::Debug for RelationalDB {
 impl Drop for RelationalDB {
     fn drop(&mut self) {
         // Attempt to flush the outstanding transactions.
-        if let Some(worker) = self.durability.take() {
-            worker.spawn_close(self.database_identity);
+        if let (Some(durability), Some(runtime)) = (self.durability.take(), self.durability_runtime.take()) {
+            spawn_durability_close(durability, &runtime, self.database_identity);
         }
     }
 }
@@ -150,18 +151,12 @@ impl RelationalDB {
         let workload_type_to_exec_counters =
             Arc::new(EnumMap::from_fn(|ty| ExecutionCounters::new(&ty, &database_identity)));
 
-        let (durability, local_durability, disk_size_fn, snapshot_worker, rt) = Persistence::unzip(persistence);
-        let durability = match (local_durability, durability, rt) {
-            (Some(local_durability), _, Some(rt)) => {
-                Some(DurabilityWorker::new_local(database_identity, local_durability, rt))
-            }
-            (None, Some(durability), Some(rt)) => Some(DurabilityWorker::new(database_identity, durability, rt)),
-            _ => None,
-        };
+        let (durability, disk_size_fn, snapshot_worker, durability_runtime) = Persistence::unzip(persistence);
 
         Self {
             inner,
             durability,
+            durability_runtime,
             snapshot_worker,
 
             database_identity,
@@ -811,9 +806,7 @@ impl RelationalDB {
         let reducer_context = tx.ctx.reducer_context().cloned();
         // TODO: Never returns `None` -- should it?
         let Some((tx_offset, tx_data, tx_metrics, reducer)) = self.inner.commit_mut_tx_and_then(tx, |tx_data| {
-            if let Some(durability) = &self.durability {
-                durability.request_durability(reducer_context, tx_data);
-            }
+            self.request_durability(reducer_context, tx_data);
         })?
         else {
             return Ok(None);
@@ -830,9 +823,7 @@ impl RelationalDB {
 
         let reducer_context = tx.ctx.reducer_context().cloned();
         let (tx_data, tx_metrics, tx) = self.inner.commit_mut_tx_downgrade_and_then(tx, workload, |tx_data| {
-            if let Some(durability) = &self.durability {
-                durability.request_durability(reducer_context, tx_data);
-            }
+            self.request_durability(reducer_context, tx_data);
         });
 
         self.maybe_do_snapshot(&tx_data);
@@ -848,6 +839,12 @@ impl RelationalDB {
             .map(|durability| durability.durable_tx_offset())
     }
 
+    fn request_durability(&self, reducer_context: Option<ReducerContext>, tx_data: &Arc<TxData>) {
+        if let Some(durability) = &self.durability {
+            request_durability(durability.as_ref(), reducer_context, tx_data);
+        }
+    }
+
     /// Decide based on the `committed_state.next_tx_offset`
     /// whether to request that the [`SnapshotWorker`] in `self` capture a snapshot of the database.
     ///
@@ -1970,7 +1967,6 @@ pub mod tests_utils {
 
             let persistence = Persistence {
                 durability: local.clone(),
-                local_durability: Some(local.clone()),
                 disk_size: disk_size_fn,
                 snapshots,
                 runtime: rt,
@@ -2092,7 +2088,6 @@ pub mod tests_utils {
             let history = local.as_history();
             let persistence = Persistence {
                 durability: local.clone(),
-                local_durability: Some(local.clone()),
                 disk_size: disk_size_fn,
                 snapshots,
                 runtime: rt,

crates/core/src/db/relational_db.rs

@@ -1900,7 +1900,7 @@ mod tests {
     use spacetimedb_commitlog::{commitlog, repo};
     use spacetimedb_data_structures::map::{HashCollectionExt as _, HashMap};
     use spacetimedb_datastore::system_tables::{StRowLevelSecurityRow, ST_ROW_LEVEL_SECURITY_ID};
-    use spacetimedb_durability::{Durability, EmptyHistory, Transaction, TxOffset};
+    use spacetimedb_durability::{Durability, EmptyHistory, TxOffset};
     use spacetimedb_execution::dml::MutDatastore;
     use spacetimedb_lib::bsatn::ToBsatn;
     use spacetimedb_lib::db::auth::StAccess;
@@ -1987,7 +1987,8 @@ mod tests {
     impl Durability for ManualDurability {
         type TxData = Txdata;
 
-        fn append_tx(&self, tx: Transaction<Self::TxData>) {
+        fn append_tx(&self, tx: spacetimedb_durability::PreparedTx<Self::TxData>) {
+            let tx = tx.into_transaction();
             let mut commitlog = self.commitlog.write().unwrap();
             commitlog.commit([tx]).expect("commit failed");
             commitlog.flush().expect("error flushing commitlog");
@@ -2043,7 +2044,6 @@ mod tests {
             EmptyHistory::new(),
             Some(Persistence {
                 durability: durability.clone(),
-                local_durability: None,
                 disk_size: Arc::new(|| Ok(<_>::default())),
                 snapshots: None,
                 runtime: rt,

crates/core/src/subscription/module_subscription_actor.rs

@@ -22,7 +22,7 @@ use tokio::{
 };
 use tracing::{instrument, Span};
 
-use crate::{Close, Durability, DurableOffset, History, TxOffset};
+use crate::{Close, Durability, DurableOffset, History, PreparedTx, TxOffset};
 
 pub use spacetimedb_commitlog::repo::{OnNewSegmentFn, SizeOnDisk};
 
@@ -75,33 +75,6 @@ pub enum OpenError {
 
 type ShutdownReply = oneshot::Sender<OwnedNotified>;
 
-enum QueueItem<T> {
-    Ready(Transaction<Txdata<T>>),
-    Deferred { prepare: Box<dyn DeferredTx<T>> },
-}
-
-impl<T> QueueItem<T> {
-    fn prepare(self) -> Transaction<Txdata<T>> {
-        match self {
-            Self::Ready(tx) => tx,
-            Self::Deferred { prepare } => prepare.prepare(),
-        }
-    }
-}
-
-trait DeferredTx<T>: Send + Sync {
-    fn prepare(self: Box<Self>) -> Transaction<Txdata<T>>;
-}
-
-impl<T, F> DeferredTx<T> for F
-where
-    F: FnOnce() -> Transaction<Txdata<T>> + Send + Sync + 'static,
-{
-    fn prepare(self: Box<Self>) -> Transaction<Txdata<T>> {
-        self()
-    }
-}
-
 /// [`Durability`] implementation backed by a [`Commitlog`] on local storage.
 ///
 /// The commitlog is constrained to store the canonical [`Txdata`] payload,
@@ -124,7 +97,7 @@ pub struct Local<T> {
     ///
     /// The queue is bounded to
     /// `Options::QUEUE_CAPACITY_MULTIPLIER * Options::batch_capacity`.
-    queue: mpsc::Sender<QueueItem<T>>,
+    queue: mpsc::Sender<PreparedTx<Txdata<T>>>,
     /// How many transactions are pending durability, including items buffered
     /// in the queue and items currently being written by the actor.
     ///
@@ -245,7 +218,7 @@ impl<T: Encode + Send + Sync + 'static> Actor<T> {
     #[instrument(name = "durability::local::actor", skip_all)]
     async fn run(
         self,
-        mut transactions_rx: mpsc::Receiver<QueueItem<T>>,
+        mut transactions_rx: mpsc::Receiver<PreparedTx<Txdata<T>>>,
         mut shutdown_rx: mpsc::Receiver<oneshot::Sender<OwnedNotified>>,
     ) {
         info!("starting durability actor");
@@ -282,9 +255,9 @@ impl<T: Encode + Send + Sync + 'static> Actor<T> {
                     let clog = self.clog.clone();
                     let ready_len = tx_buf.len();
                     self.queue_depth.fetch_sub(ready_len as u64, Relaxed);
-                    tx_buf = spawn_blocking(move || -> io::Result<Vec<QueueItem<T>>> {
-                        for item in tx_buf.drain(..) {
-                            clog.commit([item.prepare()])?;
+                    tx_buf = spawn_blocking(move || -> io::Result<Vec<PreparedTx<Txdata<T>>>> {
+                        for tx in tx_buf.drain(..) {
+                            clog.commit([tx.into_transaction()])?;
                         }
                         Ok(tx_buf)
                     })
@@ -370,8 +343,8 @@ impl Drop for Lock {
 impl<T: Send + Sync + 'static> Durability for Local<T> {
     type TxData = Txdata<T>;
 
-    fn append_tx(&self, tx: Transaction<Self::TxData>) {
-        let _ = self.send_item(QueueItem::Ready(tx));
+    fn append_tx(&self, tx: PreparedTx<Self::TxData>) {
+        let _ = self.send_tx(tx);
     }
 
     fn durable_tx_offset(&self) -> DurableOffset {
@@ -411,17 +384,18 @@ impl<T: Send + Sync + 'static> Durability for Local<T> {
 }
 
 impl<T: Send + Sync + 'static> Local<T> {
-    fn send_item(&self, item: QueueItem<T>) -> bool {
+    fn send_tx(&self, tx: PreparedTx<Txdata<T>>) -> bool {
+        let mut tx = Some(tx);
         let blocked = match self.queue.try_reserve() {
             Ok(permit) => {
-                permit.send(item);
+                permit.send(tx.take().expect("tx already sent"));
                 false
             }
             Err(mpsc::error::TrySendError::Closed(_)) => {
                 panic!("durability actor crashed");
             }
             Err(mpsc::error::TrySendError::Full(_)) => {
-                let send = || self.queue.blocking_send(item);
+                let mut send = || self.queue.blocking_send(tx.take().expect("tx already sent"));
                 if tokio::runtime::Handle::try_current().is_ok() {
                     tokio::task::block_in_place(send)
                 } else {
@@ -435,12 +409,6 @@ impl<T: Send + Sync + 'static> Local<T> {
         self.queue_depth.fetch_add(1, Relaxed);
         blocked
     }
-
-    pub fn append_tx_deferred(&self, prepare: impl FnOnce() -> Transaction<Txdata<T>> + Send + Sync + 'static) -> bool {
-        self.send_item(QueueItem::Deferred {
-            prepare: Box::new(prepare),
-        })
-    }
 }
 
 impl<T: Encode + 'static> History for Commitlog<Txdata<T>> {

crates/durability/src/imp/local.rs

@@ -15,7 +15,7 @@ mod testing {
     use futures::FutureExt as _;
     use tokio::sync::watch;
 
-    use crate::{Close, Durability, DurableOffset, Transaction, TxOffset};
+    use crate::{Close, Durability, DurableOffset, PreparedTx, TxOffset};
 
     /// A [`Durability`] impl that sends all transactions into the void.
     ///
@@ -41,7 +41,7 @@ mod testing {
     impl<T: Send + Sync> Durability for NoDurability<T> {
         type TxData = T;
 
-        fn append_tx(&self, _: Transaction<Self::TxData>) {
+        fn append_tx(&self, _: PreparedTx<Self::TxData>) {
             if self.closed.load(Ordering::Relaxed) {
                 panic!("`close` was called on this `NoDurability` instance");
             }

crates/durability/src/imp/mod.rs

@@ -91,6 +91,29 @@ impl From<watch::Receiver<Option<TxOffset>>> for DurableOffset {
 /// can be used as a trait object without knowing the type of the `close` future.
 pub type Close = BoxFuture<'static, Option<TxOffset>>;
 
+/// Object-safe conversion into an owned [Transaction].
+pub trait IntoTransaction<T>: Send {
+    fn into_transaction(self: Box<Self>) -> Transaction<T>;
+}
+
+impl<T: Send + 'static> IntoTransaction<T> for Transaction<T> {
+    fn into_transaction(self: Box<Self>) -> Transaction<T> {
+        *self
+    }
+}
+
+impl<T, F> IntoTransaction<T> for F
+where
+    T: Send + 'static,
+    F: FnOnce() -> Transaction<T> + Send + 'static,
+{
+    fn into_transaction(self: Box<Self>) -> Transaction<T> {
+        self()
+    }
+}
+
+pub type PreparedTx<T> = Box<dyn IntoTransaction<T>>;
+
 /// The durability API.
 ///
 /// NOTE: This is a preliminary definition, still under consideration.
@@ -105,7 +128,7 @@ pub trait Durability: Send + Sync {
     /// The payload representing a single transaction.
     type TxData;
 
-    /// Submit a [Transaction] to be made durable.
+    /// Submit work that yields a [Transaction] to be made durable.
     ///
     /// This method must never block, and accept new transactions even if they
     /// cannot be made durable immediately.
@@ -121,7 +144,7 @@ pub trait Durability: Send + Sync {
     // (i.e. a torn write will corrupt all transactions contained in it), and it
     // is very unclear when it is both correct and beneficial to bundle more
     // than a single transaction into a commit.
-    fn append_tx(&self, tx: Transaction<Self::TxData>);
+    fn append_tx(&self, tx: PreparedTx<Self::TxData>);
 
     /// Obtain a handle to the [DurableOffset].
     fn durable_tx_offset(&self) -> DurableOffset;