prom metrics (#760)

* wip

* add prometheus metrics probes

* expose metrics admin endpoint

* move metric route to /metrics

* use histogram macro

Author: MarinPostma

Cargo.lock

@@ -71,6 +71,8 @@ rustls-pemfile = "1.0.3"
 rustls = "0.21.7"
 async-stream = "0.3.5"
 libsql = { git = "https://github.com/tursodatabase/libsql.git", rev = "bea8863", optional = true }
+metrics = "0.21.1"
+metrics-exporter-prometheus = "0.12.1"
 
 [dev-dependencies]
 proptest = "1.0.0"

sqld/Cargo.toml

@@ -3,6 +3,7 @@ use std::path::{Path, PathBuf};
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 
+use metrics::histogram;
 use parking_lot::{Mutex, RwLock};
 use rusqlite::{DatabaseName, ErrorCode, OpenFlags, StatementStatus};
 use sqld_libsql_bindings::wal_hook::{TransparentMethods, WalMethodsHook};
@@ -12,6 +13,7 @@ use tokio::time::{Duration, Instant};
 use crate::auth::{Authenticated, Authorized, Permission};
 use crate::error::Error;
 use crate::libsql_bindings::wal_hook::WalHook;
+use crate::metrics::{READ_QUERY_COUNT, WRITE_QUERY_COUNT};
 use crate::query::Query;
 use crate::query_analysis::{State, StmtKind};
 use crate::query_result_builder::{QueryBuilderConfig, QueryResultBuilder};
@@ -247,21 +249,34 @@ struct TxnSlot<T: WalHook> {
     /// is stolen.
     conn: Arc<Mutex<Connection<T>>>,
     /// Time at which the transaction can be stolen
-    timeout_at: tokio::time::Instant,
+    created_at: tokio::time::Instant,
     /// The transaction lock was stolen
     is_stolen: AtomicBool,
 }
 
+impl<T: WalHook> TxnSlot<T> {
+    #[inline]
+    fn expires_at(&self) -> Instant {
+        self.created_at + TXN_TIMEOUT
+    }
+
+    fn abort(&self) {
+        let conn = self.conn.lock();
+        // we have a lock on the connection, we don't need mode than a
+        // Relaxed store.
+        conn.rollback();
+        histogram!("write_txn_duration", self.created_at.elapsed())
+        // WRITE_TXN_DURATION.record(self.created_at.elapsed());
+    }
+}
+
 impl<T: WalHook> std::fmt::Debug for TxnSlot<T> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let stolen = self.is_stolen.load(Ordering::Relaxed);
-        let time_left = self
-            .timeout_at
-            .duration_since(tokio::time::Instant::now())
-            .as_millis();
+        let time_left = self.expires_at().duration_since(Instant::now());
         write!(
             f,
-            "(conn: {:?}, timeout_ms: {time_left}, stolen: {stolen})",
+            "(conn: {:?}, timeout: {time_left:?}, stolen: {stolen})",
             self.conn
         )
     }
