feat: add DatabaseMetrics trait for generic DB in reth node (#5690)

bin/reth/src/node/mod.rs

@@ -38,7 +38,7 @@ use reth_config::{
     config::{PruneConfig, StageConfig},
     Config,
 };
-use reth_db::{database::Database, init_db, DatabaseEnv};
+use reth_db::{database::Database, database_metrics::DatabaseMetrics, init_db};
 use reth_downloaders::{
     bodies::bodies::BodiesDownloaderBuilder,
     headers::reverse_headers::ReverseHeadersDownloaderBuilder,
@@ -693,11 +693,14 @@ impl<Ext: RethCliExt> NodeCommand<Ext> {
         prometheus_exporter::install_recorder()
     }
 
-    async fn start_metrics_endpoint(
+    async fn start_metrics_endpoint<Metrics>(
         &self,
         prometheus_handle: PrometheusHandle,
-        db: Arc<DatabaseEnv>,
-    ) -> eyre::Result<()> {
+        db: Metrics,
+    ) -> eyre::Result<()>
+    where
+        Metrics: DatabaseMetrics + 'static + Send + Sync,
+    {
         if let Some(listen_addr) = self.metrics {
             info!(target: "reth::cli", addr = %listen_addr, "Starting metrics endpoint");
             prometheus_exporter::serve(

bin/reth/src/prometheus_exporter.rs

@@ -7,7 +7,7 @@ use hyper::{
 use metrics::{describe_gauge, gauge};
 use metrics_exporter_prometheus::{PrometheusBuilder, PrometheusHandle};
 use metrics_util::layers::{PrefixLayer, Stack};
-use reth_db::{database::Database, tables, DatabaseEnv};
+use reth_db::database_metrics::DatabaseMetrics;
 use reth_metrics::metrics::Unit;
 use std::{convert::Infallible, net::SocketAddr, sync::Arc};
 use tracing::error;
@@ -74,54 +74,21 @@ async fn start_endpoint<F: Hook + 'static>(
 }
 
 /// Serves Prometheus metrics over HTTP with database and process metrics.
-pub(crate) async fn serve(
+pub(crate) async fn serve<Metrics>(
     listen_addr: SocketAddr,
     handle: PrometheusHandle,
-    db: Arc<DatabaseEnv>,
+    db: Metrics,
     process: metrics_process::Collector,
-) -> eyre::Result<()> {
-    let db_stats = move || {
-        // TODO: A generic stats abstraction for other DB types to deduplicate this and `reth db
-        //  stats`
-        let _ = db.view(|tx| {
-            for table in tables::Tables::ALL.iter().map(|table| table.name()) {
-                let table_db =
-                    tx.inner.open_db(Some(table)).wrap_err("Could not open db.")?;
-
-                let stats = tx
-                    .inner
-                    .db_stat(&table_db)
-                    .wrap_err(format!("Could not find table: {table}"))?;
-
-                let page_size = stats.page_size() as usize;
-                let leaf_pages = stats.leaf_pages();
-                let branch_pages = stats.branch_pages();
-                let overflow_pages = stats.overflow_pages();
-                let num_pages = leaf_pages + branch_pages + overflow_pages;
-                let table_size = page_size * num_pages;
-                let entries = stats.entries();
-
-                gauge!("db.table_size", table_size as f64, "table" => table);
-                gauge!("db.table_pages", leaf_pages as f64, "table" => table, "type" => "leaf");
-                gauge!("db.table_pages", branch_pages as f64, "table" => table, "type" => "branch");
-                gauge!("db.table_pages", overflow_pages as f64, "table" => table, "type" => "overflow");
-                gauge!("db.table_entries", entries as f64, "table" => table);
-            }
-
-            Ok::<(), eyre::Report>(())
-        }).map_err(|error| error!(?error, "Failed to read db table stats"));
-
-        if let Ok(freelist) =
-            db.freelist().map_err(|error| error!(?error, "Failed to read db.freelist"))
-        {
-            gauge!("db.freelist", freelist as f64);
-        }
-    };
+) -> eyre::Result<()>
+where
+    Metrics: DatabaseMetrics + 'static + Send + Sync,
+{
+    let db_metrics_hook = move || db.report_metrics();
 
     // Clone `process` to move it into the hook and use the original `process` for describe below.
     let cloned_process = process.clone();
     let hooks: Vec<Box<dyn Hook<Output = ()>>> = vec![
-        Box::new(db_stats),
+        Box::new(db_metrics_hook),
         Box::new(move || cloned_process.collect()),
         Box::new(collect_memory_stats),
     ];

crates/storage/db/src/abstraction/database_metrics.rs

@@ -0,0 +1,14 @@
+use std::sync::Arc;
+
+/// Represents a type that can report metrics, used mainly with the database. The `report_metrics`
+/// method can be used as a prometheus hook.
+pub trait DatabaseMetrics {
+    /// Reports metrics for the database.
+    fn report_metrics(&self);
+}
+
+impl<DB: DatabaseMetrics> DatabaseMetrics for Arc<DB> {
+    fn report_metrics(&self) {
+        <DB as DatabaseMetrics>::report_metrics(self)
+    }
+}

crates/storage/db/src/abstraction/mod.rs

@@ -4,6 +4,8 @@ pub mod common;
 pub mod cursor;
 /// Database traits.
 pub mod database;
+/// Database metrics trait extensions.
+pub mod database_metrics;
 /// mock
 pub mod mock;
 /// Table traits

crates/storage/db/src/implementation/mdbx/mod.rs

@@ -2,14 +2,18 @@
 
 use crate::{
     database::Database,
+    database_metrics::DatabaseMetrics,
     tables::{TableType, Tables},
     utils::default_page_size,
     DatabaseError,
 };
+use eyre::Context;
+use metrics::gauge;
 use reth_interfaces::db::LogLevel;
 use reth_libmdbx::{
     DatabaseFlags, Environment, EnvironmentFlags, Geometry, Mode, PageSize, SyncMode, RO, RW,
 };
+use reth_tracing::tracing::error;
 use std::{ops::Deref, path::Path};
 use tx::Tx;
 
@@ -59,6 +63,44 @@ impl Database for DatabaseEnv {
     }
 }
 
+impl DatabaseMetrics for DatabaseEnv {
+    fn report_metrics(&self) {
+        let _ = self.view(|tx| {
+            for table in Tables::ALL.iter().map(|table| table.name()) {
+                let table_db =
+                    tx.inner.open_db(Some(table)).wrap_err("Could not open db.")?;
+
+                let stats = tx
+                    .inner
+                    .db_stat(&table_db)
+                    .wrap_err(format!("Could not find table: {table}"))?;
+
+                let page_size = stats.page_size() as usize;
+                let leaf_pages = stats.leaf_pages();
+                let branch_pages = stats.branch_pages();
+                let overflow_pages = stats.overflow_pages();
+                let num_pages = leaf_pages + branch_pages + overflow_pages;
+                let table_size = page_size * num_pages;
+                let entries = stats.entries();
+
+                gauge!("db.table_size", table_size as f64, "table" => table);
+                gauge!("db.table_pages", leaf_pages as f64, "table" => table, "type" => "leaf");
+                gauge!("db.table_pages", branch_pages as f64, "table" => table, "type" => "branch");
+                gauge!("db.table_pages", overflow_pages as f64, "table" => table, "type" => "overflow");
+                gauge!("db.table_entries", entries as f64, "table" => table);
+            }
+
+            Ok::<(), eyre::Report>(())
+        }).map_err(|error| error!(?error, "Failed to read db table stats"));
+
+        if let Ok(freelist) =
+            self.freelist().map_err(|error| error!(?error, "Failed to read db.freelist"))
+        {
+            gauge!("db.freelist", freelist as f64);
+        }
+    }
+}
+
 impl DatabaseEnv {
     /// Opens the database at the specified path with the given `EnvKind`.
     ///

crates/storage/db/src/lib.rs

@@ -153,7 +153,7 @@ pub fn open_db(path: &Path, log_level: Option<LogLevel>) -> eyre::Result<Databas
 #[cfg(any(test, feature = "test-utils"))]
 pub mod test_utils {
     use super::*;
-    use crate::database::Database;
+    use crate::{database::Database, database_metrics::DatabaseMetrics};
     use std::{path::PathBuf, sync::Arc};
 
     /// Error during database open
@@ -210,6 +210,12 @@ pub mod test_utils {
         }
     }
 
+    impl<DB: DatabaseMetrics> DatabaseMetrics for TempDatabase<DB> {
+        fn report_metrics(&self) {
+            self.db().report_metrics()
+        }
+    }
+
     /// Create read/write database for testing
     pub fn create_test_rw_db() -> Arc<TempDatabase<DatabaseEnv>> {
         let path = tempfile::TempDir::new().expect(ERROR_TEMPDIR).into_path();