make compatible to private

Author: Shubham8287

crates/core/src/host/host_controller.rs

@@ -904,7 +904,12 @@ impl Host {
             db::Storage::Disk => {
                 // Replay from the local state.
                 let history = relational_db::local_history(&replica_dir, &runtime).await?;
-                let persistence = persistence.persistence(database.database_identity, replica_id).await?;
+                let persistence_db = db::persistence::Database {
+                    id: database.id,
+                    database_identity: database.database_identity,
+                    owner_identity: database.owner_identity,
+                };
+                let persistence = persistence.persistence(&persistence_db, replica_id).await?;
                 // Loading a database from persistent storage involves heavy
                 // blocking I/O. `asyncify` to avoid blocking the async worker.
                 let (db, clients) = asyncify({

crates/core/src/subscription/module_subscription_actor.rs

@@ -1744,6 +1744,12 @@ impl ModuleSubscriptions {
                 }
             }
         }
+        if update_metrics.bytes_sent_to_clients > 0 {
+            WORKER_METRICS
+                .bytes_sent_to_clients
+                .with_label_values(&WorkloadType::Subscribe, &self.relational_db.database_identity())
+                .inc_by(update_metrics.bytes_sent_to_clients as u64);
+        }
         read_tx.metrics.merge(update_metrics);
         Ok(Ok(CommitAndBroadcastEventSuccess {
             tx_offset: extra_tx_offset,

crates/core/src/worker_metrics/mod.rs

@@ -43,6 +43,11 @@ metrics_group!(
         #[labels(database_identity: Identity, protocol: str)]
         pub websocket_request_msg_size: HistogramVec,
 
+        #[name = spacetime_num_bytes_sent_to_clients_total]
+        #[help = "The cumulative number of bytes sent to clients"]
+        #[labels(workload: WorkloadType, db: Identity)]
+        pub bytes_sent_to_clients: IntCounterVec,
+
         #[name = jemalloc_active_bytes]
         #[help = "Number of bytes in jemallocs heap"]
         #[labels(node_id: str)]

crates/engine/src/db/persistence.rs

@@ -85,6 +85,13 @@ pub type Durability = dyn spacetimedb_durability::Durability<TxData = Txdata>;
 /// configured or the database is in follower state.
 pub type DiskSizeFn = Arc<dyn Fn() -> io::Result<SizeOnDisk> + Send + Sync>;
 
+#[derive(Clone, Copy, Debug)]
+pub struct Database {
+    pub id: u64,
+    pub database_identity: Identity,
+    pub owner_identity: Identity,
+}
+
 /// Persistence services for a database.
 pub struct Persistence {
     /// The [Durability] to use, for persisting transactions.
@@ -182,7 +189,7 @@ impl Persistence {
 /// This is an `async_trait` to allow it to be used as a trait object.
 #[async_trait]
 pub trait PersistenceProvider: Send + Sync {
-    async fn persistence(&self, database_identity: Identity, replica_id: u64) -> anyhow::Result<Persistence>;
+    async fn persistence(&self, database: &Database, replica_id: u64) -> anyhow::Result<Persistence>;
 }
 
 /// The standard [PersistenceProvider] for non-replicated databases.
@@ -215,7 +222,8 @@ impl LocalPersistenceProvider {
 
 #[async_trait]
 impl PersistenceProvider for LocalPersistenceProvider {
-    async fn persistence(&self, database_identity: Identity, replica_id: u64) -> anyhow::Result<Persistence> {
+    async fn persistence(&self, database: &Database, replica_id: u64) -> anyhow::Result<Persistence> {
+        let database_identity = database.database_identity;
         let replica_dir = self.data_dir.replica(replica_id);
         let snapshot_dir = replica_dir.snapshots();
         let runtime = Handle::tokio_current();

crates/engine/src/metrics.rs

@@ -1,5 +1,5 @@
 use once_cell::sync::Lazy;
-use prometheus::{GaugeVec, HistogramVec, IntCounter, IntCounterVec, IntGaugeVec};
+use prometheus::{GaugeVec, HistogramVec, IntCounter, IntGaugeVec};
 use spacetimedb_datastore::{
     db_metrics::DB_METRICS, execution_context::WorkloadType, locking_tx_datastore::datastore::MetricsRecorder,
 };
@@ -8,11 +8,6 @@ use spacetimedb_metrics::metrics_group;
 
 metrics_group!(
     pub struct EngineMetrics {
-        #[name = spacetime_num_bytes_sent_to_clients_total]
-        #[help = "The cumulative number of bytes sent to clients"]
-        #[labels(txn_type: WorkloadType, db: Identity)]
-        pub bytes_sent_to_clients: IntCounterVec,
-
         #[name = spacetime_replay_total_time_seconds]
         #[help = "Total time spent replaying a database upon restart, including snapshot read, snapshot restore and commitlog replay"]
         #[labels(db: Identity)]
@@ -121,7 +116,6 @@ pub struct ExecutionCounters {
     rdb_num_rows_scanned: IntCounter,
     rdb_num_bytes_scanned: IntCounter,
     rdb_num_bytes_written: IntCounter,
-    bytes_sent_to_clients: IntCounter,
     delta_queries_matched: IntCounter,
     delta_queries_evaluated: IntCounter,
     duplicate_rows_evaluated: IntCounter,
@@ -135,7 +129,6 @@ impl ExecutionCounters {
             rdb_num_rows_scanned: DB_METRICS.rdb_num_rows_scanned.with_label_values(workload, db),
             rdb_num_bytes_scanned: DB_METRICS.rdb_num_bytes_scanned.with_label_values(workload, db),
             rdb_num_bytes_written: DB_METRICS.rdb_num_bytes_written.with_label_values(workload, db),
-            bytes_sent_to_clients: ENGINE_METRICS.bytes_sent_to_clients.with_label_values(workload, db),
             delta_queries_matched: DB_METRICS.delta_queries_matched.with_label_values(db),
             delta_queries_evaluated: DB_METRICS.delta_queries_evaluated.with_label_values(db),
             duplicate_rows_evaluated: DB_METRICS.duplicate_rows_evaluated.with_label_values(db),
@@ -156,9 +149,6 @@ impl ExecutionCounters {
         if metrics.bytes_written > 0 {
             self.rdb_num_bytes_written.inc_by(metrics.bytes_written as u64);
         }
-        if metrics.bytes_sent_to_clients > 0 {
-            self.bytes_sent_to_clients.inc_by(metrics.bytes_sent_to_clients as u64);
-        }
         if metrics.delta_queries_matched > 0 {
             self.delta_queries_matched.inc_by(metrics.delta_queries_matched);
         }