use crate::types::{RepoMetadata, RepoState};
use fjall::{Database, Keyspace, PersistMode};
use miette::{Context, IntoDiagnostic, Result};
use scc::HashMap;
use smol_str::SmolStr;
use std::collections::BTreeSet;
use std::sync::atomic::AtomicU64;
#[cfg(test)]
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use tokio::sync::watch;
use url::Url;
pub mod compaction;
pub mod counts;
#[cfg(feature = "indexer")]
use counts::CountDeltaReservation;
pub use counts::{CountDeltas, load_count_delta_watermark, set_ks_count};
pub mod ephemeral;
pub mod filter;
pub mod keys;
#[cfg(feature = "indexer")]
mod lifecycle_counts;
pub mod migration;
pub mod pds_meta;
#[cfg(feature = "indexer")]
pub(crate) mod storage_mode;
pub mod types;
pub mod keyspaces;
mod open;
pub(crate) mod outbox;
pub mod registry;
pub mod schema;
#[cfg(any(feature = "indexer_stream", feature = "relay"))]
pub(crate) mod sequencer;
#[cfg(feature = "indexer_stream")]
pub(crate) mod stream_chunks;
mod train;
mod txn;
#[cfg(feature = "indexer")]
pub use keyspaces::IndexerDb;
#[cfg(feature = "jetstream")]
pub(crate) use keyspaces::JetstreamDb;
#[cfg(feature = "relay")]
pub(crate) use keyspaces::RelayDb;
#[cfg(feature = "indexer_stream")]
pub use keyspaces::StreamDb;
pub use schema::Ks;
#[cfg(feature = "indexer")]
use lifecycle_counts::LifecycleCountBatch;
#[cfg(all(test, feature = "indexer_stream"))]
pub(crate) use txn::RecordTxn;
pub(crate) use txn::{Txn, TxnCommitTimings};
use tracing::error;
#[derive(Clone)]
pub struct Db {
pub inner: Arc<Database>,
pub path: std::path::PathBuf,
pub repos: Ks<schema::Repos>,
pub repo_metadata: Ks<schema::RepoMetadata>,
pub cursors: Ks<schema::Cursors>,
pub counts: Ks<schema::Counts>,
pub filter: Ks<schema::Filter>,
pub crawler: Ks<schema::Crawler>,
#[cfg(feature = "indexer")]
pub indexer: IndexerDb,
#[cfg(feature = "indexer_stream")]
pub stream: StreamDb,
#[cfg(feature = "jetstream")]
pub(crate) jetstream: JetstreamDb,
#[cfg(feature = "relay")]
pub(crate) relay: RelayDb,
#[cfg(feature = "backlinks")]
pub(crate) backlinks: Ks<schema::Backlinks>,
pub counts_map: Arc<HashMap<SmolStr, u64>>,
next_count_delta_id: Arc<AtomicU64>,
count_delta_checkpoint_watermark: Arc<AtomicU64>,
count_delta_gc_watermark: Arc<AtomicU64>,
count_delta_in_flight: Arc<Mutex<BTreeSet<u64>>>,
pub(crate) compaction_running: Arc<std::sync::atomic::AtomicBool>,
pub(crate) poison: Poison,
#[cfg(test)]
pub(crate) persist_failures: Arc<AtomicUsize>,
/// runs as [`Self::persist`] starts, so a test can see what it's called under
#[cfg(test)]
pub(crate) before_persist: Arc<parking_lot::Mutex<Option<TestHook>>>,
/// a compaction's blocking work waits for this lock, so a test can keep it running
#[cfg(test)]
pub(crate) compaction_hold: Arc<Mutex<()>>,
/// the only source of stream positions, see [`sequencer`]
#[cfg(any(feature = "indexer_stream", feature = "relay"))]
pub(crate) sequencer: Arc<sequencer::Sequencer>,
}
/// [`Db::run`] work that never ran, because its runtime shut down first
#[derive(Debug, miette::Diagnostic, thiserror::Error)]
#[error("the runtime shut down before this database work ran")]
pub(crate) struct Cancelled;
#[cfg(test)]
pub(crate) type TestHook = Arc<dyn Fn() -> Result<()> + Send + Sync>;
impl Db {
/// runs synchronous database work without blocking the async runtime. fails with
/// [`Cancelled`] if the runtime shut down before the work started
pub async fn run<T, F>(&self, f: F) -> Result<T>
where
T: Send + 'static,
F: FnOnce(&Db) -> Result<T> + Send + 'static,
{
let db = self.clone();
match tokio::task::spawn_blocking(move || f(&db)).await {
Ok(result) => result,
// nothing holds the handle to abort it, so only a runtime shutting down cancels it
Err(e) if e.is_cancelled() => Err(Cancelled.into()),
Err(e) => Err(e).into_diagnostic(),
}
}
pub fn persist(&self) -> Result<()> {
#[cfg(test)]
if let Some(before_persist) = self.before_persist.lock().clone() {
before_persist()?;
}
#[cfg(test)]
if self
.persist_failures
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |remaining| {
if remaining > 0 {
Some(remaining - 1)
} else {
None
}
})
.is_ok()
{
return Err(miette::miette!("injected persist failure"));
}
#[cfg(not(feature = "__persist_sync_all"))]
const MODE: PersistMode = PersistMode::Buffer;
#[cfg(feature = "__persist_sync_all")]
const MODE: PersistMode = PersistMode::SyncAll;
self.inner.persist(MODE).into_diagnostic()?;
Ok(())
}
#[cfg(test)]
pub(crate) fn fail_next_persist_for_test(&self) {
self.persist_failures.fetch_add(1, Ordering::Release);
}
pub async fn compact(&self) -> Result<()> {
use std::sync::atomic::Ordering;
if self
.compaction_running
.compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
.is_err()
{
return Err(miette::miette!("compaction already in progress"));
}
struct Guard(Arc<std::sync::atomic::AtomicBool>);
impl Drop for Guard {
fn drop(&mut self) {
self.0.store(false, Ordering::Release);
}
}
// every blocking compaction holds a clone, since dropping this future (a maintenance
// timeout does) can't stop them, and a second compaction must wait until they end
let guard = Arc::new(Guard(self.compaction_running.clone()));
let compact = |name: &'static str, ks: Keyspace| {
let guard = guard.clone();
#[cfg(test)]
let hold = self.compaction_hold.clone();
async move {
tracing::info!("compacting keyspace {name}...");
tokio::task::spawn_blocking(move || {
let _guard = guard;
#[cfg(test)]
drop(hold.lock());
ks.major_compact().into_diagnostic()
})
.await
.into_diagnostic()??;
tracing::info!("compacted keyspace {name}");
Ok::<(), miette::Report>(())
}
};
let tasks: Vec<_> = self
.all_keyspaces()
.into_iter()
.map(|(name, ks)| compact(name, ks))
.collect();
drop(guard);
futures::future::try_join_all(tasks).await?;
Ok(())
}
}
#[cfg(feature = "indexer")]
mod indexer;
#[cfg(feature = "indexer")]
pub use indexer::*;
#[derive(serde::Serialize, serde::Deserialize, Default)]
pub(crate) struct FirehoseSourceMeta {
#[serde(default)]
pub(crate) is_pds: bool,
}
pub fn set_firehose_cursor(db: &Db, relay: &Url, cursor: i64) -> Result<()> {
db.cursors
.insert(
keys::firehose_cursor_key_from_url(relay),
cursor.to_be_bytes(),
)
.into_diagnostic()
}
pub async fn get_firehose_cursor(db: &Db, relay: &Url) -> Result<Option<i64>> {
let key = keys::firehose_cursor_key_from_url(relay);
db.run(move |db| db.cursors.get(key).into_diagnostic())
.await?
.map(|v| {
Ok(i64::from_be_bytes(
v.as_ref()
.try_into()
.into_diagnostic()
.wrap_err("cursor is not 8 bytes")?,
))
})
.transpose()
}
#[cfg(feature = "indexer")]
pub fn ser_repo_meta(state: &RepoMetadata) -> Result<Vec<u8>> {
rmp_serde::to_vec(&state).into_diagnostic()
}
pub fn deser_repo_meta(bytes: &[u8]) -> Result<RepoMetadata> {
rmp_serde::from_slice(bytes).into_diagnostic()
}
pub fn ser_repo_state(state: &RepoState) -> Result<Vec<u8>> {
rmp_serde::to_vec(&state).into_diagnostic()
}
pub fn deser_repo_state<'b>(bytes: &'b [u8]) -> Result<RepoState<'b>> {
rmp_serde::from_slice(bytes).into_diagnostic()
}
/// fjall refuses every write once a disk write failed, and the first error saying so
/// ends [`Hydrant::run`](crate::control::Hydrant::run) instead of the whole process,
/// since hydrant may share it with the program embedding it.
#[derive(Clone)]
pub struct Poison(watch::Sender<bool>);
impl Default for Poison {
fn default() -> Self {
Self(watch::Sender::new(false))
}
}
impl Poison {
pub fn check(&self, e: &fjall::Error) {
if matches!(e, fjall::Error::Poisoned) && !self.0.send_replace(true) {
error!("!!! DATABASE POISONED !!! stopping");
}
}
pub fn check_report(&self, e: &miette::Report) {
// `into_diagnostic` hides the fjall error behind a private wrapper that only
// passes its message on, so a downcast never finds it
let poisoned = fjall::Error::Poisoned.to_string();
if e.chain().any(|cause| cause.to_string() == poisoned) {
self.check(&fjall::Error::Poisoned);
}
}
/// resolves once the database is poisoned.
pub async fn wait(&self) {
// only errs once the sender is dropped, and `self` holds it
let _ = self.0.subscribe().wait_for(|poisoned| *poisoned).await;
}
}
/// Load the persisted `(day, count)` pair for the daily PDS add counter.
/// Returns `None` if no entry exists or the stored data is malformed.
pub fn load_pds_daily_adds(db: &Db) -> Result<Option<(u64, u64)>> {
let Some(val) = db.cursors.get(keys::PDS_DAILY_ADDS_KEY).into_diagnostic()? else {
return Ok(None);
};
if val.len() < 16 {
miette::bail!("malformed pds daily limit value");
}
let day = u64::from_be_bytes(val[..8].try_into().into_diagnostic()?);
let count = u64::from_be_bytes(val[8..].try_into().into_diagnostic()?);
Ok(Some((day, count)))
}
/// Stage the daily PDS add counter alongside the source it admits.
pub(crate) fn stage_pds_daily_adds(
batch: &mut fjall::OwnedWriteBatch,
db: &Db,
day: u64,
count: u64,
) {
let mut value = [0u8; 16];
value[..8].copy_from_slice(&day.to_be_bytes());
value[8..].copy_from_slice(&count.to_be_bytes());
batch.insert(&db.cursors, keys::PDS_DAILY_ADDS_KEY, value);
}
/// Loads persisted sources and their explicit trusted-admin provenance bit.
/// A missing marker is deliberately public, fail-closed on legacy rows.
pub fn load_persisted_firehose_sources(
db: &crate::db::Db,
) -> Result<Vec<(crate::config::FirehoseSource, bool)>> {
use crate::db::keys::{FIREHOSE_SOURCE_PREFIX, trusted_firehose_source_key};
let mut sources = Vec::new();
for entry in db.crawler.prefix(FIREHOSE_SOURCE_PREFIX) {
let (key, val) = entry.into_inner().into_diagnostic()?;
let url_bytes = &key[FIREHOSE_SOURCE_PREFIX.len()..];
let url_str = std::str::from_utf8(url_bytes).into_diagnostic()?;
let url = Url::parse(url_str).into_diagnostic()?;
let meta: FirehoseSourceMeta = rmp_serde::from_slice(&val)
.map_err(|e| miette::miette!("failed to deserialize firehose source meta: {e}"))?;
let trusted = db
.crawler
.contains_key(trusted_firehose_source_key(url_str))
.into_diagnostic()?;
sources.push((
crate::config::FirehoseSource {
url,
is_pds: meta.is_pds,
},
trusted,
));
}
Ok(sources)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::Config;
use std::sync::atomic::Ordering;
// holding the gate across the awaits is the point, it keeps the compaction's blocking work stuck
#[allow(clippy::await_holding_lock)]
#[tokio::test(flavor = "multi_thread")]
async fn compaction_stays_running_after_its_caller_gives_up() {
let tmp = tempfile::tempdir().unwrap();
let db = Db::open(&Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
})
.unwrap();
let hold = db.compaction_hold.clone();
let held = hold.lock().unwrap();
// what a maintenance timeout does: drop the compaction while its blocking work runs
let timed_out =
tokio::time::timeout(std::time::Duration::from_millis(200), db.compact()).await;
assert!(timed_out.is_err(), "compaction finished while held");
// without the guard this second one starts and blocks on the hold too
let second = tokio::time::timeout(std::time::Duration::from_secs(5), db.compact())
.await
.expect("a second compaction started while the first was still running");
let err = second.unwrap_err();
assert!(err.to_string().contains("already in progress"), "{err}");
drop(held);
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
while db.compaction_running.load(Ordering::Acquire) {
assert!(
std::time::Instant::now() < deadline,
"compaction never ended"
);
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
db.compact().await.unwrap();
}
#[tokio::test]
async fn only_a_poisoned_error_trips_the_poison() {
let poison = Poison::default();
let wait = std::time::Duration::from_millis(50);
poison.check(&fjall::Error::Locked);
poison.check_report(&miette::miette!("not a fjall error"));
assert!(tokio::time::timeout(wait, poison.wait()).await.is_err());
// wrapped the way callsites hand it over, through `into_diagnostic` and `context`
let report = Err::<(), _>(fjall::Error::Poisoned)
.into_diagnostic()
.wrap_err("committing a batch")
.unwrap_err();
poison.clone().check_report(&report);
tokio::time::timeout(std::time::Duration::from_secs(5), poison.wait())
.await
.expect("a poisoned error should trip every clone");
}
#[test]
fn work_a_shutdown_cut_off_says_so() -> Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let db = Db::open(&Config {
database_path: tmp.path().to_path_buf(),
..Config::default()
})?;
let runtime = tokio::runtime::Runtime::new().into_diagnostic()?;
let handle = runtime.handle().clone();
runtime.shutdown_background();
let err = handle.block_on(db.run(|_| Ok(()))).unwrap_err();
assert!(err.downcast_ref::<Cancelled>().is_some(), "{err}");
Ok(())
}
#[tokio::test]
async fn test_db_compact_concurrency_guard() -> Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let cfg = Config {
database_path: tmp.path().to_path_buf(),
..Config::default()
};
let db = Db::open(&cfg)?;
// Manually mark compaction as running
db.compaction_running.store(true, Ordering::SeqCst);
// Attempting to compact should now fail
let res = db.compact().await;
assert!(res.is_err());
assert!(res.unwrap_err().to_string().contains("already in progress"));
// Release the lock
db.compaction_running.store(false, Ordering::SeqCst);
// Compacting should now succeed
let res = db.compact().await;
assert!(res.is_ok());
Ok(())
}
#[tokio::test]
async fn legacy_firehose_row_without_trusted_marker_is_public() -> Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let config = Config {
database_path: tmp.path().to_path_buf(),
..Config::default()
};
let db = Db::open(&config)?;
let url = "wss://legacy-pds.example/";
let value = rmp_serde::to_vec(&FirehoseSourceMeta { is_pds: true }).into_diagnostic()?;
db.crawler
.insert(keys::firehose_source_key(url), value)
.into_diagnostic()?;
let sources = load_persisted_firehose_sources(&db)?;
assert_eq!(sources.len(), 1);
assert_eq!(sources[0].0.url, Url::parse(url).unwrap());
assert!(!sources[0].1, "missing marker must fail closed to public");
Ok(())
}
#[test]
fn firehose_source_meta_golden_fixture() {
let encoded = rmp_serde::to_vec(&FirehoseSourceMeta { is_pds: true }).unwrap();
assert_eq!(crate::util::hex::encode(&encoded), "91c3");
let decoded: FirehoseSourceMeta = rmp_serde::from_slice(&encoded).unwrap();
assert!(decoded.is_pds);
let legacy_default: FirehoseSourceMeta = rmp_serde::from_slice(&[0x90]).unwrap();
assert!(!legacy_default.is_pds);
}
#[tokio::test]
async fn firehose_cursor_big_endian_roundtrip() -> Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let cfg = Config {
database_path: tmp.path().to_path_buf(),
..Config::default()
};
let db = Db::open(&cfg)?;
let relay = Url::parse("https://relay.example").into_diagnostic()?;
let cursor = i64::from_be_bytes(*b"stinkpot");
set_firehose_cursor(&db, &relay, cursor)?;
let raw = db
.cursors
.get(keys::firehose_cursor_key_from_url(&relay))
.into_diagnostic()?
.unwrap();
assert_eq!(raw.as_ref(), b"stinkpot");
assert_eq!(get_firehose_cursor(&db, &relay).await?, Some(cursor));
Ok(())
}
}