@@ -315,7 +330,7 @@ unsafe extern "C" fn busy_handler<W: WalHook>(state: *mut c_void, _retries: c_in
     tokio::runtime::Handle::current().block_on(async move {
         let timeout = {
             let slot = lock.as_ref().unwrap();
-            let timeout_at = slot.timeout_at;
+            let timeout_at = slot.expires_at();
             drop(lock);
             tokio::time::sleep_until(timeout_at)
         };
@@ -337,11 +352,7 @@ unsafe extern "C" fn busy_handler<W: WalHook>(state: *mut c_void, _retries: c_in
                         // steal.
                         assert!(lock.take().is_some());
 
-                        let conn = slot.conn.lock();
-                        // we have a lock on the connection, we don't need mode than a
-                        // Relaxed store.
-                        conn.rollback();
-
+                        slot.abort();
                         tracing::info!("stole transaction lock");
                     }
                 }
@@ -422,7 +433,7 @@ impl<W: WalHook> Connection<W> {
             let mut lock = this.lock();
 
             if let Some(slot) = &lock.slot {
-                if slot.is_stolen.load(Ordering::Relaxed) || Instant::now() > slot.timeout_at {
+                if slot.is_stolen.load(Ordering::Relaxed) || Instant::now() > slot.expires_at() {
                     // we mark ourselves as stolen to notify any waiting lock thief.
                     slot.is_stolen.store(true, Ordering::Relaxed);
                     lock.rollback();
@@ -447,7 +458,7 @@ impl<W: WalHook> Connection<W> {
                 (Tx::None | Tx::Read, Tx::Write) => {
                     let slot = Arc::new(TxnSlot {
                         conn: this.clone(),
-                        timeout_at: Instant::now() + TXN_TIMEOUT,
+                        created_at: Instant::now(),
                         is_stolen: AtomicBool::new(false),
                     });
 
@@ -542,6 +553,11 @@ impl<W: WalHook> Connection<W> {
         }
 
         let mut stmt = self.conn.prepare(&query.stmt.stmt)?;
+        if stmt.readonly() {
+            READ_QUERY_COUNT.increment(1);
+        } else {
+            WRITE_QUERY_COUNT.increment(1);
+        }
 
         let cols = stmt.columns();
         let cols_count = cols.len();

sqld/src/connection/libsql.rs

@@ -1,12 +1,14 @@
+use metrics::histogram;
 use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
 use std::sync::Arc;
-use tokio::time::Duration;
+use tokio::time::{Duration, Instant};
 
 use futures::Future;
 use tokio::{sync::Semaphore, time::timeout};
 
 use crate::auth::Authenticated;
 use crate::error::Error;
+use crate::metrics::CONCCURENT_CONNECTIONS_COUNT;
 use crate::query::{Params, Query};
 use crate::query_analysis::{State, Statement};
 use crate::query_result_builder::{IgnoreResult, QueryResultBuilder};
@@ -249,6 +251,7 @@ impl<F: MakeConnection> MakeConnection for MakeThrottledConnection<F> {
     type Connection = TrackedConnection<F::Connection>;
 
     async fn create(&self) -> Result<Self::Connection, Error> {
+        let before_create = Instant::now();
         // If the memory pressure is high, request more units to reduce concurrency.
         tracing::trace!(
             "Available semaphore units: {}",
@@ -279,10 +282,16 @@ impl<F: MakeConnection> MakeConnection for MakeThrottledConnection<F> {
         }
 
         let inner = self.connection_maker.create().await?;
+
+        CONCCURENT_CONNECTIONS_COUNT.increment(1.0);
+        // CONNECTION_CREATE_TIME.record(before_create.elapsed());
+        histogram!("connection_create_time", before_create.elapsed());
+
         Ok(TrackedConnection {
             permit,
             inner,
             atime: AtomicU64::new(now_millis()),
+            created_at: Instant::now(),
         })
     }
 }
@@ -293,6 +302,15 @@ pub struct TrackedConnection<DB> {
     #[allow(dead_code)] // just hold on to it
     permit: tokio::sync::OwnedSemaphorePermit,
     atime: AtomicU64,
+    created_at: Instant,
+}
+
+impl<T> Drop for TrackedConnection<T> {
+    fn drop(&mut self) {
+        CONCCURENT_CONNECTIONS_COUNT.decrement(1.0);
+        histogram!("connection_create_time", self.created_at.elapsed());
+        // CONNECTION_ALIVE_DURATION.record();
+    }
 }
 
 impl<DB: Connection> TrackedConnection<DB> {

sqld/src/connection/mod.rs

@@ -12,6 +12,7 @@ use uuid::Uuid;
 
 use crate::auth::Authenticated;
 use crate::error::Error;
+use crate::metrics::REQUESTS_PROXIED;
 use crate::namespace::NamespaceName;
 use crate::query::Value;
 use crate::query_analysis::State;
@@ -197,6 +198,8 @@ impl WriteProxyConnection {
         auth: Authenticated,
         builder: B,
     ) -> Result<(B, State)> {
+        REQUESTS_PROXIED.increment(1);
+
         self.stats.inc_write_requests_delegated();
         let mut client = self.write_proxy.clone();
 

sqld/src/connection/write_proxy.rs

@@ -1,10 +1,11 @@
 use anyhow::Context as _;
-use axum::extract::{Path, State};
+use axum::extract::{FromRef, Path, State};
 use axum::routing::delete;
 use axum::Json;
 use chrono::NaiveDateTime;
 use futures::TryStreamExt;
 use hyper::Body;
+use metrics_exporter_prometheus::{PrometheusBuilder, PrometheusHandle};
 use serde::{Deserialize, Serialize};
 use std::io::ErrorKind;
 use std::path::PathBuf;
@@ -24,10 +25,24 @@ pub mod stats;
 type UserHttpServer<M> =
     Arc<hrana::http::Server<<<M as MakeNamespace>::Database as Database>::Connection>>;
 
+#[derive(Clone)]
+struct Metrics {
+    handle: PrometheusHandle,
+}
+
 struct AppState<M: MakeNamespace, C> {
     namespaces: NamespaceStore<M>,
     user_http_server: UserHttpServer<M>,
     connector: C,
+    metrics: PrometheusHandle,
+}
+
+impl<M: MakeNamespace, C> FromRef<Arc<AppState<M, C>>> for Metrics {
+    fn from_ref(input: &Arc<AppState<M, C>>) -> Self {
+        Metrics {
+            handle: input.metrics.clone(),
+        }
+    }
 }
 
 pub async fn run<M, A, C>(
@@ -59,10 +74,12 @@ where
         .route("/v1/namespaces/:namespace", delete(handle_delete_namespace))
         .route("/v1/namespaces/:namespace/stats", get(stats::handle_stats))
         .route("/v1/diagnostics", get(handle_diagnostics))
+        .route("/metrics", get(handle_metrics))
         .with_state(Arc::new(AppState {
             namespaces,
             connector,
             user_http_server,
+            metrics: PrometheusBuilder::new().install_recorder().unwrap(),
         }));
 
     hyper::server::Server::builder(acceptor)
@@ -76,6 +93,10 @@ async fn handle_get_index() -> &'static str {
     "Welcome to the sqld admin API"
 }
 
+async fn handle_metrics(State(metrics): State<Metrics>) -> String {
+    metrics.handle.render()
+}
+
 async fn handle_get_config<M: MakeNamespace, C: Connector>(
     State(app_state): State<Arc<AppState<M, C>>>,
     Path(namespace): Path<String>,

sqld/src/http/admin/mod.rs

@@ -55,6 +55,7 @@ mod h2c;
 mod heartbeat;
 mod hrana;
 mod http;
+mod metrics;
 mod migration;
 mod namespace;
 mod query;

sqld/src/lib.rs

@@ -0,0 +1,47 @@
+#![allow(dead_code)]
+use metrics::{
+    describe_counter, describe_gauge, describe_histogram, register_counter, register_gauge,
+    register_histogram, Counter, Gauge, Histogram,
+};
+use once_cell::sync::Lazy;
+
+pub static WRITE_QUERY_COUNT: Lazy<Counter> = Lazy::new(|| {
+    const NAME: &str = "writes_count";
+    describe_counter!(NAME, "number of write statements");
+    register_counter!(NAME)
+});
+pub static READ_QUERY_COUNT: Lazy<Counter> = Lazy::new(|| {
+    const NAME: &str = "read_count";
+    describe_counter!(NAME, "number of read statements");
+    register_counter!(NAME)
+});
+pub static REQUESTS_PROXIED: Lazy<Counter> = Lazy::new(|| {
+    const NAME: &str = "requests_proxied";
+    describe_counter!(NAME, "number of proxied requests");
+    register_counter!(NAME)
+});
+pub static CONCCURENT_CONNECTIONS_COUNT: Lazy<Gauge> = Lazy::new(|| {
+    const NAME: &str = "conccurent_connections";
+    describe_gauge!(NAME, "mumber of conccurent connections");
+    register_gauge!(NAME)
+});
+pub static NAMESPACE_LOAD_LATENCY: Lazy<Histogram> = Lazy::new(|| {
+    const NAME: &str = "namespace_load_latency";
+    describe_histogram!(NAME, "latency is us when loading a namespace");
+    register_histogram!(NAME)
+});
+pub static CONNECTION_CREATE_TIME: Lazy<Histogram> = Lazy::new(|| {
+    const NAME: &str = "connection_create_time";
+    describe_histogram!(NAME, "time to create a connection");
+    register_histogram!(NAME)
+});
+pub static CONNECTION_ALIVE_DURATION: Lazy<Histogram> = Lazy::new(|| {
+    const NAME: &str = "connection_alive_duration";
+    describe_histogram!(NAME, "duration for which a connection was kept alive");
+    register_histogram!(NAME)
+});
+pub static WRITE_TXN_DURATION: Lazy<Histogram> = Lazy::new(|| {
+    const NAME: &str = "write_txn_duration";
+    describe_histogram!(NAME, "duration for which a write transaction was kept open");
+    register_histogram!(NAME)
+});

sqld/src/metrics.rs

@@ -12,13 +12,14 @@ use chrono::NaiveDateTime;
 use enclose::enclose;
 use futures_core::Stream;
 use hyper::Uri;
+use metrics::histogram;
 use parking_lot::Mutex;
 use rusqlite::ErrorCode;
 use sqld_libsql_bindings::wal_hook::TRANSPARENT_METHODS;
 use tokio::io::AsyncBufReadExt;
 use tokio::sync::watch;
 use tokio::task::JoinSet;
-use tokio::time::Duration;
+use tokio::time::{Duration, Instant};
 use tokio_util::io::StreamReader;
 use tonic::transport::Channel;
 use uuid::Uuid;
@@ -433,6 +434,7 @@ impl<M: MakeNamespace> NamespaceStore<M> {
     where
         Fun: FnOnce(&Namespace<M::Database>) -> R,
     {
+        let before_load = Instant::now();
         let lock = self.inner.store.upgradable_read().await;
         if let Some(ns) = lock.get(&namespace) {
             Ok(f(ns))
@@ -451,6 +453,10 @@ impl<M: MakeNamespace> NamespaceStore<M> {
             let ret = f(&ns);
             tracing::info!("loaded namespace: `{namespace}`");
             lock.insert(namespace, ns);
+
+            // NAMESPACE_LOAD_LATENCY.record(before_load.elapsed());
+            histogram!("namespace_load_latency", before_load.elapsed());
+
             Ok(ret)
         }
     }