Check node health in antithesis. Reduce metric granularity for buffered changes. Admin command for bookie consistency check

Author: gorbak25

crates/corro-admin/src/lib.rs

@@ -1,12 +1,13 @@
 use std::{
+    collections::HashMap,
     fmt::Display,
     time::{Duration, Instant},
 };
 
 use camino::Utf8PathBuf;
 use corro_types::{
     actor::{ActorId, ClusterId},
-    agent::{Agent, Bookie},
+    agent::{Agent, BookedVersions, Bookie},
     base::{CrsqlDbVersion, CrsqlSeq},
     broadcast::{FocaCmd, FocaInput},
     sqlite::SqlitePoolError,
@@ -113,6 +114,7 @@ pub enum Command {
 pub enum SyncCommand {
     Generate,
     ReconcileGaps,
+    CheckBookieConsistency,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -282,6 +284,54 @@ async fn handle_conn(
 
                     send_success(&mut stream).await;
                 }
+                Command::Sync(SyncCommand::CheckBookieConsistency) => {
+                    info_log(
+                        &mut stream,
+                        "checking in-memory and DB bookie consistency...",
+                    )
+                    .await;
+
+                    let rw_conn = match agent.pool().write_low().await {
+                        Ok(conn) => conn,
+                        Err(e) => {
+                            send_error(&mut stream, e).await;
+                            continue;
+                        }
+                    };
+
+                    let report = match block_in_place(|| check_bookie_consistency(&rw_conn, bookie))
+                    {
+                        Ok(report) => report,
+                        Err(e) => {
+                            send_error(&mut stream, e).await;
+                            continue;
+                        }
+                    };
+
+                    match serde_json::to_value(&report) {
+                        Ok(json) => send(&mut stream, Response::Json(json)).await,
+                        Err(e) => {
+                            send_error(&mut stream, e).await;
+                            continue;
+                        }
+                    }
+
+                    if !report.ok {
+                        send_error(
+                            &mut stream,
+                            format!(
+                                "bookie mismatch: value_mismatch={}, only_in_memory={}, only_in_db={}",
+                                report.value_mismatch_actors.len(),
+                                report.only_in_memory_actors.len(),
+                                report.only_in_db_actors.len()
+                            ),
+                        )
+                        .await;
+                        continue;
+                    }
+
+                    send_success(&mut stream).await;
+                }
                 Command::Cluster(ClusterCommand::Rejoin) => {
                     let (cb_tx, cb_rx) = oneshot::channel();
 
@@ -707,3 +757,66 @@ fn collapse_gaps(
 
     Ok((deleted, inserted))
 }
+
+#[derive(Debug, Serialize)]
+struct BookieConsistencyReport {
+    ok: bool,
+    in_memory_actors: usize,
+    db_actors: usize,
+    value_mismatch_actors: Vec<ActorId>,
+    only_in_memory_actors: Vec<ActorId>,
+    only_in_db_actors: Vec<ActorId>,
+}
+
+fn check_bookie_consistency(
+    conn: &rusqlite::Connection,
+    bookie: &Bookie,
+) -> rusqlite::Result<BookieConsistencyReport> {
+    let _bookie_write_guard = bookie.write_lock_blocking();
+
+    let db_map = BookedVersions::load_all_from_conn(conn)?;
+
+    let in_memory_map: HashMap<ActorId, BookedVersions> = {
+        let guard = bookie.owned_guard();
+        bookie
+            .iter(&guard)
+            .map(|(actor_id, booked)| {
+                let read = booked.read();
+                (*actor_id, (**read).clone())
+            })
+            .collect()
+    };
+
+    let mut value_mismatch_actors = Vec::new();
+    let mut only_in_memory_actors = Vec::new();
+    let mut only_in_db_actors = Vec::new();
+
+    for (actor_id, mem) in &in_memory_map {
+        match db_map.get(actor_id) {
+            Some(db) if db != mem => value_mismatch_actors.push(*actor_id),
+            Some(_) => {}
+            None => only_in_memory_actors.push(*actor_id),
+        }
+    }
+
+    for actor_id in db_map.keys() {
+        if !in_memory_map.contains_key(actor_id) {
+            only_in_db_actors.push(*actor_id);
+        }
+    }
+
+    value_mismatch_actors.sort_unstable();
+    only_in_memory_actors.sort_unstable();
+    only_in_db_actors.sort_unstable();
+
+    Ok(BookieConsistencyReport {
+        ok: value_mismatch_actors.is_empty()
+            && only_in_memory_actors.is_empty()
+            && only_in_db_actors.is_empty(),
+        in_memory_actors: in_memory_map.len(),
+        db_actors: db_map.len(),
+        value_mismatch_actors,
+        only_in_memory_actors,
+        only_in_db_actors,
+    })
+}

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

@@ -1,25 +1,28 @@
 use crate::transport::Transport;
-use corro_types::{actor::ActorId, agent::Agent};
+use antithesis_sdk::assert_always;
+use corro_types::agent::Agent;
 use metrics::gauge;
+use serde_json::json;
 use std::time::Duration;
 use tokio::task::block_in_place;
 use tracing::error;
 use tripwire::Tripwire;
 
 pub async fn metrics_loop(agent: Agent, transport: Transport, mut tripwire: Tripwire) {
     let mut metrics_interval = tokio::time::interval(Duration::from_secs(10));
+    let on_antithesis = std::env::var("ANTITHESIS_OUTPUT_DIR").is_ok();
 
     loop {
         tokio::select! {
             _ = metrics_interval.tick() => {},
             _ = &mut tripwire => break,
         }
 
-        block_in_place(|| collect_metrics(&agent, &transport));
+        block_in_place(|| collect_metrics(&agent, &transport, on_antithesis));
     }
 }
 
-pub fn collect_metrics(agent: &Agent, transport: &Transport) {
+pub fn collect_metrics(agent: &Agent, transport: &Transport, on_antithesis: bool) {
     agent.pool().emit_metrics();
     transport.emit_metrics();
 
@@ -50,38 +53,66 @@ pub fn collect_metrics(agent: &Agent, transport: &Transport) {
         }
     }
 
-    // TODO: collect from bookie?
+    // Buffered changes stats across all actors
     match conn
-        .prepare_cached("SELECT actor_id, (select count(site_id) FROM __corro_buffered_changes WHERE site_id = actor_id) FROM __corro_members")
+        .prepare_cached(
+            "
+            SELECT
+                COALESCE(json_extract(m.foca_state, '$.state'), 'orphan') AS peer_state,
+                count(*) AS total,
+                strftime('%s', 'now') - (min(bc.ts) >> 32) as oldest_age_seconds
+            FROM __corro_buffered_changes bc
+            LEFT JOIN __corro_members m ON m.actor_id = bc.site_id
+            GROUP BY peer_state
+            ",
+        )
         .and_then(|mut prepped| {
             prepped
-                .query_map((), |row| {
-                    Ok((row.get::<_, ActorId>(0)?, row.get::<_, i64>(1)?))
+                .query_map([], |row| {
+                    Ok((
+                        row.get::<_, String>(0)?,
+                        row.get::<_, i64>(1)?,
+                        row.get::<_, i64>(2)?,
+                    ))
                 })
                 .and_then(|mapped| mapped.collect::<Result<Vec<_>, _>>())
         }) {
-        Ok(mapped) => {
-            for (actor_id, count) in mapped {
-                gauge!("corro.db.buffered.changes.rows.total", "actor_id" => actor_id.to_string()).set(count as f64)
+        Ok(rows) => {
+            for (peer_state, total_count, oldest_age_seconds) in rows {
+                // v2 has O(N) time series while the v1 metric had O(N^2) time series
+                // The name got changed on purpose so the metric has a fresh namespace
+                gauge!("corro.db.buffered.changes.v2.total", "peer_state" => peer_state.clone())
+                    .set(total_count as f64);
+                gauge!("corro.db.buffered.changes.v2.oldest_age_seconds", "peer_state" => peer_state)
+                    .set(oldest_age_seconds as f64);
             }
         }
         Err(e) => {
-            error!("could not query count for buffered changes: {e}");
+            error!("could not query buffered changes v2 stats: {e}");
         }
     }
 
     match conn
-        .prepare_cached("select actor_id, sum((end - start) + 1) from __corro_bookkeeping_gaps group by actor_id")
+        .prepare_cached(
+            "
+            SELECT
+                COALESCE(json_extract(m.foca_state, '$.state'), 'orphan') AS peer_state,
+                COALESCE(SUM((g.end - g.start) + 1), 0) AS gaps_sum
+            FROM __corro_bookkeeping_gaps g
+            LEFT JOIN __corro_members m ON m.actor_id = g.actor_id
+            GROUP BY peer_state
+            ",
+        )
         .and_then(|mut prepped| {
             prepped
-                .query_map((), |row| {
-                    Ok((row.get::<_, ActorId>(0)?, row.get::<_, u64>(1)?))
+                .query_map([], |row| {
+                    Ok((row.get::<_, String>(0)?, row.get::<_, u64>(1)?))
                 })
                 .and_then(|mapped| mapped.collect::<Result<Vec<_>, _>>())
         }) {
-        Ok(mapped) => {
-            for (actor_id, sum) in mapped {
-                gauge!("corro.db.gaps.sum", "actor_id" => actor_id.to_string()).set(sum as f64)
+        Ok(rows) => {
+            for (peer_state, sum) in rows {
+                gauge!("corro.db.gaps.sum", "peer_state" => peer_state).set(sum as f64)
             }
         }
         Err(e) => {
@@ -109,4 +140,38 @@ pub fn collect_metrics(agent: &Agent, transport: &Transport) {
             gauge!("corro.db.wal.size").set(meta.len() as f64);
         }
     }
+
+    if on_antithesis {
+        let gaps = conn
+            .prepare_cached(
+                "SELECT COALESCE(SUM(end - start + 1), 0) FROM __corro_bookkeeping_gaps",
+            )
+            .and_then(|mut prepped| prepped.query_row([], |row| row.get::<_, i64>(0)));
+
+        if let Ok(gaps) = gaps {
+            const MAX_GAPS: i64 = 100;
+            const MAX_QUEUE: u64 = 16_000;
+            const MAX_P99_LAG: f64 = 3.0;
+
+            let p99_lag = agent.metrics_tracker().quantile_lag(0.99).unwrap_or(0.0);
+            let queue_size = agent.metrics_tracker().queue_size();
+            let is_healthy = gaps <= MAX_GAPS && queue_size <= MAX_QUEUE && p99_lag <= MAX_P99_LAG;
+            let details = json!({
+                "gaps": gaps,
+                "queue_size": queue_size,
+                "p99_lag": p99_lag,
+                "max_gaps": MAX_GAPS,
+                "max_queue": MAX_QUEUE,
+                "max_p99_lag": MAX_P99_LAG,
+            });
+
+            assert_always!(
+                is_healthy,
+                "Corrosion node remains healthy in antithesis",
+                &details
+            );
+        } else if let Err(e) = gaps {
+            error!("could not query gaps for antithesis health assertion: {e}");
+        }
+    }
 }

crates/corrosion/src/main.rs

@@ -499,6 +499,13 @@ async fn process_cli(cli: Cli) -> eyre::Result<()> {
             ))
             .await?;
         }
+        Command::Sync(SyncCommand::CheckBookieConsistency) => {
+            let mut conn = AdminConn::connect(cli.admin_path()).await?;
+            conn.send_command(corro_admin::Command::Sync(
+                corro_admin::SyncCommand::CheckBookieConsistency,
+            ))
+            .await?;
+        }
         Command::Template { template, flags } => {
             command::tpl::run(cli.api_addr()?, template, flags).await?;
         }
@@ -780,6 +787,8 @@ enum ConsulCommand {
 enum SyncCommand {
     /// Generate a sync message from the current agent
     Generate,
+    /// Check in-memory bookie state against DB-loaded bookie state
+    CheckBookieConsistency,
     ReconcileGaps,
 }