fix(quaint): forward subscriber and recorder to spawned tasks (#5033)

* Move query-engine-metrics to libs

* Use prisma_metrics in quaint

* Merge metrics-guards into prisma-metrics

* Encapsulate metrics inside prisma-metrics

* Forward subscriber and recorder in quaint tasks

* [integration]

* rustfmt

* [integration]

* Forward subscriber and recorder to mobc

* Remove the global recorder fallback

* [integration]

Cargo.toml

@@ -20,7 +20,6 @@ members = [
   "query-engine/black-box-tests",
   "query-engine/dmmf",
   "query-engine/driver-adapters",
-  "query-engine/metrics",
   "query-engine/query-structure",
   "query-engine/query-engine",
   "query-engine/query-engine-node-api",
@@ -65,14 +64,14 @@ napi = { version = "2.15.1", default-features = false, features = [
   "serde-json",
 ] }
 napi-derive = "2.15.0"
-metrics = "0.23.0"
 js-sys = { version = "0.3" }
 pin-project = "1"
 rand = { version = "0.8" }
 regex = { version = "1", features = ["std"] }
 serde_repr = { version = "0.1.17" }
 serde-wasm-bindgen = { version = "0.5" }
 tracing = { version = "0.1" }
+tracing-futures = "0.2"
 tsify = { version = "0.4.5" }
 wasm-bindgen = { version = "0.2.93" }
 wasm-bindgen-futures = { version = "0.4" }

query-engine/metrics/Cargo.toml → libs/metrics/Cargo.toml

@@ -1,19 +1,19 @@
 [package]
-name = "query-engine-metrics"
+name = "prisma-metrics"
 version = "0.1.0"
 edition = "2021"
 
 [dependencies]
 futures.workspace = true
 derive_more.workspace = true
-metrics.workspace = true
+metrics = "0.23.0"
 metrics-util = "0.17.0"
 metrics-exporter-prometheus = { version = "0.15.3", default-features = false }
 once_cell = "1.3"
 serde.workspace = true
 serde_json.workspace = true
 tracing.workspace = true
-tracing-futures = "0.2"
+tracing-futures.workspace = true
 tracing-subscriber = "0.3.11"
 parking_lot = "0.12"
 pin-project.workspace = true

libs/metrics-guards/src/lib.rs → libs/metrics/src/guards.rs

@@ -1,6 +1,6 @@
 use std::sync::atomic::{AtomicBool, Ordering};
 
-use metrics::gauge;
+use crate::gauge;
 
 pub struct GaugeGuard {
     name: &'static str,

query-engine/metrics/src/lib.rs → libs/metrics/src/lib.rs

@@ -1,5 +1,8 @@
-//! Query Engine Metrics
-//! This crate is responsible for capturing and recording metrics in the Query Engine.
+//! # Prisma Metrics
+//!
+//! This crate is responsible for capturing and recording metrics in the Query Engine and its
+//! dependencies.
+//!
 //! Metrics is broken into two parts, `MetricsRecorder` and `MetricsRegistry`, and uses our tracing framework to communicate.
 //! An example best explains this system.
 //! When the engine boots up, the `MetricRegistry` is added to our tracing as a layer and The MetricRecorder is
@@ -25,6 +28,8 @@ mod instrument;
 mod recorder;
 mod registry;
 
+pub mod guards;
+
 use once_cell::sync::Lazy;
 use serde::Deserialize;
 use std::collections::HashMap;

libs/query-engine-common/Cargo.toml

@@ -16,12 +16,12 @@ psl.workspace = true
 async-trait.workspace = true
 tracing.workspace = true
 tracing-subscriber = { version = "0.3" }
-tracing-futures = "0.2"
+tracing-futures.workspace = true
 tracing-opentelemetry = "0.17.3"
 opentelemetry = { version = "0.17" }
 
 [target.'cfg(all(not(target_arch = "wasm32")))'.dependencies]
-query-engine-metrics = { path = "../../query-engine/metrics" }
+prisma-metrics.path = "../metrics"
 napi.workspace = true
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]

libs/query-engine-common/src/engine.rs

@@ -59,7 +59,7 @@ pub struct EngineBuilder {
 pub struct ConnectedEngineNative {
     pub config_dir: PathBuf,
     pub env: HashMap<String, String>,
-    pub metrics: Option<query_engine_metrics::MetricRegistry>,
+    pub metrics: Option<prisma_metrics::MetricRegistry>,
 }
 
 /// Internal structure for querying and reconnecting with the engine.

quaint/Cargo.toml

@@ -72,12 +72,11 @@ fmt-sql = ["sqlformat"]
 connection-string = "0.2"
 percent-encoding = "2"
 tracing.workspace = true
-tracing-core = "0.1"
+tracing-futures.workspace = true
 async-trait.workspace = true
 thiserror = "1.0"
 num_cpus = "1.12"
-metrics.workspace = true
-metrics-guards.path = "../libs/metrics-guards"
+prisma-metrics.path = "../libs/metrics"
 futures.workspace = true
 url.workspace = true
 hex = "0.4"

quaint/src/connector/metrics.rs

@@ -1,3 +1,4 @@
+use prisma_metrics::{counter, histogram};
 use tracing::{info_span, Instrument};
 
 use crate::ast::{Params, Value};

quaint/src/connector/postgres/native/mod.rs

@@ -28,6 +28,7 @@ use lru_cache::LruCache;
 use native_tls::{Certificate, Identity, TlsConnector};
 use postgres_native_tls::MakeTlsConnector;
 use postgres_types::{Kind as PostgresKind, Type as PostgresType};
+use prisma_metrics::WithMetricsInstrumentation;
 use std::hash::{DefaultHasher, Hash, Hasher};
 use std::{
     fmt::{Debug, Display},
@@ -38,6 +39,7 @@ use std::{
 };
 use tokio::sync::OnceCell;
 use tokio_postgres::{config::ChannelBinding, Client, Config, Statement};
+use tracing_futures::WithSubscriber;
 use websocket::connect_via_websocket;
 
 /// The underlying postgres driver. Only available with the `expose-drivers`
@@ -245,12 +247,15 @@ impl PostgreSql {
         let is_cockroachdb = conn.parameter("crdb_version").is_some();
         let is_materialize = conn.parameter("mz_version").is_some();
 
-        tokio::spawn(conn.map(|r| match r {
-            Ok(_) => (),
-            Err(e) => {
-                tracing::error!("Error in PostgreSQL connection: {:?}", e);
-            }
-        }));
+        tokio::spawn(
+            conn.map(|r| {
+                if let Err(e) = r {
+                    tracing::error!("Error in PostgreSQL connection: {e:?}");
+                }
+            })
+            .with_current_subscriber()
+            .with_current_recorder(),
+        );
 
         // On Postgres, we set the SEARCH_PATH and client-encoding through client connection parameters to save a network roundtrip on connection.
         // We can't always do it for CockroachDB because it does not expect quotes for unsafe identifiers (https://github.com/cockroachdb/cockroach/issues/101328), which might change once the issue is fixed.