fix(metrics): ensure consistent metrics for one engine (#5028)

libs/metrics-guards/Cargo.toml

@@ -0,0 +1,7 @@
+[package]
+name = "metrics-guards"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+metrics.workspace = true

libs/metrics-guards/src/lib.rs

@@ -0,0 +1,31 @@
+use std::sync::atomic::{AtomicBool, Ordering};
+
+use metrics::gauge;
+
+pub struct GaugeGuard {
+    name: &'static str,
+    decremented: AtomicBool,
+}
+
+impl GaugeGuard {
+    pub fn increment(name: &'static str) -> Self {
+        gauge!(name).increment(1.0);
+
+        Self {
+            name,
+            decremented: AtomicBool::new(false),
+        }
+    }
+
+    pub fn decrement(&self) {
+        if !self.decremented.swap(true, Ordering::Relaxed) {
+            gauge!(self.name).decrement(1.0);
+        }
+    }
+}
+
+impl Drop for GaugeGuard {
+    fn drop(&mut self) {
+        self.decrement();
+    }
+}

quaint/Cargo.toml

@@ -77,6 +77,7 @@ async-trait.workspace = true
 thiserror = "1.0"
 num_cpus = "1.12"
 metrics.workspace = true
+metrics-guards.path = "../libs/metrics-guards"
 futures.workspace = true
 url.workspace = true
 hex = "0.4"

quaint/src/connector/transaction.rs

@@ -1,13 +1,13 @@
+use std::{fmt, str::FromStr};
+
+use async_trait::async_trait;
+use metrics_guards::GaugeGuard;
+
 use super::*;
 use crate::{
     ast::*,
     error::{Error, ErrorKind},
 };
-use async_trait::async_trait;
-use metrics::gauge;
-use std::{fmt, str::FromStr};
-
-extern crate metrics as metrics;
 
 #[async_trait]
 pub trait Transaction: Queryable {
@@ -36,6 +36,7 @@ pub(crate) struct TransactionOptions {
 /// transaction object will panic.
 pub struct DefaultTransaction<'a> {
     pub inner: &'a dyn Queryable,
+    gauge: GaugeGuard,
 }
 
 impl<'a> DefaultTransaction<'a> {
@@ -44,7 +45,10 @@ impl<'a> DefaultTransaction<'a> {
         begin_stmt: &str,
         tx_opts: TransactionOptions,
     ) -> crate::Result<DefaultTransaction<'a>> {
-        let this = Self { inner };
+        let this = Self {
+            inner,
+            gauge: GaugeGuard::increment("prisma_client_queries_active"),
+        };
 
         if tx_opts.isolation_first {
             if let Some(isolation) = tx_opts.isolation_level {
@@ -62,7 +66,6 @@ impl<'a> DefaultTransaction<'a> {
 
         inner.server_reset_query(&this).await?;
 
-        gauge!("prisma_client_queries_active").increment(1.0);
         Ok(this)
     }
 }
@@ -71,15 +74,15 @@ impl<'a> DefaultTransaction<'a> {
 impl<'a> Transaction for DefaultTransaction<'a> {
     /// Commit the changes to the database and consume the transaction.
     async fn commit(&self) -> crate::Result<()> {
-        gauge!("prisma_client_queries_active").decrement(1.0);
+        self.gauge.decrement();
         self.inner.raw_cmd("COMMIT").await?;
 
         Ok(())
     }
 
     /// Rolls back the changes to the database.
     async fn rollback(&self) -> crate::Result<()> {
-        gauge!("prisma_client_queries_active").decrement(1.0);
+        self.gauge.decrement();
         self.inner.raw_cmd("ROLLBACK").await?;
 
         Ok(())

query-engine/query-engine-node-api/src/logger.rs

@@ -67,6 +67,15 @@ impl Logger {
         let (metrics, recorder) = if enable_metrics {
             let registry = MetricRegistry::new();
             let recorder = MetricRecorder::new(registry.clone()).with_initialized_prisma_metrics();
+
+            // FIXME: we attempt to install the recorder globally because some of the mobc metrics
+            // are being modified outside of the async context with active local metric recorder,
+            // causing them to be lost. This workaround ensures we have a global fallback for that
+            // case, but installing the global recorder will only work for the first engine
+            // instance, so the metrics will still be inconsistent if there are multiple
+            // PrismaClient instances in the app. We need to fix this to be able to GA metrics.
+            _ = recorder.install_globally();
+
             (Some(registry), Some(recorder))
         } else {
             (None, None)