Add health endpoint that returns commit lag and gaps (#423)

* Add a simple lag tracker for health checks

Signed-off-by: Somtochi Onyekwere <somtochionyekwere@gmail.com>

* Use rolling-window bucket histogram so that track lag over a longer
period without increasing space usage.

* reduce window chunks

Signed-off-by: Somtochi Onyekwere <somtochionyekwere@gmail.com>

* track queue size too

Signed-off-by: Somtochi Onyekwere <somtochionyekwere@gmail.com>

* fix up comment

Signed-off-by: Somtochi Onyekwere <somtochionyekwere@gmail.com>

---------

Signed-off-by: Somtochi Onyekwere <somtochionyekwere@gmail.com>

Author: somtochiama

Cargo.lock

@@ -631,6 +631,7 @@ pub async fn handle_changes(
     ));
 
     let max_seen_cache_len: usize = max_queue_len;
+    let metrics_tracker = agent.metrics_tracker();
 
     // unlikely, but max_seen_cache_len can be less than 10, in that case we want to just clear the whole cache
     // (todo): put some validation in config instead
@@ -697,11 +698,11 @@ pub async fn handle_changes(
 
             _ = max_wait.tick() => {
                 // got a wait interval tick...
-
                 gauge!("corro.agent.changes.in_queue").set(buf_cost as f64);
                 gauge!("corro.agent.changesets.in_queue").set(queue.len() as f64);
                 gauge!("corro.agent.changes.processing.jobs").set(join_set.len() as f64);
 
+                metrics_tracker.observe_queue_size(queue.len() as u64);
                 if buf_cost < max_changes_chunk && !queue.is_empty() && join_set.len() < MAX_CONCURRENT {
                     // we can process this right away
                     debug!(%buf_cost, "spawning processing multiple changes from max wait interval");

crates/corro-agent/src/agent/handlers.rs

@@ -42,6 +42,7 @@ use corro_types::{
     channel::{bounded, CorroReceiver},
     config::Config,
     members::Members,
+    metrics_tracker::MetricsTracker,
     pubsub::{Matcher, SubsManager},
     schema::{init_schema, Schema},
     sqlite::CrConn,
@@ -243,6 +244,8 @@ pub async fn setup(conf: Config, tripwire: Tripwire) -> eyre::Result<(Agent, Age
         }
     });
 
+    let metrics_tracker = MetricsTracker::new(Duration::from_secs(120), 5)?;
+
     let opts = AgentOptions {
         gossip_server_endpoint,
         transport: transport.clone(),
@@ -280,6 +283,7 @@ pub async fn setup(conf: Config, tripwire: Tripwire) -> eyre::Result<(Agent, Age
         cluster_id,
         subs_manager,
         updates_manager,
+        metrics_tracker,
         tripwire,
     });
 

crates/corro-agent/src/agent/setup.rs

@@ -8,7 +8,7 @@
 use crate::{
     agent::{handlers, CountedExecutor, TO_CLEAR_COUNT},
     api::public::{
-        api_v1_db_schema, api_v1_queries, api_v1_table_stats, api_v1_transactions,
+        api_v1_db_schema, api_v1_health, api_v1_queries, api_v1_table_stats, api_v1_transactions,
         pubsub::{api_v1_sub_by_id, api_v1_subs},
         update::SharedUpdateBroadcastCache,
     },
@@ -285,6 +285,20 @@ pub async fn setup_http_api_handler(
                     .layer(ConcurrencyLimitLayer::new(4)),
             ),
         )
+        .route(
+            "/v1/health",
+            get(api_v1_health).route_layer(
+                tower::ServiceBuilder::new()
+                    .layer(HandleErrorLayer::new(|_error: BoxError| async {
+                        Ok::<_, Infallible>((
+                            StatusCode::SERVICE_UNAVAILABLE,
+                            "max concurrency limit reached".to_string(),
+                        ))
+                    }))
+                    .layer(LoadShedLayer::new())
+                    .layer(ConcurrencyLimitLayer::new(4)),
+            ),
+        )
         .layer(axum::middleware::from_fn(require_authz))
         .layer(
             tower::ServiceBuilder::new()
@@ -1017,6 +1031,7 @@ pub async fn process_multiple_changes(
             if let Some(ts) = changeset.ts() {
                 let dur = (agent.clock().new_timestamp().get_time() - ts.0).to_duration();
                 histogram!("corro.agent.changes.commit.lag.seconds").record(dur);
+                agent.metrics_tracker().observe_lag(dur.as_secs_f64());
             }
         }
 

crates/corro-agent/src/agent/util.rs

@@ -7,14 +7,18 @@ use std::{
 
 use crate::api::utils::CountedBody;
 use antithesis_sdk::assert_sometimes;
-use axum::{extract::ConnectInfo, response::IntoResponse, Extension};
+use axum::{
+    extract::{ConnectInfo, Query},
+    response::IntoResponse,
+    Extension,
+};
 use bytes::{BufMut, BytesMut};
 use compact_str::ToCompactString;
 use corro_types::{
     agent::{Agent, ChangeError},
     api::{
-        ColumnName, ExecResponse, ExecResult, QueryEvent, Statement, TableStatRequest,
-        TableStatResponse,
+        ColumnName, ExecResponse, ExecResult, HealthQuery, HealthResponse, QueryEvent, Statement,
+        TableStatRequest, TableStatResponse,
     },
     base::CrsqlDbVersion,
     broadcast::Timestamp,
@@ -643,6 +647,64 @@ pub async fn api_v1_db_schema(
     )
 }
 
+pub async fn api_v1_health(
+    Extension(agent): Extension<Agent>,
+    Query(query): Query<HealthQuery>,
+) -> (StatusCode, axum::Json<HealthResponse>) {
+    match check_health(&agent).await {
+        Ok((gaps, members)) => {
+            let p99_lag = agent.metrics_tracker().quantile_lag(0.99).unwrap_or(0.0);
+            let queue_size = agent.metrics_tracker().queue_size();
+            let status = if query.gaps.is_some_and(|max| gaps > max)
+                // we use queue size and p99 lag as a stronger metric for an unhealthy node
+                // since a different node that is slow to send out changes can cause worse commit lag
+                // even though the node is perfectly fine.
+                || (query.p99_lag.is_some_and(|max| p99_lag > max)
+                    && query.queue_size.is_some_and(|max| queue_size > max))
+            {
+                StatusCode::SERVICE_UNAVAILABLE
+            } else {
+                StatusCode::OK
+            };
+            (
+                status,
+                axum::Json(HealthResponse::Response {
+                    gaps,
+                    members,
+                    p99_lag,
+                    queue_size,
+                }),
+            )
+        }
+        Err(e) => {
+            error!("could not check health: {e}");
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                axum::Json(HealthResponse::Error(e.to_string())),
+            )
+        }
+    }
+}
+
+async fn check_health(agent: &Agent) -> eyre::Result<(i64, i64)> {
+    let read_conn = match agent.pool().read().await {
+        Ok(conn) => conn,
+        Err(e) => {
+            error!("could not acquire read connection for health check: {e}");
+            return Err(eyre::eyre!("unable to grab write conn"));
+        }
+    };
+
+    let gaps = read_conn
+        .prepare_cached("SELECT COALESCE(SUM(end - start + 1), 0) FROM __corro_bookkeeping_gaps")?
+        .query_row([], |row| row.get::<_, i64>(0))?;
+
+    let members = read_conn.prepare_cached(r#"
+            SELECT COALESCE(COUNT(*), 0) FROM __corro_members WHERE json_extract(foca_state, "$.state") = "Alive""#)?
+        .query_row([], |row| row.get::<_, i64>(0))?;
+
+    Ok((gaps, members))
+}
 /// Query the table status of the current node
 ///
 /// Currently this endpoint only supports querying the row count for a

crates/corro-agent/src/api/public/mod.rs

@@ -241,6 +241,26 @@ pub struct TableStatResponse {
     pub invalid_tables: Vec<String>,
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+pub struct HealthQuery {
+    pub gaps: Option<i64>,
+    pub p99_lag: Option<f64>,
+    pub queue_size: Option<u64>,
+}
+
+/// Contains status information about the node
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum HealthResponse {
+    Response {
+        gaps: i64,
+        members: i64,
+        p99_lag: f64,
+        queue_size: u64,
+    },
+    Error(String),
+}
+
 #[derive(Debug, Copy, Clone, PartialEq)]
 pub struct SqliteValueRef<'a>(pub ValueRef<'a>);
 

crates/corro-api-types/src/lib.rs

@@ -23,6 +23,7 @@ corro-api-types = { version = "0.1.0-alpha.1", path = "../corro-api-types" }
 corro-base-types = { version = "0.1.0-alpha.1", path = "../corro-base-types" }
 deadpool = { workspace = true }
 enquote = { workspace = true }
+eyre = { workspace = true }
 fallible-iterator = { workspace = true }
 foca = { workspace = true }
 futures = { workspace = true }

crates/corro-types/Cargo.toml

@@ -48,6 +48,7 @@ use crate::{
     broadcast::{BroadcastInput, ChangeSource, ChangeV1, FocaInput, Timestamp},
     channel::{bounded, CorroSender},
     config::Config,
+    metrics_tracker::MetricsTracker,
     pubsub::SubsManager,
     schema::Schema,
     sqlite::{
@@ -89,6 +90,8 @@ pub struct AgentConfig {
 
     pub updates_manager: UpdatesManager,
 
+    pub metrics_tracker: MetricsTracker,
+
     pub tripwire: Tripwire,
 }
 
@@ -100,6 +103,7 @@ pub struct AgentInner {
     external_addr: Option<SocketAddr>,
     api_addr: SocketAddr,
     members: RwLock<Members>,
+    metrics_tracker: MetricsTracker,
     clock: Arc<uhlc::HLC>,
     booked: Booked,
     tx_bcast: CorroSender<BroadcastInput>,
@@ -130,6 +134,7 @@ impl Agent {
             external_addr: config.external_addr,
             api_addr: config.api_addr,
             members: config.members,
+            metrics_tracker: config.metrics_tracker,
             clock: config.clock,
             booked: config.booked,
             tx_bcast: config.tx_bcast,
@@ -231,6 +236,10 @@ impl Agent {
         &self.0.schema
     }
 
+    pub fn metrics_tracker(&self) -> &MetricsTracker {
+        &self.0.metrics_tracker
+    }
+
     pub fn db_path(&self) -> Utf8PathBuf {
         self.0.config.load().db.path.clone()
     }

crates/corro-types/src/agent.rs

@@ -8,6 +8,7 @@ pub mod channel;
 pub mod config;
 pub mod gauge;
 pub mod members;
+pub mod metrics_tracker;
 pub mod pubsub;
 pub mod schema;
 pub mod sqlite;