Namespace stats

sqld/src/connection/libsql.rs

@@ -28,7 +28,7 @@ pub struct LibSqlDbFactory<W: WalHook + 'static> {
     db_path: PathBuf,
     hook: &'static WalMethodsHook<W>,
     ctx_builder: Box<dyn Fn() -> W::Context + Sync + Send + 'static>,
-    stats: Stats,
+    stats: Arc<Stats>,
     config_store: Arc<DatabaseConfigStore>,
     extensions: Arc<[PathBuf]>,
     max_response_size: u64,
@@ -49,7 +49,7 @@ where
         db_path: PathBuf,
         hook: &'static WalMethodsHook<W>,
         ctx_builder: F,
-        stats: Stats,
+        stats: Arc<Stats>,
         config_store: Arc<DatabaseConfigStore>,
         extensions: Arc<[PathBuf]>,
         max_response_size: u64,
@@ -168,7 +168,7 @@ impl LibSqlConnection {
         extensions: Arc<[PathBuf]>,
         wal_hook: &'static WalMethodsHook<W>,
         hook_ctx: W::Context,
-        stats: Stats,
+        stats: Arc<Stats>,
         config_store: Arc<DatabaseConfigStore>,
         builder_config: QueryBuilderConfig,
     ) -> crate::Result<Self>
@@ -242,7 +242,7 @@ struct Connection<'a> {
     timeout_deadline: Option<Instant>,
     conn: sqld_libsql_bindings::Connection<'a>,
     timed_out: bool,
-    stats: Stats,
+    stats: Arc<Stats>,
     config_store: Arc<DatabaseConfigStore>,
     builder_config: QueryBuilderConfig,
 }
