Store SubscriptionMetrics for Update, Subscribe, Unsubscribe in ModuleSubscriptions (#3821)

# Description of Changes

This showed up in flamegraphs, in particular for `Update`, so let's
cache the `SubscriptionMetrics`s.

# API and ABI breaking changes

None

# Expected complexity level and risk

1

# Testing

Covered by existing tests.

Author: Centril

crates/core/src/subscription/module_subscription_actor.rs

@@ -42,6 +42,7 @@ use spacetimedb_lib::identity::AuthCtx;
 use spacetimedb_lib::metrics::ExecutionMetrics;
 use spacetimedb_lib::Identity;
 use spacetimedb_primitives::ArgId;
+use spacetimedb_table::static_assert_size;
 use std::collections::HashSet;
 use std::{sync::Arc, time::Instant};
 use tokio::sync::oneshot;
@@ -56,10 +57,11 @@ pub struct ModuleSubscriptions {
     subscriptions: Subscriptions,
     broadcast_queue: BroadcastQueue,
     stats: Arc<SubscriptionGauges>,
+    metrics: Arc<SubscriptionMetricsForWorkloads>,
 }
 
 #[derive(Debug, Clone)]
-pub struct SubscriptionGauges {
+struct SubscriptionGauges {
     db_identity: Identity,
     num_queries: IntGauge,
     num_connections: IntGauge,
@@ -111,17 +113,36 @@ impl SubscriptionGauges {
     }
 }
 
-pub struct SubscriptionMetrics {
-    pub lock_waiters: IntGauge,
-    pub lock_wait_time: Histogram,
-    pub compilation_time: Histogram,
-    pub num_queries_subscribed: IntCounter,
-    pub num_new_queries_subscribed: IntCounter,
-    pub num_queries_evaluated: IntCounter,
+struct SubscriptionMetricsForWorkloads {
+    update: SubscriptionMetrics,
+    subscribe: SubscriptionMetrics,
+    unsubscribe: SubscriptionMetrics,
 }
 
+impl SubscriptionMetricsForWorkloads {
+    fn new(db: &Identity) -> Self {
+        Self {
+            update: SubscriptionMetrics::new(db, WorkloadType::Update),
+            subscribe: SubscriptionMetrics::new(db, WorkloadType::Subscribe),
+            unsubscribe: SubscriptionMetrics::new(db, WorkloadType::Unsubscribe),
+        }
+    }
+}
+
+struct SubscriptionMetrics {
+    lock_waiters: IntGauge,
+    lock_wait_time: Histogram,
+    compilation_time: Histogram,
+    num_queries_subscribed: IntCounter,
+    num_new_queries_subscribed: IntCounter,
+    num_queries_evaluated: IntCounter,
+}
+
+static_assert_size!(SubscriptionMetrics, 48);
+
 impl SubscriptionMetrics {
-    pub fn new(db: &Identity, workload: &WorkloadType) -> Self {
+    fn new(db: &Identity, workload: WorkloadType) -> Self {
+        let workload = &workload;
         Self {
             lock_waiters: DB_METRICS.subscription_lock_waiters.with_label_values(db, workload),
             lock_wait_time: DB_METRICS.subscription_lock_wait_time.with_label_values(db, workload),
@@ -218,12 +239,14 @@ impl ModuleSubscriptions {
     ) -> Self {
         let db = &relational_db.database_identity();
         let stats = Arc::new(SubscriptionGauges::new(db));
+        let metrics = Arc::new(SubscriptionMetricsForWorkloads::new(db));
 
         Self {
             relational_db,
             subscriptions,
             broadcast_queue,
             stats,
+            metrics,
         }
     }
 
@@ -579,8 +602,7 @@ impl ModuleSubscriptions {
             )
         };
 
-        let database_identity = self.relational_db.database_identity();
-        let subscription_metrics = SubscriptionMetrics::new(&database_identity, &WorkloadType::Unsubscribe);
+        let subscription_metrics = &self.metrics.unsubscribe;
 
         // Always lock the db before the subscription lock to avoid deadlocks.
         let (mut_tx, _) = self.begin_mut_tx(Workload::Unsubscribe);
@@ -780,12 +802,9 @@ impl ModuleSubscriptions {
             );
         };
 
-        let num_queries = request.query_strings.len();
-
-        let database_identity = self.relational_db.database_identity();
-        let subscription_metrics = SubscriptionMetrics::new(&database_identity, &WorkloadType::Subscribe);
-
         // How many queries make up this subscription?
+        let subscription_metrics = &self.metrics.subscribe;
+        let num_queries = request.query_strings.len();
         subscription_metrics.num_queries_subscribed.inc_by(num_queries as _);
 
         let (queries, auth, mut_tx, compile_timer) = return_on_err!(
@@ -794,7 +813,7 @@ impl ModuleSubscriptions {
                 auth,
                 &request.query_strings,
                 num_queries,
-                &subscription_metrics
+                subscription_metrics
             ),
             send_err_msg,
             None
@@ -885,19 +904,17 @@ impl ModuleSubscriptions {
         timer: Instant,
         _assert: Option<AssertTxFn>,
     ) -> Result<ExecutionMetrics, DBError> {
-        let num_queries = subscription.query_strings.len();
-        let database_identity = self.relational_db.database_identity();
-        let subscription_metrics = SubscriptionMetrics::new(&database_identity, &WorkloadType::Subscribe);
-
         // How many queries make up this subscription?
+        let subscription_metrics = &self.metrics.subscribe;
+        let num_queries = subscription.query_strings.len();
         subscription_metrics.num_queries_subscribed.inc_by(num_queries as _);
 
         let (queries, auth, mut_tx, compile_timer) = self.compile_queries(
             sender.id.identity,
             auth,
             &subscription.query_strings,
             num_queries,
-            &subscription_metrics,
+            subscription_metrics,
         )?;
 
         let (tx, tx_offset) = self
@@ -929,7 +946,7 @@ impl ModuleSubscriptions {
             )?,
         };
 
-        record_query_metrics(&database_identity, query_metrics);
+        record_query_metrics(&self.relational_db.database_identity(), query_metrics);
 
         // It acquires the subscription lock after `eval`, allowing `add_subscription` to run concurrently.
         // This also makes it possible for `broadcast_event` to get scheduled before the subsequent part here
@@ -994,8 +1011,7 @@ impl ModuleSubscriptions {
         mut event: ModuleEvent,
         tx: MutTx,
     ) -> Result<CommitAndBroadcastEventResult, DBError> {
-        let database_identity = self.relational_db.database_identity();
-        let subscription_metrics = SubscriptionMetrics::new(&database_identity, &WorkloadType::Update);
+        let subscription_metrics = &self.metrics.update;
 
         // Take a read lock on `subscriptions` before committing tx
         // else it can result in subscriber receiving duplicate updates.