Merge #315

315: introduce stats route r=MarinPostma a=MarinPostma

This PR introduces the `GET /v1/stats` route that return the total amount of rows read/written by an instance:

```
{
    "rows_read_count": 0,
    "rows_written_count": 0
}
```

The number of rows read/written persists accross restarts, thanks to a sync loop that flushes the information to disk about every 5s.


Co-authored-by: ad hoc <postma.marin@protonmail.com>
Co-authored-by: Piotr Sarna <sarna@chiselstrike.com>

Author: 3 people

Cargo.lock

@@ -9,7 +9,7 @@ edition = "2021"
 anyhow = "1.0.66"
 mvfs = { git = "https://github.com/psarna/mvsqlite", branch = "mwal", optional = true }
 mwal = { git = "https://github.com/psarna/mvsqlite", branch = "mwal", optional = true }
-rusqlite = { version = "0.29.0", git = "https://github.com/psarna/rusqlite", rev = "63b7aabfccbc21738", default-features = false, features = [
+rusqlite = { version = "0.29.0", git = "https://github.com/psarna/rusqlite", rev = "a6332e530f30dc2d47110", default-features = false, features = [
     "buildtime_bindgen",
     "bundled-libsql-wasm-experimental",
     "column_decltype"

sqld-libsql-bindings/Cargo.toml

@@ -33,7 +33,7 @@ pin-project-lite = "0.2.9"
 postgres-protocol = "0.6.4"
 prost = "0.11.3"
 regex = "1.7.0"
-rusqlite = { version = "0.29.0", git = "https://github.com/psarna/rusqlite", rev = "63b7aabfccbc21738", default-features = false, features = [
+rusqlite = { version = "0.29.0", git = "https://github.com/psarna/rusqlite", rev = "a6332e530f30dc2d47110", default-features = false, features = [
     "buildtime_bindgen",
     "bundled-libsql-wasm-experimental",
     "column_decltype"

sqld/Cargo.toml

@@ -3,14 +3,15 @@ use std::str::FromStr;
 use std::time::{Duration, Instant};
 
 use crossbeam::channel::RecvTimeoutError;
-use rusqlite::OpenFlags;
+use rusqlite::{OpenFlags, StatementStatus};
 use tokio::sync::oneshot;
 use tracing::warn;
 
 use crate::error::Error;
 use crate::libsql::wal_hook::WalHook;
 use crate::query::{Column, Query, QueryResponse, QueryResult, ResultSet, Row};
 use crate::query_analysis::{State, Statement};
+use crate::stats::Stats;
 use crate::Result;
 
 use super::{Cond, Database, Program, Step, TXN_TIMEOUT_SECS};
@@ -143,11 +144,13 @@ impl LibSqlDb {
         path: impl AsRef<Path> + Send + 'static,
         wal_hook: impl WalHook + Send + Clone + 'static,
         with_bottomless: bool,
+        stats: Stats,
     ) -> crate::Result<Self> {
         let (sender, receiver) = crossbeam::channel::unbounded::<Message>();
 
         tokio::task::spawn_blocking(move || {
-            let mut connection = Connection::new(path.as_ref(), wal_hook, with_bottomless).unwrap();
+            let mut connection =
+                Connection::new(path.as_ref(), wal_hook, with_bottomless, stats).unwrap();
             loop {
                 let Message { pgm, resp } = match connection.state.deadline() {
                     Some(deadline) => match receiver.recv_deadline(deadline) {
@@ -189,18 +192,21 @@ struct Connection {
     state: ConnectionState,
     conn: rusqlite::Connection,
     timed_out: bool,
+    stats: Stats,
 }
 
 impl Connection {
     fn new(
         path: &Path,
         wal_hook: impl WalHook + Send + Clone + 'static,
         with_bottomless: bool,
+        stats: Stats,
     ) -> anyhow::Result<Self> {
         Ok(Self {
             conn: open_db(path, wal_hook, with_bottomless)?,
             state: ConnectionState::initial(),
             timed_out: false,
+            stats,
         })
     }
 
@@ -245,8 +251,8 @@ impl Connection {
 
     fn execute_query_inner(&self, query: &Query) -> QueryResult {
         let mut rows = vec![];
-        let mut prepared = self.conn.prepare(&query.stmt.stmt)?;
-        let columns = prepared
+        let mut stmt = self.conn.prepare(&query.stmt.stmt)?;
+        let columns = stmt
             .columns()
             .iter()
             .map(|col| Column {
@@ -262,10 +268,10 @@ impl Connection {
 
         query
             .params
-            .bind(&mut prepared)
+            .bind(&mut stmt)
             .map_err(Error::LibSqlInvalidQueryParams)?;
 
-        let mut qresult = prepared.raw_query();
+        let mut qresult = stmt.raw_query();
         while let Some(row) = qresult.next()? {
             let mut values = vec![];
             for (i, _) in columns.iter().enumerate() {
@@ -288,6 +294,10 @@ impl Connection {
             false => None,
         };
 
+        drop(qresult);
+
+        self.update_stats(&stmt);
+
         Ok(QueryResponse::ResultSet(ResultSet {
             columns,
             rows,
@@ -302,6 +312,13 @@ impl Connection {
             .execute("rollback transaction;", ())
             .expect("failed to rollback");
     }
+
+    fn update_stats(&self, stmt: &rusqlite::Statement) {
+        self.stats
+            .inc_rows_read(stmt.get_status(StatementStatus::RowsRead) as usize);
+        self.stats
+            .inc_rows_written(stmt.get_status(StatementStatus::RowsWritten) as usize);
+    }
 }
 
 fn eval_cond(cond: &Cond, results: &[Option<QueryResult>]) -> Result<bool> {

sqld/src/database/libsql.rs

@@ -16,6 +16,7 @@ use crate::rpc::proxy::rpc::proxy_client::ProxyClient;
 use crate::rpc::proxy::rpc::query_result::RowResult;
 use crate::rpc::proxy::rpc::DisconnectMessage;
 use crate::rpc::replication_log::rpc::replication_log_client::ReplicationLogClient;
+use crate::stats::Stats;
 use crate::Result;
 
 use super::Program;
@@ -25,6 +26,7 @@ use super::{libsql::LibSqlDb, service::DbFactory, Database};
 pub struct WriteProxyDbFactory {
     write_proxy: ProxyClient<Channel>,
     db_path: PathBuf,
+    stats: Stats,
     /// abort handle: abort db update loop on drop
     _abort_handle: crossbeam::channel::Sender<()>,
 }
@@ -37,6 +39,7 @@ impl WriteProxyDbFactory {
         key_path: Option<PathBuf>,
         ca_cert_path: Option<PathBuf>,
         db_path: PathBuf,
+        stats: Stats,
     ) -> anyhow::Result<(Self, JoinHandle<anyhow::Result<()>>)> {
         let mut endpoint = Channel::from_shared(addr.to_string())?;
         if tls {
@@ -84,6 +87,7 @@ impl WriteProxyDbFactory {
         let this = Self {
             write_proxy,
             db_path,
+            stats,
             _abort_handle,
         };
         Ok((this, handle))
@@ -93,7 +97,11 @@ impl WriteProxyDbFactory {
 #[async_trait::async_trait]
 impl DbFactory for WriteProxyDbFactory {
     async fn create(&self) -> Result<Arc<dyn Database>> {
-        let db = WriteProxyDatabase::new(self.write_proxy.clone(), self.db_path.clone())?;
+        let db = WriteProxyDatabase::new(
+            self.write_proxy.clone(),
+            self.db_path.clone(),
+            self.stats.clone(),
+        )?;
         Ok(Arc::new(db))
     }
 }
@@ -106,8 +114,8 @@ pub struct WriteProxyDatabase {
 }
 
 impl WriteProxyDatabase {
-    fn new(write_proxy: ProxyClient<Channel>, path: PathBuf) -> Result<Self> {
-        let read_db = LibSqlDb::new(path, (), false)?;
+    fn new(write_proxy: ProxyClient<Channel>, path: PathBuf, stats: Stats) -> Result<Self> {
+        let read_db = LibSqlDb::new(path, (), false, stats)?;
         Ok(Self {
             read_db,
             write_proxy,

sqld/src/database/write_proxy.rs

@@ -1,4 +1,5 @@
 mod hrana_over_http;
+mod stats;
 mod types;
 
 use std::future::poll_fn;
@@ -29,6 +30,7 @@ use crate::hrana;
 use crate::http::types::HttpQuery;
 use crate::query::{self, Query, QueryResult, ResultSet};
 use crate::query_analysis::{predict_final_state, State, Statement};
+use crate::stats::Stats;
 use crate::utils::services::idle_shutdown::IdleShutdownLayer;
 
 use self::types::QueryObject;
@@ -236,6 +238,7 @@ async fn handle_request(
     upgrade_tx: mpsc::Sender<hrana::Upgrade>,
     db_factory: Arc<dyn DbFactory>,
     enable_console: bool,
+    stats: Stats,
 ) -> anyhow::Result<Response<Body>> {
     if hyper_tungstenite::is_upgrade_request(&req) {
         return Ok(handle_upgrade(&upgrade_tx, req).await);
@@ -257,6 +260,7 @@ async fn handle_request(
         (&Method::GET, "/v1") => hrana_over_http::handle_index(req).await,
         (&Method::POST, "/v1/execute") => hrana_over_http::handle_execute(db_factory, req).await,
         (&Method::POST, "/v1/batch") => hrana_over_http::handle_batch(db_factory, req).await,
+        (&Method::GET, "/v1/stats") => Ok(stats::handle_stats(&stats)),
         _ => Ok(Response::builder().status(404).body(Body::empty()).unwrap()),
     }
 }
@@ -266,6 +270,8 @@ fn handle_version() -> Response<Body> {
     Response::new(Body::from(version.as_bytes()))
 }
 
+// TODO: refactor
+#[allow(clippy::too_many_arguments)]
 pub async fn run_http<F>(
     addr: SocketAddr,
     auth: Arc<Auth>,
@@ -274,6 +280,7 @@ pub async fn run_http<F>(
     upgrade_tx: mpsc::Sender<hrana::Upgrade>,
     enable_console: bool,
     idle_shutdown_layer: Option<IdleShutdownLayer>,
+    stats: Stats,
 ) -> anyhow::Result<()>
 where
     F: MakeService<(), Vec<Query>> + Send + 'static,
@@ -317,6 +324,7 @@ where
                 upgrade_tx.clone(),
                 db_factory.clone(),
                 enable_console,
+                stats.clone(),
             )
         });
 

sqld/src/http/mod.rs

@@ -0,0 +1,23 @@
+use hyper::{Body, Response};
+use serde::Serialize;
+
+use crate::stats::Stats;
+
+#[derive(Serialize)]
+struct StatsResponse {
+    rows_read_count: usize,
+    rows_written_count: usize,
+}
+
+pub fn handle_stats(stats: &Stats) -> Response<Body> {
+    let resp = StatsResponse {
+        rows_read_count: stats.rows_read(),
+        rows_written_count: stats.rows_written(),
+    };
+
+    let payload = serde_json::to_vec(&resp).unwrap();
+    Response::builder()
+        .header("Content-Type", "application/json")
+        .body(Body::from(payload))
+        .unwrap()
+}

sqld/src/http/stats.rs

@@ -25,6 +25,7 @@ use crate::auth::Auth;
 use crate::error::Error;
 use crate::postgres::service::PgConnectionFactory;
 use crate::server::Server;
+use crate::stats::Stats;
 
 pub use sqld_libsql_bindings as libsql;
 
@@ -39,6 +40,7 @@ mod query_analysis;
 mod replication;
 pub mod rpc;
 mod server;
+mod stats;
 mod utils;
 
 #[derive(clap::ValueEnum, Clone, Debug, PartialEq)]
@@ -97,6 +99,7 @@ async fn run_service(
     config: &Config,
     join_set: &mut JoinSet<anyhow::Result<()>>,
     idle_shutdown_layer: Option<IdleShutdownLayer>,
+    stats: Stats,
 ) -> anyhow::Result<()> {
     let auth = get_auth(config)?;
 
@@ -122,6 +125,7 @@ async fn run_service(
             upgrade_tx,
             config.enable_http_console,
             idle_shutdown_layer.clone(),
+            stats,
         ));
     }
 
@@ -183,6 +187,7 @@ async fn start_replica(
     join_set: &mut JoinSet<anyhow::Result<()>>,
     addr: &str,
     idle_shutdown_layer: Option<IdleShutdownLayer>,
+    stats: Stats,
 ) -> anyhow::Result<()> {
     let (factory, handle) = WriteProxyDbFactory::new(
         addr,
@@ -191,6 +196,7 @@ async fn start_replica(
         config.writer_rpc_key.clone(),
         config.writer_rpc_ca_cert.clone(),
         config.db_path.clone(),
+        stats.clone(),
     )
     .await
     .context("failed to start WriteProxy DB")?;
@@ -201,7 +207,7 @@ async fn start_replica(
     join_set.spawn(async move { handle.await.expect("WriteProxy DB task failed") });
 
     let service = DbFactoryService::new(Arc::new(factory));
-    run_service(service, config, join_set, idle_shutdown_layer).await?;
+    run_service(service, config, join_set, idle_shutdown_layer, stats).await?;
 
     Ok(())
 }
@@ -214,6 +220,7 @@ async fn start_primary(
     config: &Config,
     join_set: &mut JoinSet<anyhow::Result<()>>,
     idle_shutdown_layer: Option<IdleShutdownLayer>,
+    stats: Stats,
 ) -> anyhow::Result<()> {
     let is_fresh_db = check_fresh_db(&config.db_path);
     let logger = Arc::new(ReplicationLogger::open(
@@ -237,10 +244,12 @@ async fn start_primary(
         dump_loader.load_dump(path.into()).await?;
     }
 
+    let stats_clone = stats.clone();
     let db_factory = Arc::new(move || {
         let db_path = path_clone.clone();
         let hook = hook.clone();
-        async move { LibSqlDb::new(db_path, hook, enable_bottomless) }
+        let stats_clone = stats_clone.clone();
+        async move { LibSqlDb::new(db_path, hook, enable_bottomless, stats_clone) }
     });
     let service = DbFactoryService::new(db_factory.clone());
     if let Some(ref addr) = config.rpc_server_addr {
@@ -256,7 +265,7 @@ async fn start_primary(
         ));
     }
 
-    run_service(service, config, join_set, idle_shutdown_layer).await?;
+    run_service(service, config, join_set, idle_shutdown_layer, stats).await?;
 
     Ok(())
 }
@@ -309,11 +318,13 @@ pub async fn run_server(config: Config) -> anyhow::Result<()> {
             .idle_shutdown_timeout
             .map(|d| IdleShutdownLayer::new(d, shutdown_notify.clone()));
 
+        let stats = Stats::new(&config.db_path)?;
+
         match config.writer_rpc_addr {
             Some(ref addr) => {
-                start_replica(&config, &mut join_set, addr, idle_shutdown_layer).await?
+                start_replica(&config, &mut join_set, addr, idle_shutdown_layer, stats).await?
             }
-            None => start_primary(&config, &mut join_set, idle_shutdown_layer).await?,
+            None => start_primary(&config, &mut join_set, idle_shutdown_layer, stats).await?,
         }
 
         let reset = HARD_RESET.clone();