Record transaction metrics off the main thread

Author: joshua-spacetime

crates/core/src/db/relational_db.rs

@@ -21,6 +21,7 @@ use crate::db::datastore::system_tables::StModuleRow;
 use crate::error::{DBError, DatabaseError, RestoreSnapshotError};
 use crate::execution_context::{ReducerContext, Workload, WorkloadType};
 use crate::messages::control_db::HostType;
+use crate::subscription::module_subscription_manager::{spawn_metrics_recorder, MetricsRecorderQueue};
 use crate::subscription::ExecutionCounters;
 use crate::util::{asyncify, spawn_rayon};
 use anyhow::{anyhow, Context};
@@ -118,6 +119,9 @@ pub struct RelationalDB {
 
     /// A map from workload types to their cached prometheus counters.
     workload_type_to_exec_counters: Arc<EnumMap<WorkloadType, ExecutionCounters>>,
+
+    /// An async queue for recording transaction metrics off the main thread
+    metrics_recorder_queue: MetricsRecorderQueue,
 }
 
 #[derive(Clone)]
@@ -237,6 +241,7 @@ impl RelationalDB {
             snapshot_repo.map(|repo| SnapshotWorker::new(inner.committed_state.clone(), repo.clone()));
         let workload_type_to_exec_counters =
             Arc::new(EnumMap::from_fn(|ty| ExecutionCounters::new(&ty, &database_identity)));
+        let metrics_recorder_queue = spawn_metrics_recorder();
 
         Self {
             inner,
@@ -251,6 +256,7 @@ impl RelationalDB {
 
             _lock: lock,
             workload_type_to_exec_counters,
+            metrics_recorder_queue,
         }
     }
 
@@ -749,6 +755,11 @@ impl RelationalDB {
         Ok(AlgebraicValue::decode(col_ty, &mut &*bytes)?)
     }
 
+    /// Returns the execution counters for this database.
+    pub fn exec_counter_map(&self) -> Arc<EnumMap<WorkloadType, ExecutionCounters>> {
+        self.workload_type_to_exec_counters.clone()
+    }
+
     /// Returns the execution counters for `workload_type` for this database.
     pub fn exec_counters_for(&self, workload_type: WorkloadType) -> &ExecutionCounters {
         &self.workload_type_to_exec_counters[workload_type]
@@ -988,7 +999,7 @@ impl RelationalDB {
         let mut tx = self.begin_tx(workload);
         let res = f(&mut tx);
         let (tx_metrics, reducer) = self.release_tx(tx);
-        self.report_tx_metricses(&reducer, None, None, &tx_metrics);
+        self.report_tx_metrics(reducer, Arc::new(None), None, Some(tx_metrics));
         res
     }
 
@@ -999,11 +1010,11 @@ impl RelationalDB {
     {
         if res.is_err() {
             let (tx_metrics, reducer) = self.rollback_mut_tx(tx);
-            self.report(&reducer, &tx_metrics, None);
+            self.report(reducer, tx_metrics, None);
         } else {
             match self.commit_tx(tx).map_err(E::from)? {
                 Some((tx_data, tx_metrics, reducer)) => {
-                    self.report(&reducer, &tx_metrics, Some(&tx_data));
+                    self.report(reducer, tx_metrics, Some(tx_data));
                 }
                 None => panic!("TODO: retry?"),
             }
@@ -1018,7 +1029,7 @@ impl RelationalDB {
         match res {
             Err(e) => {
                 let (tx_metrics, reducer) = self.rollback_mut_tx(tx);
-                self.report(&reducer, &tx_metrics, None);
+                self.report(reducer, tx_metrics, None);
 
                 Err(e)
             }
@@ -1042,17 +1053,20 @@ impl RelationalDB {
     /// Reports the `TxMetrics`s passed.
     ///
     /// Should only be called after the tx lock has been fully released.
-    pub(crate) fn report_tx_metricses(
+    pub(crate) fn report_tx_metrics(
         &self,
-        reducer: &str,
-        tx_data: Option<&TxData>,
-        metrics_mut: Option<&TxMetrics>,
-        metrics_read: &TxMetrics,
+        reducer: String,
+        tx_data: Arc<Option<TxData>>,
+        metrics_for_writer: Option<TxMetrics>,
+        metrics_for_reader: Option<TxMetrics>,
     ) {
-        if let Some(metrics_mut) = metrics_mut {
-            self.report(reducer, metrics_mut, tx_data);
-        }
-        self.report(reducer, metrics_read, None);
+        self.metrics_recorder_queue.send_metrics(
+            reducer,
+            metrics_for_writer,
+            metrics_for_reader,
+            tx_data,
+            self.exec_counter_map(),
+        );
     }
 }
 
@@ -1403,8 +1417,8 @@ impl RelationalDB {
     }
 
     /// Reports the metrics for `reducer`, using counters provided by `db`.
-    pub fn report(&self, reducer: &str, metrics: &TxMetrics, tx_data: Option<&TxData>) {
-        metrics.report(tx_data, reducer, |wl: WorkloadType| self.exec_counters_for(wl));
+    pub fn report(&self, reducer: String, metrics: TxMetrics, tx_data: Option<TxData>) {
+        self.report_tx_metrics(reducer, Arc::new(tx_data), Some(metrics), None);
     }
 
     /// Read the value of [ST_VARNAME_ROW_LIMIT] from `st_var`

crates/core/src/host/module_host.rs

@@ -185,7 +185,6 @@ pub struct ModuleEvent {
 }
 
 /// Information about a running module.
-#[derive(Debug)]
 pub struct ModuleInfo {
     /// The definition of the module.
     /// Loaded by loading the module's program from the system tables, extracting its definition,
@@ -205,6 +204,17 @@ pub struct ModuleInfo {
     pub metrics: ModuleMetrics,
 }
 
+impl fmt::Debug for ModuleInfo {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("ModuleInfo")
+            .field("module_def", &self.module_def)
+            .field("owner_identity", &self.owner_identity)
+            .field("database_identity", &self.database_identity)
+            .field("module_hash", &self.module_hash)
+            .finish()
+    }
+}
+
 #[derive(Debug)]
 pub struct ModuleMetrics {
     pub connected_clients: IntGauge,
@@ -368,7 +378,7 @@ fn init_database(
     let rcr = match module_def.lifecycle_reducer(Lifecycle::Init) {
         None => {
             if let Some((tx_data, tx_metrics, reducer)) = stdb.commit_tx(tx)? {
-                stdb.report(&reducer, &tx_metrics, Some(&tx_data));
+                stdb.report(reducer, tx_metrics, Some(tx_data));
             }
             None
         }

crates/core/src/host/wasm_common/module_host_actor.rs

@@ -293,12 +293,12 @@ impl<T: WasmInstance> ModuleInstance for WasmModuleInstance<T> {
                 log::warn!("Database update failed: {} @ {}", e, stdb.database_identity());
                 self.system_logger().warn(&format!("Database update failed: {e}"));
                 let (tx_metrics, reducer) = stdb.rollback_mut_tx(tx);
-                stdb.report(&reducer, &tx_metrics, None);
+                stdb.report(reducer, tx_metrics, None);
                 Ok(UpdateDatabaseResult::ErrorExecutingMigration(e))
             }
             Ok(()) => {
                 if let Some((tx_data, tx_metrics, reducer)) = stdb.commit_tx(tx)? {
-                    stdb.report(&reducer, &tx_metrics, Some(&tx_data));
+                    stdb.report(reducer, tx_metrics, Some(tx_data));
                 }
                 self.system_logger().info("Database updated");
                 log::info!("Database updated, {}", stdb.database_identity());

crates/core/src/sql/execute.rs

@@ -1,3 +1,4 @@
+use std::sync::Arc;
 use std::time::Duration;
 
 use super::ast::SchemaViewer;
@@ -202,7 +203,12 @@ pub fn run(
             // Release the tx on drop, so that we record metrics.
             let mut tx = scopeguard::guard(tx, |tx| {
                 let (tx_metrics_downgrade, reducer) = db.release_tx(tx);
-                db.report_tx_metricses(&reducer, Some(&tx_data), Some(&tx_metrics_mut), &tx_metrics_downgrade);
+                db.report_tx_metrics(
+                    reducer,
+                    Arc::new(Some(tx_data)),
+                    Some(tx_metrics_mut),
+                    Some(tx_metrics_downgrade),
+                );
             });
 
             // Compute the header for the result set
@@ -247,7 +253,7 @@ pub fn run(
                 let metrics = tx.metrics;
                 return db.commit_tx(tx).map(|tx_opt| {
                     if let Some((tx_data, tx_metrics, reducer)) = tx_opt {
-                        db.report(&reducer, &tx_metrics, Some(&tx_data));
+                        db.report(reducer, tx_metrics, Some(tx_data));
                     }
                     SqlResult { rows: vec![], metrics }
                 });

crates/core/src/subscription/module_subscription_actor.rs

@@ -37,7 +37,7 @@ use std::{sync::Arc, time::Instant};
 
 type Subscriptions = Arc<RwLock<SubscriptionManager>>;
 
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 pub struct ModuleSubscriptions {
     relational_db: Arc<RelationalDB>,
     /// If taking a lock (tx) on the db at the same time, ALWAYS lock the db first.
@@ -318,7 +318,7 @@ impl ModuleSubscriptions {
 
         let tx = scopeguard::guard(self.relational_db.begin_tx(Workload::Subscribe), |tx| {
             let (tx_metrics, reducer) = self.relational_db.release_tx(tx);
-            self.relational_db.report(&reducer, &tx_metrics, None);
+            self.relational_db.report(reducer, tx_metrics, None);
         });
 
         let existing_query = {
@@ -420,7 +420,7 @@ impl ModuleSubscriptions {
 
         let tx = scopeguard::guard(self.relational_db.begin_tx(Workload::Unsubscribe), |tx| {
             let (tx_metrics, reducer) = self.relational_db.release_tx(tx);
-            self.relational_db.report(&reducer, &tx_metrics, None);
+            self.relational_db.report(reducer, tx_metrics, None);
         });
         let auth = AuthCtx::new(self.owner_identity, sender.id.identity);
         let (table_rows, metrics) = return_on_err_with_sql!(
@@ -482,7 +482,7 @@ impl ModuleSubscriptions {
         // Always lock the db before the subscription lock to avoid deadlocks.
         let tx = scopeguard::guard(self.relational_db.begin_tx(Workload::Unsubscribe), |tx| {
             let (tx_metrics, reducer) = self.relational_db.release_tx(tx);
-            self.relational_db.report(&reducer, &tx_metrics, None);
+            self.relational_db.report(reducer, tx_metrics, None);
         });
 
         let removed_queries = {
@@ -578,7 +578,7 @@ impl ModuleSubscriptions {
         // We always get the db lock before the subscription lock to avoid deadlocks.
         let tx = scopeguard::guard(self.relational_db.begin_tx(Workload::Subscribe), |tx| {
             let (tx_metrics, reducer) = self.relational_db.release_tx(tx);
-            self.relational_db.report(&reducer, &tx_metrics, None);
+            self.relational_db.report(reducer, tx_metrics, None);
         });
 
         let compile_timer = metrics.compilation_time.start_timer();
@@ -680,7 +680,7 @@ impl ModuleSubscriptions {
         );
         let tx = scopeguard::guard(tx, |tx| {
             let (tx_metrics, reducer) = self.relational_db.release_tx(tx);
-            self.relational_db.report(&reducer, &tx_metrics, None);
+            self.relational_db.report(reducer, tx_metrics, None);
         });
 
         // We minimize locking so that other clients can add subscriptions concurrently.
@@ -774,7 +774,7 @@ impl ModuleSubscriptions {
         )?;
         let tx = scopeguard::guard(tx, |tx| {
             let (tx_metrics, reducer) = self.relational_db.release_tx(tx);
-            self.relational_db.report(&reducer, &tx_metrics, None);
+            self.relational_db.report(reducer, tx_metrics, None);
         });
 
         check_row_limit(
@@ -884,13 +884,17 @@ impl ModuleSubscriptions {
             }
         };
 
+        let tx_data = Arc::new(tx_data);
+
         // When we're done with this method, release the tx and report metrics.
         let mut read_tx = scopeguard::guard(read_tx, |tx| {
             let (tx_metrics_read, reducer) = self.relational_db.release_tx(tx);
-            stdb.report_tx_metricses(&reducer, tx_data.as_ref(), Some(&tx_metrics_mut), &tx_metrics_read);
+            self.relational_db
+                .report_tx_metrics(reducer, tx_data.clone(), Some(tx_metrics_mut), Some(tx_metrics_read));
         });
         // Create the delta transaction we'll use to eval updates against.
         let delta_read_tx = tx_data
+            .as_ref()
             .as_ref()
             .map(|tx_data| DeltaTx::new(&read_tx, tx_data, subscriptions.index_ids_for_subscriptions()))
             .unwrap_or_else(|| DeltaTx::from(&*read_tx));

crates/core/src/subscription/module_subscription_manager.rs

@@ -5,13 +5,18 @@ use crate::client::messages::{
     TransactionUpdateMessage,
 };
 use crate::client::{ClientConnectionSender, Protocol};
+use crate::db::datastore::locking_tx_datastore::datastore::TxMetrics;
 use crate::db::datastore::locking_tx_datastore::state_view::StateView;
+use crate::db::datastore::traits::TxData;
 use crate::error::DBError;
+use crate::execution_context::WorkloadType;
 use crate::host::module_host::{DatabaseTableUpdate, ModuleEvent, UpdatesRelValue};
 use crate::messages::websocket::{self as ws, TableUpdate};
 use crate::subscription::delta::eval_delta;
+use crate::subscription::ExecutionCounters;
 use crate::worker_metrics::WORKER_METRICS;
 use core::mem;
+use enum_map::EnumMap;
 use hashbrown::hash_map::OccupiedError;
 use hashbrown::{HashMap, HashSet};
 use parking_lot::RwLock;
@@ -26,6 +31,7 @@ use spacetimedb_lib::{AlgebraicValue, ConnectionId, Identity, ProductValue};
 use spacetimedb_primitives::{ColId, IndexId, TableId};
 use spacetimedb_subscription::{JoinEdge, SubscriptionPlan, TableName};
 use std::collections::BTreeMap;
+use std::fmt::Debug;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use tokio::sync::mpsc;
@@ -1283,6 +1289,62 @@ impl SubscriptionManager {
     }
 }
 
+struct MetricsMessage {
+    reducer: String,
+    metrics_for_writer: Option<TxMetrics>,
+    metrics_for_reader: Option<TxMetrics>,
+    tx_data: Arc<Option<TxData>>,
+    counters: Arc<EnumMap<WorkloadType, ExecutionCounters>>,
+}
+
+#[derive(Clone)]
+pub struct MetricsRecorderQueue {
+    tx: mpsc::UnboundedSender<MetricsMessage>,
+}
+
+impl MetricsRecorderQueue {
+    pub fn send_metrics(
+        &self,
+        reducer: String,
+        metrics_for_writer: Option<TxMetrics>,
+        metrics_for_reader: Option<TxMetrics>,
+        tx_data: Arc<Option<TxData>>,
+        counters: Arc<EnumMap<WorkloadType, ExecutionCounters>>,
+    ) {
+        if let Err(err) = self.tx.send(MetricsMessage {
+            reducer,
+            metrics_for_writer,
+            metrics_for_reader,
+            tx_data,
+            counters,
+        }) {
+            log::warn!("failed to send metrics: {err}");
+        }
+    }
+}
+
+pub fn spawn_metrics_recorder() -> MetricsRecorderQueue {
+    let (tx, mut rx) = mpsc::unbounded_channel();
+    tokio::spawn(async move {
+        while let Some(MetricsMessage {
+            reducer,
+            metrics_for_writer,
+            metrics_for_reader,
+            tx_data,
+            counters,
+        }) = rx.recv().await
+        {
+            if let Some(tx_metrics) = metrics_for_writer {
+                tx_metrics.report(tx_data.as_ref().as_ref(), &reducer, |wl| &counters[wl]);
+            }
+            if let Some(tx_metrics) = metrics_for_reader {
+                tx_metrics.report(tx_data.as_ref().as_ref(), &reducer, |wl| &counters[wl]);
+            }
+        }
+    });
+    MetricsRecorderQueue { tx }
+}
+
 struct SendWorkerClient {
     /// This flag is set if an error occurs during a tx update.
     /// It will be cleaned up async or on resubscribe.