Merge pull request #563 from tursodatabase/lucio/sim-test-metrics

sim-tests: add asserts for metrics

libsql-server/Cargo.toml

@@ -86,6 +86,7 @@ rand = "0.8.5"
 tempfile = "3.7.0"
 turmoil = "0.5.6"
 url = "2.3"
+metrics-util = "0.15"
 
 [build-dependencies]
 prost-build = "0.12.0"

libsql-server/src/connection/libsql.rs

@@ -3,7 +3,7 @@ use std::path::{Path, PathBuf};
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 
-use metrics::histogram;
+use metrics::{histogram, increment_counter};
 use parking_lot::{Mutex, RwLock};
 use rusqlite::{DatabaseName, ErrorCode, OpenFlags, StatementStatus, TransactionState};
 use sqld_libsql_bindings::wal_hook::{TransparentMethods, WalMethodsHook};
@@ -616,6 +616,9 @@ impl<W: WalHook> Connection<W> {
         builder: &mut impl QueryResultBuilder,
     ) -> Result<(u64, Option<i64>)> {
         tracing::trace!("executing query: {}", query.stmt.stmt);
+
+        increment_counter!("libsql_server_libsql_query_execute");
+
         let start = Instant::now();
         let config = self.config_store.get();
         let blocked = match query.stmt.kind {
@@ -857,6 +860,8 @@ where
         builder: B,
         _replication_index: Option<FrameNo>,
     ) -> Result<B> {
+        increment_counter!("libsql_server_libsql_execute_program");
+
         check_program_auth(auth, &pgm)?;
         let conn = self.inner.clone();
         tokio::task::spawn_blocking(move || Connection::run(conn, pgm, builder))

libsql-server/tests/cluster/mod.rs

@@ -12,7 +12,7 @@ use turmoil::{Builder, Sim};
 
 use common::net::{init_tracing, TestServer, TurmoilAcceptor, TurmoilConnector};
 
-use crate::common::{http::Client, net::SimServer};
+use crate::common::{http::Client, net::SimServer, snapshot_metrics};
 
 fn make_cluster(sim: &mut Sim, num_replica: usize, disable_namespaces: bool) {
     init_tracing();
@@ -101,6 +101,8 @@ fn proxy_write() {
             Value::Integer(1)
         ));
 
+        snapshot_metrics().assert_gauge("libsql_server_current_frame_no", 2.0);
+
         Ok(())
     });
 

libsql-server/tests/common/mod.rs

@@ -1,4 +1,83 @@
 #![allow(dead_code)]
 
+use std::collections::HashMap;
+
+use metrics::{SharedString, Unit};
+use metrics_util::{
+    debugging::{DebugValue, Snapshotter},
+    CompositeKey, MetricKind,
+};
+
 pub mod http;
 pub mod net;
+
+pub fn print_metrics() {
+    let snapshot = MetricsSnapshot::current();
+
+    eprintln!("{:?}", snapshot);
+}
+
+#[track_caller]
+pub fn snapshot_metrics() -> MetricsSnapshot {
+    MetricsSnapshot::current()
+}
+
+#[derive(Debug)]
+pub struct MetricsSnapshot {
+    snapshot: HashMap<CompositeKey, (Option<Unit>, Option<SharedString>, DebugValue)>,
+}
+
+impl MetricsSnapshot {
+    #[track_caller]
+    pub fn current() -> Self {
+        let snapshot = Snapshotter::current_thread_snapshot()
+            .expect("No snapshot available")
+            .into_hashmap();
+
+        MetricsSnapshot { snapshot }
+    }
+
+    pub fn get_counter(&self, metric_name: &str) -> Option<u64> {
+        for (key, (_, _, val)) in &self.snapshot {
+            if key.kind() == MetricKind::Counter && key.key().name() == metric_name {
+                match val {
+                    DebugValue::Counter(v) => return Some(*v),
+                    _ => unreachable!(),
+                }
+            }
+        }
+
+        None
+    }
+
+    pub fn get_gauge(&self, metric_name: &str) -> Option<f64> {
+        for (key, (_, _, val)) in &self.snapshot {
+            if key.kind() == MetricKind::Gauge && key.key().name() == metric_name {
+                match val {
+                    DebugValue::Gauge(v) => return Some(v.0),
+                    _ => unreachable!(),
+                }
+            }
+        }
+
+        None
+    }
+
+    #[track_caller]
+    pub fn assert_gauge(&self, metric_name: &str, value: f64) -> &Self {
+        let val = self.get_gauge(metric_name).expect("metric does not exist");
+
+        assert_eq!(val, value);
+        self
+    }
+
+    #[track_caller]
+    pub fn assert_counter(&self, metric_name: &str, value: u64) -> &Self {
+        let val = self
+            .get_counter(metric_name)
+            .expect("metric does not exist");
+
+        assert_eq!(val, value);
+        self
+    }
+}

libsql-server/tests/common/net.rs

@@ -8,6 +8,7 @@ use futures_core::Future;
 use hyper::client::connect::Connected;
 use hyper::server::accept::Accept as HyperAccept;
 use hyper::Uri;
+use metrics_util::debugging::DebuggingRecorder;
 use tokio::io::{AsyncRead, AsyncWrite};
 use tower::Service;
 use tracing_subscriber::{fmt, prelude::*, EnvFilter};
@@ -152,9 +153,17 @@ impl SimServer for TestServer {
         let db = libsql::Database::open_in_memory().unwrap();
         db.connect().unwrap();
 
+        // Ignore the result because we may set it many times in a single process.
+        let _ = DebuggingRecorder::per_thread().install();
+
         let user_api = TurmoilAcceptor::bind(([0, 0, 0, 0], user_api_port as u16)).await?;
         self.user_api_config.http_acceptor = Some(user_api);
 
+        // Disable prom metrics since we already created our recorder.
+        if let Some(admin_api) = &mut self.admin_api_config {
+            admin_api.disable_metrics = true;
+        }
+
         self.start().await?;
 
         Ok(())

libsql-server/tests/standalone/mod.rs

@@ -1,6 +1,6 @@
 //! Tests for standalone primary configuration
 
-use crate::common::net::SimServer;
+use crate::common::{net::SimServer, snapshot_metrics};
 
 use super::common;
 
@@ -57,6 +57,36 @@ fn basic_query() {
     sim.run().unwrap();
 }
 
+#[test]
+fn basic_metrics() {
+    let mut sim = turmoil::Builder::new().build();
+
+    sim.host("primary", make_standalone_server);
+
+    sim.client("test", async {
+        let db = Database::open_remote_with_connector("http://primary:8080", "", TurmoilConnector)?;
+        let conn = db.connect()?;
+
+        conn.execute("create table test (x)", ()).await?;
+        conn.execute("insert into test values (12)", ()).await?;
+
+        let mut rows = conn.query("select count(*) from test", ()).await?;
+
+        assert!(matches!(
+            rows.next().unwrap().unwrap().get_value(0).unwrap(),
+            libsql::Value::Integer(1)
+        ));
+
+        snapshot_metrics()
+            .assert_gauge("libsql_server_current_frame_no", 2.0)
+            .assert_counter("libsql_server_libsql_execute_program", 3);
+
+        Ok(())
+    });
+
+    sim.run().unwrap();
+}
+
 #[test]
 fn primary_serializability() {
     let mut sim = turmoil::Builder::new().build();

libsql-server/tests/tests.rs

@@ -1,7 +1,9 @@
 #![allow(clippy::disallowed_names)]
 
-mod cluster;
+#[macro_use]
 mod common;
+
+mod cluster;
 mod embedded_replica;
 mod namespaces;
 mod standalone;