introduce stats route

Author: MarinPostma

sqld/src/database/libsql.rs

@@ -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,18 @@ impl Connection {
             .execute("rollback transaction;", ())
             .expect("failed to rollback");
     }
+
+    fn update_stats(&self, stmt: &rusqlite::Statement) {
+        const STMT_ROWS_READ: i32 = 1024 + 1;
+        const STMT_ROWS_WRITTEN: i32 = 1024 + 2;
+
+        let rows_read: StatementStatus = unsafe { std::mem::transmute(STMT_ROWS_READ) };
+        let rows_written: StatementStatus = unsafe { std::mem::transmute(STMT_ROWS_WRITTEN) };
+
+        self.stats.inc_row_read(stmt.get_status(rows_read) as usize);
+        self.stats
+            .inc_row_written(stmt.get_status(rows_written) as usize);
+    }
 }
 
 fn eval_cond(cond: &Cond, results: &[Option<QueryResult>]) -> Result<bool> {

sqld/src/database/write_proxy.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/http/mod.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()),
     }
 }
@@ -274,6 +278,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 +322,7 @@ where
                 upgrade_tx.clone(),
                 db_factory.clone(),
                 enable_console,
+                stats.clone(),
             )
         });
 

sqld/src/lib.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();

sqld/src/stats.rs

@@ -0,0 +1,67 @@
+use std::fs::{File, OpenOptions};
+use std::io::Seek;
+use std::path::Path;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+use std::time::Duration;
+
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone)]
+pub struct Stats {
+    inner: Arc<StatsInner>,
+}
+
+#[derive(Serialize, Deserialize, Default)]
+struct StatsInner {
+    rows_written: AtomicUsize,
+    rows_read: AtomicUsize,
+}
+
+impl Stats {
+    pub fn new(db_path: &Path) -> anyhow::Result<Self> {
+        let stats_path = db_path.join("stats.json");
+        let stats_file = OpenOptions::new()
+            .read(true)
+            .write(true)
+            .create(true)
+            .open(stats_path)?;
+
+        let stats_inner =
+            serde_json::from_reader(&stats_file).unwrap_or_else(|_| StatsInner::default());
+        let inner = Arc::new(stats_inner);
+
+        spawn_stats_persist_thread(inner.clone(), stats_file);
+
+        Ok(Self { inner })
+    }
+
+    /// increments the number of written rows by n
+    pub fn inc_row_written(&self, n: usize) {
+        self.inner.rows_written.fetch_add(n, Ordering::Relaxed);
+    }
+
+    /// increments the number of read rows by n
+    pub fn inc_row_read(&self, n: usize) {
+        self.inner.rows_read.fetch_add(n, Ordering::Relaxed);
+    }
+
+    /// returns the total number of rows read since this database was created
+    pub fn row_read(&self) -> usize {
+        self.inner.rows_read.load(Ordering::Relaxed)
+    }
+
+    /// returns the total number of rows written since this database was created
+    pub fn rows_written(&self) -> usize {
+        self.inner.rows_written.load(Ordering::Relaxed)
+    }
+}
+
+fn spawn_stats_persist_thread(stats: Arc<StatsInner>, mut file: File) {
+    std::thread::spawn(move || loop {
+        if file.seek(std::io::SeekFrom::Start(0)).is_ok() {
+            let _ = serde_json::to_writer(&mut file, &stats);
+        }
+        std::thread::sleep(Duration::from_secs(5));
+    });
+}