Something went wrong. Try again.
fork of hydrant
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394use 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 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 mod registry;pub mod schema;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;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>, #[cfg(test)] pub(crate) persist_failures: Arc<AtomicUsize>, /// 256 lock-sharded mutexes keyed by one byte of the trimmed DID. this is /// synchronization only, not storage partitioning: unrelated DIDs may /// conservatively share a lock. record mutations hold the selected lock /// from read through commit so operator redaction cannot race them. #[cfg(feature = "indexer")] pub(crate) record_write_locks: Arc<[std::sync::Mutex<()>; 256]>,}
/// lock index for a DID. skip the method tag so all PLC DIDs do not collapse/// onto one mutex.#[cfg(feature = "indexer")]pub(crate) fn record_lock_index_for_did(did: &jacquard_common::types::did::Did) -> u8 { let prefix = keys::record_prefix_did(did); match prefix.first() { Some(&types::TAG_PLC) | Some(&types::TAG_WEB) => prefix.get(1).copied().unwrap_or(0), Some(&b) => b, None => 0, }}
impl Db { /// runs synchronous database work without blocking the async runtime. 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(); tokio::task::spawn_blocking(move || f(&db)) .await .into_diagnostic()? }
pub fn persist(&self) -> Result<()> { #[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); } } let _guard = Guard(self.compaction_running.clone());
let compact = |name: &'static str, ks: Keyspace| async move { tracing::info!("compacting keyspace {name}..."); tokio::task::spawn_blocking(move || 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();
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()}
pub fn check_poisoned(e: &fjall::Error) { if matches!(e, fjall::Error::Poisoned) { error!("!!! DATABASE POISONED !!! exiting"); std::process::exit(10); }}
pub fn check_poisoned_report(e: &miette::Report) { let Some(err) = e.downcast_ref::<fjall::Error>() else { return; }; self::check_poisoned(err);}
/// 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;
#[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!(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(()) }}