@@ -253,7 +253,7 @@ impl<'a> Connection<'a> {
         extensions: Arc<[PathBuf]>,
         wal_methods: &'static WalMethodsHook<W>,
         hook_ctx: &'a mut W::Context,
-        stats: Stats,
+        stats: Arc<Stats>,
         config_store: Arc<DatabaseConfigStore>,
         builder_config: QueryBuilderConfig,
     ) -> Result<Self> {
@@ -612,7 +612,7 @@ mod test {
             timeout_deadline: None,
             conn: sqld_libsql_bindings::Connection::test(ctx),
             timed_out: false,
-            stats: Stats::default(),
+            stats: Arc::new(Stats::default()),
             config_store: Arc::new(DatabaseConfigStore::new_test()),
             builder_config: QueryBuilderConfig::default(),
         };

sqld/src/connection/write_proxy.rs

@@ -37,7 +37,7 @@ pub struct MakeWriteProxyConnection {
     client: ProxyClient<Channel>,
     db_path: PathBuf,
     extensions: Arc<[PathBuf]>,
-    stats: Stats,
+    stats: Arc<Stats>,
     config_store: Arc<DatabaseConfigStore>,
     applied_frame_no_receiver: watch::Receiver<FrameNo>,
     max_response_size: u64,
@@ -52,7 +52,7 @@ impl MakeWriteProxyConnection {
         extensions: Arc<[PathBuf]>,
         channel: Channel,
         uri: tonic::transport::Uri,
-        stats: Stats,
+        stats: Arc<Stats>,
         config_store: Arc<DatabaseConfigStore>,
         applied_frame_no_receiver: watch::Receiver<FrameNo>,
         max_response_size: u64,
@@ -110,7 +110,7 @@ pub struct WriteProxyConnection {
     /// Notifier from the repliator of the currently applied frameno
     applied_frame_no_receiver: watch::Receiver<FrameNo>,
     builder_config: QueryBuilderConfig,
-    stats: Stats,
+    stats: Arc<Stats>,
     /// bytes representing the namespace name
     namespace: Bytes,
 }
@@ -166,7 +166,7 @@ impl WriteProxyConnection {
         write_proxy: ProxyClient<Channel>,
         db_path: PathBuf,
         extensions: Arc<[PathBuf]>,
-        stats: Stats,
+        stats: Arc<Stats>,
         config_store: Arc<DatabaseConfigStore>,
         applied_frame_no_receiver: watch::Receiver<FrameNo>,
         builder_config: QueryBuilderConfig,

sqld/src/heartbeat.rs

@@ -1,28 +1,62 @@
+use bytes::Bytes;
+use std::collections::HashMap;
+use std::sync::Weak;
 use std::time::Duration;
-use tokio::time::sleep;
+use tokio::sync::mpsc;
+use url::Url;
 
-use crate::http::stats::StatsResponse;
+use crate::http::admin::stats::StatsResponse;
 use crate::stats::Stats;
 
 pub async fn server_heartbeat(
-    url: String,
+    url: Url,
     auth: Option<String>,
     update_period: Duration,
-    stats: Stats,
+    mut stats_subs: mpsc::Receiver<(Bytes, Weak<Stats>)>,
 ) {
+    let mut watched = HashMap::new();
     let client = reqwest::Client::new();
+    let mut interval = tokio::time::interval(update_period);
+    interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
     loop {
-        sleep(update_period).await;
-        let body = StatsResponse::from(&stats);
-        let request = client.post(&url);
-        let request = if let Some(ref auth) = auth {
-            request.header("Authorization", auth.clone())
-        } else {
-            request
+        tokio::select! {
+            Some((ns, stats)) = stats_subs.recv() => {
+                watched.insert(ns, stats);
+            }
+            _ = interval.tick() => {
+                send_stats(&mut watched, &client, &url, auth.as_deref()).await;
+            }
         };
-        let request = request.json(&body);
-        if let Err(err) = request.send().await {
-            tracing::warn!("Error sending heartbeat: {}", err);
+    }
+}
+
+async fn send_stats(
+    watched: &mut HashMap<Bytes, Weak<Stats>>,
+    client: &reqwest::Client,
+    url: &Url,
+    auth: Option<&str>,
+) {
+    // first send all the stats...
+    for (ns, stats) in watched.iter() {
+        if let Some(stats) = stats.upgrade() {
+            let body = StatsResponse::from(stats.as_ref());
+            let mut url = url.clone();
+            url.path_segments_mut()
+                .unwrap()
+                .push(std::str::from_utf8(ns).unwrap());
+            let request = client.post(url);
+            let request = if let Some(ref auth) = auth {
+                request.header("Authorization", auth.to_string())
+            } else {
+                request
+            };
+            let request = request.json(&body);
+            if let Err(err) = request.send().await {
+                tracing::warn!("Error sending heartbeat: {}", err);
+            }
         }
     }
+
+    // ..and then remove all expired subscription
+    watched.retain(|_, s| s.upgrade().is_some());
 }

sqld/src/hrana/ws/handshake.rs

@@ -4,7 +4,7 @@ use futures::{SinkExt as _, StreamExt as _};
 use tokio_tungstenite::tungstenite;
 use tungstenite::http;
 
-use crate::http::db_factory::namespace_from_headers;
+use crate::http::user::db_factory::namespace_from_headers;
 use crate::net::Conn;
 
 use super::super::{Encoding, Version};

sqld/src/admin_api.rs → sqld/src/http/admin/mod.rs

@@ -15,12 +15,14 @@ use crate::connection::config::{DatabaseConfig, DatabaseConfigStore};
 use crate::error::LoadDumpError;
 use crate::namespace::{DumpStream, MakeNamespace, NamespaceStore, RestoreOption};
 
+pub mod stats;
+
 struct AppState<M: MakeNamespace> {
     db_config_store: Arc<DatabaseConfigStore>,
     namespaces: NamespaceStore<M>,
 }
 
-pub async fn run_admin_api<M, A>(
+pub async fn run<M, A>(
     acceptor: A,
     db_config_store: Arc<DatabaseConfigStore>,
     namespaces: NamespaceStore<M>,
@@ -47,6 +49,7 @@ where
             post(handle_restore_namespace),
         )
         .route("/v1/namespaces/:namespace", delete(handle_delete_namespace))
+        .route("/v1/namespaces/:namespace/stats", get(stats::handle_stats))
         .with_state(Arc::new(AppState {
             db_config_store,
             namespaces,

sqld/src/http/stats.rs → sqld/src/http/admin/stats.rs

@@ -1,9 +1,12 @@
-use hyper::{Body, Response};
+use std::sync::Arc;
+
 use serde::Serialize;
 
-use axum::extract::{FromRef, State as AxumState};
+use axum::extract::{Path, State};
+use axum::Json;
 
-use crate::{namespace::MakeNamespace, stats::Stats};
+use crate::namespace::MakeNamespace;
+use crate::stats::Stats;
 
 use super::AppState;
 
@@ -32,18 +35,12 @@ impl From<Stats> for StatsResponse {
     }
 }
 
-impl<F: MakeNamespace> FromRef<AppState<F>> for Stats {
-    fn from_ref(input: &AppState<F>) -> Self {
-        input.stats.clone()
-    }
-}
-
-pub(crate) async fn handle_stats(AxumState(stats): AxumState<Stats>) -> Response<Body> {
-    let resp: StatsResponse = stats.into();
+pub(super) async fn handle_stats<M: MakeNamespace>(
+    State(app_state): State<Arc<AppState<M>>>,
+    Path(namespace): Path<String>,
+) -> crate::Result<Json<StatsResponse>> {
+    let stats = app_state.namespaces.stats(namespace.into()).await?;
+    let resp: StatsResponse = stats.as_ref().into();
 
-    let payload = serde_json::to_vec(&resp).unwrap();
-    Response::builder()
-        .header("Content-Type", "application/json")
-        .body(Body::from(payload))
-        .unwrap()
+    Ok(Json(resp))
 }