Move the pool acquisition metric logic to the PgRepositoryFactory

Author: sandhose

Cargo.lock

@@ -4,7 +4,7 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Please see LICENSE in the repository root for full details.
 
-use std::{convert::Infallible, net::IpAddr, sync::Arc, time::Instant};
+use std::{convert::Infallible, net::IpAddr, sync::Arc};
 
 use axum::extract::{FromRef, FromRequestParts};
 use ipnetwork::IpNetwork;
@@ -19,10 +19,12 @@ use mas_keystore::{Encrypter, Keystore};
 use mas_matrix::HomeserverConnection;
 use mas_policy::{Policy, PolicyFactory};
 use mas_router::UrlBuilder;
-use mas_storage::{BoxClock, BoxRepository, BoxRepositoryFactory, BoxRng, SystemClock, RepositoryFactory};
+use mas_storage::{
+    BoxClock, BoxRepository, BoxRepositoryFactory, BoxRng, RepositoryFactory, SystemClock,
+};
 use mas_storage_pg::PgRepositoryFactory;
 use mas_templates::Templates;
-use opentelemetry::{KeyValue, metrics::Histogram};
+use opentelemetry::KeyValue;
 use rand::SeedableRng;
 use sqlx::PgPool;
 use tracing::Instrument;
@@ -47,7 +49,6 @@ pub struct AppState {
     pub activity_tracker: ActivityTracker,
     pub trusted_proxies: Vec<IpNetwork>,
     pub limiter: Limiter,
-    pub conn_acquisition_histogram: Option<Histogram<u64>>,
 }
 
 impl AppState {
@@ -76,14 +77,6 @@ impl AppState {
                 instrument.observe(i64::from(max_conn), &[]);
             })
             .build();
-
-        // Track the connection acquisition time
-        let histogram = METER
-            .u64_histogram("db.client.connections.create_time")
-            .with_description("The time it took to create a new connection.")
-            .with_unit("ms")
-            .build();
-        self.conn_acquisition_histogram = Some(histogram);
     }
 
     /// Init the metadata cache in the background
@@ -371,17 +364,7 @@ impl FromRequestParts<AppState> for BoxRepository {
         _parts: &mut axum::http::request::Parts,
         state: &AppState,
     ) -> Result<Self, Self::Rejection> {
-        let start = Instant::now();
         let repo = state.repository_factory.create().await?;
-
-        // Measure the time it took to create the connection
-        let duration = start.elapsed();
-        let duration_ms = duration.as_millis().try_into().unwrap_or(u64::MAX);
-
-        if let Some(histogram) = &state.conn_acquisition_histogram {
-            histogram.record(duration_ms, &[]);
-        }
-
         Ok(repo)
     }
 }

crates/cli/src/app_state.rs

@@ -18,7 +18,7 @@ use mas_handlers::{ActivityTracker, CookieManager, Limiter, MetadataCache};
 use mas_listener::server::Server;
 use mas_router::UrlBuilder;
 use mas_storage::SystemClock;
-use mas_storage_pg::{PgRepositoryFactory, MIGRATOR};
+use mas_storage_pg::{MIGRATOR, PgRepositoryFactory};
 use sqlx::migrate::Migrate;
 use tracing::{Instrument, info, info_span, warn};
 
@@ -242,7 +242,6 @@ impl Options {
                 activity_tracker,
                 trusted_proxies,
                 limiter,
-                conn_acquisition_histogram: None,
             };
             s.init_metrics();
             s.init_metadata_cache();

crates/cli/src/commands/server.rs

@@ -21,6 +21,7 @@ serde_json.workspace = true
 thiserror.workspace = true
 tracing.workspace = true
 futures-util.workspace = true
+opentelemetry.workspace = true
 opentelemetry-semantic-conventions.workspace = true
 
 rand.workspace = true

crates/storage-pg/Cargo.toml

@@ -175,6 +175,7 @@ pub(crate) mod iden;
 pub(crate) mod pagination;
 pub(crate) mod policy_data;
 pub(crate) mod repository;
+pub(crate) mod telemetry;
 pub(crate) mod tracing;
 
 pub(crate) use self::errors::DatabaseInconsistencyError;

crates/storage-pg/src/lib.rs

@@ -48,6 +48,7 @@ use crate::{
         job::PgQueueJobRepository, schedule::PgQueueScheduleRepository,
         worker::PgQueueWorkerRepository,
     },
+    telemetry::DB_CLIENT_CONNECTIONS_CREATE_TIME_HISTOGRAM,
     upstream_oauth2::{
         PgUpstreamOAuthLinkRepository, PgUpstreamOAuthProviderRepository,
         PgUpstreamOAuthSessionRepository,
@@ -89,10 +90,18 @@ impl PgRepositoryFactory {
 #[async_trait]
 impl RepositoryFactory for PgRepositoryFactory {
     async fn create(&self) -> Result<BoxRepository, RepositoryError> {
-        Ok(PgRepository::from_pool(&self.pool)
+        let start = std::time::Instant::now();
+        let repo = PgRepository::from_pool(&self.pool)
             .await
             .map_err(RepositoryError::from_error)?
-            .boxed())
+            .boxed();
+
+        // Measure the time it took to create the connection
+        let duration = start.elapsed();
+        let duration_ms = duration.as_millis().try_into().unwrap_or(u64::MAX);
+        DB_CLIENT_CONNECTIONS_CREATE_TIME_HISTOGRAM.record(duration_ms, &[]);
+
+        Ok(repo)
     }
 }
 

crates/storage-pg/src/repository.rs

@@ -0,0 +1,31 @@
+// Copyright 2025 New Vector Ltd.
+//
+// SPDX-License-Identifier: AGPL-3.0-only
+// Please see LICENSE in the repository root for full details.
+
+use std::sync::LazyLock;
+
+use opentelemetry::{
+    InstrumentationScope,
+    metrics::{Histogram, Meter},
+};
+use opentelemetry_semantic_conventions as semcov;
+
+static SCOPE: LazyLock<InstrumentationScope> = LazyLock::new(|| {
+    InstrumentationScope::builder(env!("CARGO_PKG_NAME"))
+        .with_version(env!("CARGO_PKG_VERSION"))
+        .with_schema_url(semcov::SCHEMA_URL)
+        .build()
+});
+
+static METER: LazyLock<Meter> =
+    LazyLock::new(|| opentelemetry::global::meter_with_scope(SCOPE.clone()));
+
+pub(crate) static DB_CLIENT_CONNECTIONS_CREATE_TIME_HISTOGRAM: LazyLock<Histogram<u64>> =
+    LazyLock::new(|| {
+        METER
+            .u64_histogram("db.client.connections.create_time")
+            .with_description("The time it took to create a new connection.")
+            .with_unit("ms")
+            .build()
+    });

crates/storage-pg/src/telemetry.rs

@@ -31,7 +31,8 @@ use crate::{
 };
 
 /// A [`RepositoryFactory`] is a factory that can create a [`BoxRepository`]
-// XXX(quenting): this could be generic over the repository type, but it's annoying to make it dyn-safe
+// XXX(quenting): this could be generic over the repository type, but it's annoying to make it
+// dyn-safe
 #[async_trait]
 pub trait RepositoryFactory {
     /// Create a new [`BoxRepository`]