use super::Hydrant;
use miette::{IntoDiagnostic, Result};
use std::collections::BTreeMap;
/// database statistics returned by [`Hydrant::stats`].
#[derive(serde::Serialize)]
pub struct StatsResponse {
/// record counts per logical category (repos, records, events, error kinds, etc.)
pub counts: BTreeMap<&'static str, u64>,
/// on-disk size in bytes per keyspace
pub sizes: BTreeMap<&'static str, u64>,
}
impl Hydrant {
/// return database counts and on-disk sizes for all keyspaces.
///
/// counts include: `repos`, `pending`, `resync`, `records`, `history`, `events`,
/// `error_ratelimited`, `error_transport`, `error_generic`. `history`, `events`
/// and the relay and jetstream event counts are approximate: they come from
/// the keyspace's table metadata, so deleted or expired entries still count
/// until compaction drops them.
///
/// sizes are in bytes, reported per keyspace.
pub async fn stats(&self) -> Result<StatsResponse> {
let state = self.state.clone();
#[allow(unused_mut)]
let mut count_keys = vec![
"repos",
"error_ratelimited",
"error_transport",
"error_generic",
];
#[cfg(feature = "indexer")]
{
count_keys.push("pending");
count_keys.push("records");
count_keys.push("resync");
}
#[cfg_attr(not(any(feature = "indexer", feature = "relay")), allow(unused_mut))]
let mut counts: BTreeMap<&'static str, u64> =
futures::future::join_all(count_keys.into_iter().map(|name| {
let state = state.clone();
async move { (name, state.db.get_count(name).await) }
}))
.await
.into_iter()
.collect();
// a write counter can't see what the retention filter drops, so this
// reads the keyspace itself
#[cfg(feature = "indexer")]
counts.insert(
"history",
state.db.indexer.approximate_history_count() as u64,
);
#[cfg(feature = "indexer_stream")]
counts.insert("events", state.db.stream.approximate_event_count() as u64);
#[cfg(feature = "indexer")]
counts.insert(
"refetch_denied",
state
.refetch_denied
.load(std::sync::atomic::Ordering::Relaxed),
);
// only meaningful when bounded history can legitimately trim bodies
#[cfg(feature = "indexer_stream")]
if state.history_ttl.is_some() {
counts.insert(
"history_trim_misses",
state
.history_trim_misses
.load(std::sync::atomic::Ordering::Relaxed),
);
}
#[cfg(feature = "relay")]
counts.insert(
"relay_events",
state.db.relay.events.approximate_len() as u64,
);
#[cfg(feature = "jetstream")]
counts.insert(
"jetstream_events",
state.db.jetstream.events.approximate_len() as u64,
);
let sizes = state
.db
.run(move |db| {
Ok(db
.all_keyspaces()
.into_iter()
.map(|(name, ks)| (name, ks.disk_space()))
.collect::<BTreeMap<_, _>>())
})
.await?;
Ok(StatsResponse { counts, sizes })
}
}
#[cfg(all(test, feature = "indexer"))]
mod tests {
use super::*;
use crate::config::Config;
use crate::db::types::{DbAction, DbRkey, DbTid};
use crate::state::AppState;
use jacquard_common::types::string::Did;
use std::time::Duration;
const COLLECTION: &str = "app.bsky.feed.post";
fn did() -> Did {
Did::new_static("did:plc:ewvi7nxzyoun6zhxrhs64oiz").unwrap()
}
fn rev_at(micros: i64) -> DbTid {
DbTid::new_from_bytes((((micros as u64) << 10) | 0x2A).to_be_bytes())
}
async fn hydrant(dir: &std::path::Path, history_ttl: Option<Duration>) -> Result<Hydrant> {
Hydrant::new(Config {
database_path: dir.to_path_buf(),
history_ttl,
..Config::default()
})
.await
}
fn put(state: &AppState, rev: &DbTid, rkey: &str, body: &'static str) -> Result<()> {
let did = did();
let mut txn = crate::db::Txn::new(&state.db);
let mut records = txn.records(state, rev, &did)?;
let block = crate::car::CarBlock::from_body(body.as_bytes());
records.put_record(COLLECTION, &DbRkey::new(rkey), &block, DbAction::Update)?;
txn.commit()
}
fn delete(state: &AppState, rev: &DbTid, rkey: &str) -> Result<()> {
let did = did();
let mut txn = crate::db::Txn::new(&state.db);
let mut records = txn.records(state, rev, &did)?;
records.delete_record(COLLECTION, &DbRkey::new(rkey))?;
txn.commit()
}
fn compact_history(state: &AppState) -> Result<()> {
let history = state
.db
.keyspace_by_name("history")
.ok_or_else(|| miette::miette!("no history keyspace"))?;
history.rotate_memtable_and_wait().into_diagnostic()?;
history.major_compact().into_diagnostic()
}
async fn history_count(hydrant: &Hydrant) -> Result<u64> {
Ok(hydrant.stats().await?.counts["history"])
}
#[tokio::test]
async fn history_count_drops_what_retention_expires() -> Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let hydrant = hydrant(tmp.path(), Some(Duration::from_secs(1))).await?;
let state = &hydrant.state;
let now = chrono::Utc::now().timestamp_micros();
put(state, &rev_at(now), "3kzbif5moe22m", "v1")?;
put(state, &rev_at(now + 1), "3kzbif5moe22m", "v2")?;
delete(state, &rev_at(now + 2), "3kzbif5moe22m")?;
assert_eq!(history_count(&hydrant).await?, 2);
tokio::time::sleep(Duration::from_millis(1100)).await;
compact_history(state)?;
assert_eq!(history_count(&hydrant).await?, 0);
Ok(())
}
#[tokio::test]
async fn history_count_follows_overwrites_and_purges() -> Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let hydrant = hydrant(tmp.path(), None).await?;
let state = &hydrant.state;
// a rev ahead of the local clock is its own death rev, so a second
// version replaced at that rev lands on the same history key
let ahead = rev_at(chrono::Utc::now().timestamp_micros() + 3_600_000_000);
put(state, &ahead, "3kzbif5moe22m", "v1")?;
put(state, &ahead, "3kzbif5moe22m", "v2")?;
put(state, &ahead, "3kzbif5moe22m", "v3")?;
compact_history(state)?;
assert_eq!(history_count(&hydrant).await?, 1);
put(state, &ahead, "3kzbif5mof33m", "other")?;
let mut txn = crate::db::Txn::new(&state.db);
crate::db::delete_repo_records(&mut txn, &did(), None, false)?;
txn.commit()?;
compact_history(state)?;
assert_eq!(history_count(&hydrant).await?, 3);
crate::db::erase_repo(&state.db, &did())?;
compact_history(state)?;
assert_eq!(history_count(&hydrant).await?, 0);
Ok(())
}
}