diff --git a/examples/statusphere.rs b/examples/statusphere.rs new file mode 100644 --- /dev/null +++ b/examples/statusphere.rs @@ -0,0 +1,171 @@ +//! a statusphere indexer: tracks xyz.statusphere.status records across the ATProto network. +//! +//! statusphere is a demo app where users set a single-emoji status on their bluesky profile. +//! this example indexes those status records in real time, maintaining the current status +//! per user and printing a periodic leaderboard of the top emoji statuses in use. +//! +//! see: https://github.com/bluesky-social/statusphere-example-app +//! +//! run with: +//! HYDRANT_DATABASE_PATH=./statusphere.db cargo run --example statusphere +//! +//! the database persists records across restarts. on each start the full event +//! history is replayed from the database to rebuild the in-memory index. + +use std::sync::Arc; +use std::time::Duration; + +use futures::StreamExt; +use hydrant::config::Config; +use hydrant::control::{EventStream, Hydrant, ReposControl}; +use hydrant::filter::FilterMode; +use scc::HashMap; + +const COLLECTION: &str = "xyz.statusphere.status"; + +struct StatusEntry { + emoji: String, + created_at: String, +} + +struct StatusIndex { + /// current status per DID: only the latest by createdAt is kept. + current: HashMap, +} + +impl StatusIndex { + fn new() -> Self { + Self { + current: HashMap::new(), + } + } + + fn set(&self, did: String, emoji: String, created_at: String) -> bool { + let is_newer = self + .current + .read_sync(&did, |_, e| created_at > e.created_at) + .unwrap_or(true); + if is_newer { + self.current + .upsert_sync(did, StatusEntry { emoji, created_at }); + } + is_newer + } + + fn delete(&self, did: &str) { + self.current.remove_sync(did); + } + + fn top(&self, n: usize) -> Vec<(String, usize)> { + use std::collections::HashMap; + let mut counts: HashMap = HashMap::with_capacity(self.current.capacity()); + self.current.iter_sync(|_, e| { + *counts.entry(e.emoji.clone()).or_default() += 1; + true + }); + let mut ranked: Vec<_> = counts.into_iter().collect(); + ranked.sort_by(|a, b| b.1.cmp(&a.1)); + ranked.truncate(n); + ranked + } +} + +async fn run_ticker(index: Arc) { + let mut interval = tokio::time::interval(Duration::from_secs(30)); + interval.tick().await; + loop { + interval.tick().await; + let top = index.top(10); + if top.is_empty() { + continue; + } + println!( + "\n--- top statuses ({} users tracked) ---", + index.current.len() + ); + for (emoji, count) in &top { + println!(" {emoji} ×{count}"); + } + println!("----------------------------------------\n"); + } +} + +async fn handle_stream(index: Arc, repos: ReposControl, mut stream: EventStream) { + while let Some(event) = stream.next().await { + if let Some(rec) = event.record { + let did = rec.did.as_str().to_owned(); + match rec.action.as_str() { + "create" | "update" => { + let Some(record) = rec.record else { continue }; + let Some(emoji) = record + .get("status") + .and_then(|v| v.as_str()) + .map(|s| s.to_owned()) + else { + continue; + }; + let created_at = record + .get("createdAt") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_owned(); + if index.set(did.clone(), emoji.clone(), created_at) { + let name = repos + .get(&rec.did) + .await + .ok() + .flatten() + .and_then(|info| info.handle) + .unwrap_or(did); + println!("[{}] {name} set status: {emoji}", event.id); + } + } + "delete" => { + let name = repos + .get(&rec.did) + .await + .ok() + .flatten() + .and_then(|info| info.handle) + .unwrap_or(did.clone()); + index.delete(&did); + println!("[{}] {name} cleared status", event.id); + } + _ => {} + } + } else if let Some(account) = event.account { + // when an account is deactivated or deleted, drop their status. + if !account.active { + index.delete(account.did.as_str()); + } + } + } +} + +#[tokio::main] +async fn main() -> miette::Result<()> { + tracing_subscriber::fmt() + .with_env_filter("hydrant=info") + .init(); + + let cfg = Config::from_env()?; + let hydrant = Hydrant::new(cfg).await?; + + // discover only repos that publish xyz.statusphere.status records, + // and only store that collection (all other record types are dropped). + hydrant.filter.set_mode(FilterMode::Filter).await?; + hydrant.filter.set_signals([COLLECTION]).await?; + hydrant.filter.set_collections([COLLECTION]).await?; + + // replay all persisted events from the start to rebuild the in-memory index, + // then switch to live tail. since the index is in-memory, we always need the + // full replay on startup. + let stream = hydrant.subscribe(Some(0)); + + let index = Arc::new(StatusIndex::new()); + tokio::select! { + r = hydrant.run() => r, + _ = run_ticker(index.clone()) => Ok(()), + _ = handle_stream(index.clone(), hydrant.repos.clone(), stream) => Ok(()), + } +} diff --git a/src/config.rs b/src/config.rs --- a/src/config.rs +++ b/src/config.rs @@ -73,13 +73,10 @@ pub ephemeral_ttl: Duration, pub cursor_save_interval: Duration, pub repo_fetch_timeout: Duration, - pub api_port: u16, pub cache_size: u64, pub backfill_concurrency_limit: usize, pub data_compression: Compression, pub journal_compression: Compression, - pub debug_port: u16, - pub enable_debug: bool, pub verify_signatures: SignatureVerification, pub identity_cache_size: u64, pub enable_firehose: bool, @@ -167,10 +164,6 @@ let data_compression = cfg!("DATA_COMPRESSION", Compression::Lz4); let journal_compression = cfg!("JOURNAL_COMPRESSION", Compression::Lz4); - let api_port = cfg!("API_PORT", 3000u16); - let enable_debug = cfg!("ENABLE_DEBUG", false); - let debug_port: u16 = api_port + 1; - let debug_port = cfg!("DEBUG_PORT", debug_port); let verify_signatures = cfg!("VERIFY_SIGNATURES", SignatureVerification::Full); let identity_cache_size = cfg!("IDENTITY_CACHE_SIZE", 1_000_000u64); let enable_firehose = cfg!("ENABLE_FIREHOSE", true); @@ -245,13 +238,10 @@ full_network, cursor_save_interval, repo_fetch_timeout, - api_port, cache_size, backfill_concurrency_limit, data_compression, journal_compression, - debug_port, - enable_debug, verify_signatures, identity_cache_size, enable_firehose, @@ -269,6 +259,34 @@ filter_collections, filter_excludes, }) + } +} + +#[derive(Debug, Clone)] +pub struct AppConfig { + pub api_port: u16, + pub enable_debug: bool, + pub debug_port: u16, +} + +impl AppConfig { + pub fn from_env() -> Self { + macro_rules! cfg { + ($key:expr, $default:expr) => { + std::env::var(concat!("HYDRANT_", $key)) + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or($default) + }; + } + let api_port = cfg!("API_PORT", 3000u16); + let enable_debug = cfg!("ENABLE_DEBUG", false); + let debug_port = cfg!("DEBUG_PORT", api_port + 1); + Self { + api_port, + enable_debug, + debug_port, + } } } @@ -304,7 +322,6 @@ config_line!(f, "cache size", format_args!("{} mb", self.cache_size))?; config_line!(f, "data compression", self.data_compression)?; config_line!(f, "journal compression", self.journal_compression)?; - config_line!(f, "api port", self.api_port)?; config_line!(f, "firehose workers", self.firehose_workers)?; config_line!(f, "db worker threads", self.db_worker_threads)?; config_line!( @@ -346,10 +363,6 @@ } if let Some(excludes) = &self.filter_excludes { config_line!(f, "filter excludes", format_args!("{:?}", excludes))?; - } - config_line!(f, "enable debug", self.enable_debug)?; - if self.enable_debug { - config_line!(f, "debug port", self.debug_port)?; } Ok(()) } diff --git a/src/control.rs b/src/control.rs new file mode 100644 --- /dev/null +++ b/src/control.rs @@ -0,0 +1,1251 @@ +use std::collections::BTreeMap; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::task::{Context, Poll}; + +use chrono::{DateTime, Utc}; +use futures::{FutureExt, Stream}; +use jacquard_common::types::cid::{ATP_CID_HASH, IpldCid}; +use jacquard_common::types::string::Did; +use jacquard_common::{CowStr, IntoStatic, RawData}; +use jacquard_repo::DAG_CBOR_CID_CODEC; +use miette::{IntoDiagnostic, Result}; +use rand::Rng; +use sha2::{Digest, Sha256}; +use tokio::sync::{mpsc, watch}; +use tracing::{debug, error, info}; + +use crate::backfill::BackfillWorker; +use crate::config::{Config, SignatureVerification}; +use crate::crawler::Crawler; +use crate::db::{self, filter as db_filter, keys, ser_repo_state}; +use crate::filter::{FilterMode, SetUpdate}; +use crate::ingest::{firehose::FirehoseIngestor, worker::FirehoseWorker}; +use crate::state::AppState; +use crate::types::{ + BroadcastEvent, GaugeState, MarshallableEvt, RecordEvt, RepoState, StoredData, StoredEvent, +}; + +/// an event emitted by the hydrant event stream. +/// +/// three variants are possible depending on the `type` field: +/// - `"record"`: a repo record was created, updated, or deleted. carries a [`RecordEvt`]. +/// - `"identity"`: a DID's handle or PDS changed. carries an [`IdentityEvt`]. ephemeral, not replayable. +/// - `"account"`: a repo's active/inactive status changed. carries an [`AccountEvt`]. ephemeral, not replayable. +/// +/// the `id` field is a monotonically increasing sequence number usable as a cursor for [`Hydrant::subscribe`]. +pub type Event = MarshallableEvt<'static>; + +/// the top-level handle to a hydrant instance. +/// +/// `Hydrant` is cheaply cloneable. all sub-handles share the same underlying state. +/// construct it via [`Hydrant::new`] or [`Hydrant::from_env`], configure the filter +/// and repos as needed, then call [`Hydrant::run`] to start all background components. +/// +/// # example +/// +/// ```rust,no_run +/// use hydrant::control::Hydrant; +/// +/// #[tokio::main] +/// async fn main() -> miette::Result<()> { +/// let hydrant = Hydrant::from_env().await?; +/// +/// tokio::select! { +/// r = hydrant.run() => r, +/// r = hydrant.serve(3000) => r, +/// } +/// } +/// ``` +#[derive(Clone)] +pub struct Hydrant { + pub crawler: CrawlerHandle, + pub firehose: FirehoseHandle, + pub backfill: BackfillHandle, + pub filter: FilterControl, + pub repos: ReposControl, + pub db: DbControl, + pub(crate) state: Arc, + config: Arc, + started: Arc, + _priv: (), +} + +impl Hydrant { + /// open the database and configure hydrant from `config`. + /// + /// this sets up the database, applies any filter configuration from `config`, and + /// initializes all sub-handles. no background tasks are started yet: call + /// [`run`](Self::run) to start all components and drive the instance. + pub async fn new(config: Config) -> Result { + info!("{config}"); + + // 1. open database and construct AppState + let state = AppState::new(&config)?; + + // 2. apply any filter config from env variables + if config.full_network + || config.filter_signals.is_some() + || config.filter_collections.is_some() + || config.filter_excludes.is_some() + { + let filter_ks = state.db.filter.clone(); + let inner = state.db.inner.clone(); + let mode = config.full_network.then_some(FilterMode::Full); + let signals = config.filter_signals.clone().map(SetUpdate::Set); + let collections = config.filter_collections.clone().map(SetUpdate::Set); + let excludes = config.filter_excludes.clone().map(SetUpdate::Set); + + tokio::task::spawn_blocking(move || { + let mut batch = inner.batch(); + db_filter::apply_patch( + &mut batch, + &filter_ks, + mode, + signals, + collections, + excludes, + )?; + batch.commit().into_diagnostic() + }) + .await + .into_diagnostic()??; + + // 3. reload the live filter into the hot-path arc-swap + let new_filter = tokio::task::spawn_blocking({ + let filter_ks = state.db.filter.clone(); + move || db_filter::load(&filter_ks) + }) + .await + .into_diagnostic()??; + state.filter.store(Arc::new(new_filter)); + } + + // 4. set crawler enabled state from config, evaluated against the post-patch filter + let post_patch_crawler = match config.enable_crawler { + Some(b) => b, + None => state.filter.load().mode == FilterMode::Full, + }; + state.crawler_enabled.send_replace(post_patch_crawler); + + let state = Arc::new(state); + + Ok(Self { + crawler: CrawlerHandle(state.clone()), + firehose: FirehoseHandle(state.clone()), + backfill: BackfillHandle(state.clone()), + filter: FilterControl(state.clone()), + repos: ReposControl(state.clone()), + db: DbControl(state.clone()), + state, + config: Arc::new(config), + started: Arc::new(AtomicBool::new(false)), + _priv: (), + }) + } + + /// reads config from environment variables and calls [`Hydrant::new`]. + pub async fn from_env() -> Result { + Self::new(Config::from_env()?).await + } + + /// start all background components and return a future that resolves when any + /// fatal component exits. + /// + /// starts the backfill worker, firehose ingestors, crawler, and worker thread. + /// resolves with `Ok(())` if a fatal component exits cleanly, or `Err(e)` if it + /// fails. intended for use in `tokio::select!` alongside [`serve`](Self::serve). + /// + /// panics if called more than once on the same `Hydrant` instance. + pub fn run(&self) -> impl Future> { + let state = self.state.clone(); + let config = self.config.clone(); + let started = self.started.clone(); + + async move { + if started.swap(true, Ordering::SeqCst) { + panic!("Hydrant::run() called more than once"); + } + + // internal buffered channel between ingestors / backfill and the firehose worker + let (buffer_tx, buffer_rx) = mpsc::unbounded_channel(); + + // 5. spawn the backfill worker + tokio::spawn({ + let state = state.clone(); + BackfillWorker::new( + state.clone(), + buffer_tx.clone(), + config.repo_fetch_timeout, + config.backfill_concurrency_limit, + matches!( + config.verify_signatures, + SignatureVerification::Full | SignatureVerification::BackfillOnly + ), + config.ephemeral, + state.backfill_enabled.subscribe(), + ) + .run() + }); + + // 6. re-queue any repos that lost their backfill state, then start the retry worker + if let Err(e) = tokio::task::spawn_blocking({ + let state = state.clone(); + move || crate::backfill::manager::queue_gone_backfills(&state) + }) + .await + .into_diagnostic()? + { + error!(err = %e, "failed to queue gone backfills"); + db::check_poisoned_report(&e); + } + + std::thread::spawn({ + let state = state.clone(); + move || crate::backfill::manager::retry_worker(state) + }); + + // 7. ephemeral GC thread + if config.ephemeral { + let state = state.clone(); + std::thread::Builder::new() + .name("ephemeral-gc".into()) + .spawn(move || crate::db::ephemeral::ephemeral_ttl_worker(state)) + .into_diagnostic()?; + } + + // 8. cursor / counts persist thread + std::thread::spawn({ + let state = state.clone(); + let persist_interval = config.cursor_save_interval; + move || loop { + std::thread::sleep(persist_interval); + + for (relay, cursor) in &state.relay_cursors { + let seq = cursor.load(Ordering::SeqCst); + if seq > 0 { + if let Err(e) = db::set_firehose_cursor(&state.db, relay, seq) { + error!(relay = %relay, err = %e, "failed to save cursor"); + db::check_poisoned_report(&e); + } + } + } + + if let Err(e) = db::persist_counts(&state.db) { + error!(err = %e, "failed to persist counts"); + db::check_poisoned_report(&e); + } + + if let Err(e) = state.db.persist() { + error!(err = %e, "db persist failed"); + db::check_poisoned_report(&e); + } + } + }); + + // 9. events/sec stats ticker + tokio::spawn({ + let state = state.clone(); + let mut last_id = state.db.next_event_id.load(Ordering::Relaxed); + let mut last_time = std::time::Instant::now(); + let mut interval = tokio::time::interval(std::time::Duration::from_secs(60)); + async move { + loop { + interval.tick().await; + + let current_id = state.db.next_event_id.load(Ordering::Relaxed); + let current_time = std::time::Instant::now(); + let delta = current_id.saturating_sub(last_id); + + if delta == 0 { + debug!("no new events in 60s"); + continue; + } + + let elapsed = current_time.duration_since(last_time).as_secs_f64(); + let rate = if elapsed > 0.0 { + delta as f64 / elapsed + } else { + 0.0 + }; + info!("{rate:.2} events/s ({delta} events in {elapsed:.1}s)"); + + last_id = current_id; + last_time = current_time; + } + } + }); + + let (fatal_tx_inner, mut fatal_rx) = watch::channel(None); + let fatal_tx = Arc::new(fatal_tx_inner); + + info!( + crawler_enabled = *state.crawler_enabled.borrow(), + firehose_enabled = *state.firehose_enabled.borrow(), + filter_mode = ?state.filter.load().mode, + "starting ingestion" + ); + + // 10. spawn one firehose ingestor per relay (fatal tasks) + let relay_hosts = config.relays.clone(); + if !relay_hosts.is_empty() { + info!( + relay_count = relay_hosts.len(), + hosts = relay_hosts + .iter() + .map(|h| h.as_str()) + .collect::>() + .join(", "), + "starting firehose ingestor(s)" + ); + for relay_url in &relay_hosts { + let ingestor = FirehoseIngestor::new( + state.clone(), + buffer_tx.clone(), + relay_url.clone(), + state.filter.clone(), + state.firehose_enabled.subscribe(), + matches!(config.verify_signatures, SignatureVerification::Full), + ); + let tx = Arc::clone(&fatal_tx); + tokio::spawn(async move { + let result = ingestor.run().await; + let _ = tx.send(Some(result.map_err(|e| e.to_string()))); + }); + } + } + + // 11. spawn the crawler if we have relay hosts to crawl + if !relay_hosts.is_empty() { + let crawler_rx = state.crawler_enabled.subscribe(); + info!( + relay_count = relay_hosts.len(), + hosts = relay_hosts + .iter() + .map(|h| h.as_str()) + .collect::>() + .join(", "), + enabled = *state.crawler_enabled.borrow(), + "starting crawler(s)" + ); + let state = state.clone(); + let max_pending = config.crawler_max_pending_repos; + let resume_pending = config.crawler_resume_pending_repos; + tokio::spawn(async move { + let crawler = + Crawler::new(state, relay_hosts, max_pending, resume_pending, crawler_rx); + if let Err(e) = crawler.run().await { + error!(err = %e, "crawler error"); + db::check_poisoned_report(&e); + } + }); + } + + // 12. spawn the firehose worker on a blocking thread (fatal task) + let handle = tokio::runtime::Handle::current(); + let firehose_worker = std::thread::spawn({ + let state = state.clone(); + move || { + FirehoseWorker::new( + state, + buffer_rx, + matches!(config.verify_signatures, SignatureVerification::Full), + config.ephemeral, + config.firehose_workers, + ) + .run(handle) + } + }); + + { + let tx = Arc::clone(&fatal_tx); + tokio::spawn( + tokio::task::spawn_blocking(move || { + firehose_worker + .join() + .map_err(|e| miette::miette!("buffer processor died: {e:?}")) + }) + .map(move |r| { + let result = r.into_diagnostic().flatten().flatten(); + let _ = tx.send(Some(result.map_err(|e| e.to_string()))); + }), + ); + } + + // drop the local fatal_tx so the watch channel is only kept alive by the + // spawned tasks. when all fatal tasks exit (and drop their tx clones), + // fatal_rx.changed() returns Err and we return Ok(()). + drop(fatal_tx); + + loop { + match fatal_rx.changed().await { + Ok(()) => { + if let Some(result) = fatal_rx.borrow().clone() { + return result.map_err(|s| miette::miette!("{s}")); + } + } + // all fatal_tx clones dropped: all tasks finished cleanly + Err(_) => return Ok(()), + } + } + } + } + + /// subscribe to the ordered event stream. + /// + /// returns an [`EventStream`] that implements [`futures::Stream`]. + /// + /// - if `cursor` is `None`, streaming starts from the current head (live tail only). + /// - if `cursor` is `Some(id)`, all persisted `record` events from that ID onward are + /// replayed first, then live events follow seamlessly. + /// + /// `identity` and `account` events are ephemeral and are never replayed from a cursor - + /// only live occurrences are delivered. use [`ReposControl::get`] to fetch current + /// identity/account state for a specific DID. + /// + /// multiple concurrent subscribers each receive a full independent copy of the stream. + /// the stream ends when the `EventStream` is dropped. + pub fn subscribe(&self, cursor: Option) -> EventStream { + let (tx, rx) = mpsc::channel(500); + let state = self.state.clone(); + let runtime = tokio::runtime::Handle::current(); + + std::thread::Builder::new() + .name("hydrant-stream".into()) + .spawn(move || { + let _g = runtime.enter(); + event_stream_thread(state, tx, cursor); + }) + .expect("failed to spawn stream thread"); + + EventStream(rx) + } + + /// return database counts and on-disk sizes for all keyspaces. + /// + /// counts include: `repos`, `pending`, `resync`, `records`, `blocks`, `events`, + /// `error_ratelimited`, `error_transport`, `error_generic`. + /// + /// sizes are in bytes, reported per keyspace. + pub async fn stats(&self) -> Result { + let db = self.state.db.clone(); + + let mut counts: BTreeMap<&'static str, u64> = futures::future::join_all( + [ + "repos", + "pending", + "resync", + "records", + "blocks", + "error_ratelimited", + "error_transport", + "error_generic", + ] + .into_iter() + .map(|name| { + let db = db.clone(); + async move { (name, db.get_count(name).await) } + }), + ) + .await + .into_iter() + .collect(); + + counts.insert("events", db.events.approximate_len() as u64); + + let sizes = tokio::task::spawn_blocking(move || { + let mut s = BTreeMap::new(); + s.insert("repos", db.repos.disk_space()); + s.insert("records", db.records.disk_space()); + s.insert("blocks", db.blocks.disk_space()); + s.insert("cursors", db.cursors.disk_space()); + s.insert("pending", db.pending.disk_space()); + s.insert("resync", db.resync.disk_space()); + s.insert("resync_buffer", db.resync_buffer.disk_space()); + s.insert("events", db.events.disk_space()); + s.insert("counts", db.counts.disk_space()); + s.insert("filter", db.filter.disk_space()); + s.insert("crawler", db.crawler.disk_space()); + s + }) + .await + .into_diagnostic()?; + + Ok(StatsResponse { counts, sizes }) + } + + /// returns a future that runs the HTTP management API server on `0.0.0.0:{port}`. + /// + /// the server exposes all management endpoints (`/filter`, `/repos`, `/ingestion`, + /// `/stream`, `/stats`, `/db/*`, `/xrpc/*`). it runs indefinitely and resolves + /// only on error. + /// + /// intended for `tokio::spawn` or inclusion in a `select!` / task list. the clone + /// of `self` is deferred until the future is first polled. + /// + /// to disable the HTTP API entirely, simply don't call this method. + pub fn serve(&self, port: u16) -> impl Future> { + let hydrant = self.clone(); + async move { crate::api::serve(hydrant, port).await } + } + + /// returns a future that runs the debug HTTP API server on `127.0.0.1:{port}`. + /// + /// exposes internal inspection endpoints (`/debug/get`, `/debug/iter`, etc.) + /// that are not safe to expose publicly. binds only to loopback. + pub fn serve_debug(&self, port: u16) -> impl Future> { + let state = self.state.clone(); + async move { crate::api::serve_debug(state, port).await } + } +} + +impl axum::extract::FromRef for Arc { + fn from_ref(h: &Hydrant) -> Self { + h.state.clone() + } +} + +// --- event stream --- + +/// a stream of [`Event`]s. returned by [`Hydrant::subscribe`]. +/// +/// implements [`futures::Stream`] and can be used with `StreamExt::next`, +/// `while let Some(evt) = stream.next().await`, `forward`, etc. +/// the stream terminates when the underlying channel closes (i.e. hydrant shuts down). +pub struct EventStream(mpsc::Receiver); + +impl Stream for EventStream { + type Item = Event; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.0.poll_recv(cx) + } +} + +// --- stats --- + +/// 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>, +} + +// --- ingestion handles --- + +/// runtime control over the crawler component. +/// +/// the crawler walks `com.atproto.sync.listRepos` on each configured relay to discover +/// repositories that have never emitted a firehose event. in `filter` mode it also +/// checks each discovered repo against the configured signal collections before +/// enqueuing it for backfill. +/// +/// disabling the crawler does not affect in-progress repo checks. each one completes +/// its current PDS request before pausing. +#[derive(Clone)] +pub struct CrawlerHandle(Arc); + +impl CrawlerHandle { + /// enable the crawler. no-op if already enabled. + pub fn enable(&self) { + self.0.crawler_enabled.send_replace(true); + } + /// disable the crawler. in-progress repo checks finish before the crawler pauses. + pub fn disable(&self) { + self.0.crawler_enabled.send_replace(false); + } + /// returns the current enabled state of the crawler. + pub fn is_enabled(&self) -> bool { + *self.0.crawler_enabled.borrow() + } +} + +/// runtime control over the firehose ingestor component. +/// +/// the firehose connects to each configured relay's `com.atproto.sync.subscribeRepos` +/// websocket and processes commit, identity, account, and sync events in real time. +/// one independent connection is maintained per relay URL. +/// +/// disabling the firehose closes the websocket after the current message is processed. +#[derive(Clone)] +pub struct FirehoseHandle(Arc); + +impl FirehoseHandle { + /// enable the firehose. no-op if already enabled. + pub fn enable(&self) { + self.0.firehose_enabled.send_replace(true); + } + /// disable the firehose. the current message finishes processing before the connection closes. + pub fn disable(&self) { + self.0.firehose_enabled.send_replace(false); + } + /// returns the current enabled state of the firehose. + pub fn is_enabled(&self) -> bool { + *self.0.firehose_enabled.borrow() + } +} + +/// runtime control over the backfill worker component. +/// +/// the backfill worker fetches full repo CAR files from each repo's PDS for any +/// repository in the pending queue, parses the MST, and inserts all matching records +/// into the database. concurrency is bounded by `HYDRANT_BACKFILL_CONCURRENCY_LIMIT`. +/// +/// disabling backfill lets any in-flight repo fetches finish before pausing. +#[derive(Clone)] +pub struct BackfillHandle(Arc); + +impl BackfillHandle { + /// enable the backfill worker. no-op if already enabled. + pub fn enable(&self) { + self.0.backfill_enabled.send_replace(true); + } + /// disable the backfill worker. in-flight repo fetches complete before pausing. + pub fn disable(&self) { + self.0.backfill_enabled.send_replace(false); + } + /// returns the current enabled state of the backfill worker. + pub fn is_enabled(&self) -> bool { + *self.0.backfill_enabled.borrow() + } +} + +// --- filter control --- + +/// a point-in-time snapshot of the filter configuration. returned by all [`FilterControl`] methods. +/// +/// because the filter is stored in the database and loaded on demand, this snapshot +/// may be stale if another caller modifies the filter concurrently. for the authoritative +/// live config use [`FilterControl::get`]. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct FilterSnapshot { + pub mode: FilterMode, + pub signals: Vec, + pub collections: Vec, + pub excludes: Vec, +} + +/// runtime control over the indexing filter. +/// +/// the filter has two orthogonal axes: +/// +/// **mode** controls discovery: +/// - [`FilterMode::Filter`]: only indexes repos whose firehose commits touch a collection +/// matching a configured `signal`. explicit [`ReposControl::track`] always works regardless. +/// - [`FilterMode::Full`]: indexes the entire network. `signals` are ignored for discovery +/// but `collections` and `excludes` still apply. +/// +/// **sets** are each independently configurable: +/// - `signals`: NSID patterns that trigger auto-discovery in `filter` mode (e.g. `app.bsky.feed.post`, `app.bsky.graph.*`) +/// - `collections`: NSID patterns that filter which records are *stored*. empty means store all. +/// - `excludes`: DIDs that are always skipped regardless of mode. +/// +/// NSID patterns support an optional `.*` suffix to match an entire namespace. +/// all mutations are persisted to the database and take effect immediately. +#[derive(Clone)] +pub struct FilterControl(Arc); + +impl FilterControl { + /// return the current filter configuration from the database. + pub async fn get(&self) -> Result { + let filter_ks = self.0.db.filter.clone(); + tokio::task::spawn_blocking(move || { + let hot = db_filter::load(&filter_ks)?; + let excludes = db_filter::read_set(&filter_ks, db_filter::EXCLUDE_PREFIX)?; + Ok(FilterSnapshot { + mode: hot.mode, + signals: hot.signals.iter().map(|s| s.to_string()).collect(), + collections: hot.collections.iter().map(|s| s.to_string()).collect(), + excludes, + }) + }) + .await + .into_diagnostic()? + } + + /// set the indexing mode. see [`FilterControl`] for mode semantics. + pub async fn set_mode(&self, mode: FilterMode) -> Result { + self.patch(Some(mode), None, None, None).await + } + + /// replace the entire signals set. existing signals are removed. + pub async fn set_signals( + &self, + signals: impl IntoIterator>, + ) -> Result { + self.patch( + None, + Some(SetUpdate::Set( + signals.into_iter().map(Into::into).collect(), + )), + None, + None, + ) + .await + } + + /// add multiple signals without disturbing existing ones. + pub async fn append_signals( + &self, + signals: impl IntoIterator>, + ) -> Result { + self.patch( + None, + Some(SetUpdate::Patch( + signals.into_iter().map(|s| (s.into(), true)).collect(), + )), + None, + None, + ) + .await + } + + /// add a single signal. no-op if already present. + pub async fn add_signal(&self, signal: impl Into) -> Result { + self.patch( + None, + Some(SetUpdate::Patch([(signal.into(), true)].into())), + None, + None, + ) + .await + } + + /// remove a single signal. no-op if not present. + pub async fn remove_signal(&self, signal: impl Into) -> Result { + self.patch( + None, + Some(SetUpdate::Patch([(signal.into(), false)].into())), + None, + None, + ) + .await + } + + /// replace the entire collections set. pass an empty iterator to store all collections. + pub async fn set_collections( + &self, + collections: impl IntoIterator>, + ) -> Result { + self.patch( + None, + None, + Some(SetUpdate::Set( + collections.into_iter().map(Into::into).collect(), + )), + None, + ) + .await + } + + /// add multiple collections without disturbing existing ones. + pub async fn append_collections( + &self, + collections: impl IntoIterator>, + ) -> Result { + self.patch( + None, + None, + Some(SetUpdate::Patch( + collections.into_iter().map(|c| (c.into(), true)).collect(), + )), + None, + ) + .await + } + + /// add a single collection filter. no-op if already present. + pub async fn add_collection(&self, collection: impl Into) -> Result { + self.patch( + None, + None, + Some(SetUpdate::Patch([(collection.into(), true)].into())), + None, + ) + .await + } + + /// remove a single collection filter. no-op if not present. + pub async fn remove_collection(&self, collection: impl Into) -> Result { + self.patch( + None, + None, + Some(SetUpdate::Patch([(collection.into(), false)].into())), + None, + ) + .await + } + + /// replace the entire excludes set. + pub async fn set_excludes( + &self, + excludes: impl IntoIterator>, + ) -> Result { + self.patch( + None, + None, + None, + Some(SetUpdate::Set( + excludes.into_iter().map(Into::into).collect(), + )), + ) + .await + } + + /// add multiple DIDs to the excludes set without disturbing existing ones. + pub async fn append_excludes( + &self, + excludes: impl IntoIterator>, + ) -> Result { + self.patch( + None, + None, + None, + Some(SetUpdate::Patch( + excludes.into_iter().map(|d| (d.into(), true)).collect(), + )), + ) + .await + } + + /// add a single DID to the excludes set. no-op if already excluded. + pub async fn add_exclude(&self, did: impl Into) -> Result { + self.patch( + None, + None, + None, + Some(SetUpdate::Patch([(did.into(), true)].into())), + ) + .await + } + + /// remove a single DID from the excludes set. no-op if not present. + pub async fn remove_exclude(&self, did: impl Into) -> Result { + self.patch( + None, + None, + None, + Some(SetUpdate::Patch([(did.into(), false)].into())), + ) + .await + } + + /// apply a batch patch atomically. all provided fields are updated in a single db transaction. + /// returns the updated [`FilterSnapshot`]. this is the primitive all other `FilterControl` methods delegate to. + pub async fn patch( + &self, + mode: Option, + signals: Option, + collections: Option, + excludes: Option, + ) -> Result { + let filter_ks = self.0.db.filter.clone(); + let inner = self.0.db.inner.clone(); + let filter_handle = self.0.filter.clone(); + + let new_filter = tokio::task::spawn_blocking(move || { + let mut batch = inner.batch(); + db_filter::apply_patch(&mut batch, &filter_ks, mode, signals, collections, excludes)?; + batch.commit().into_diagnostic()?; + db_filter::load(&filter_ks) + }) + .await + .into_diagnostic()??; + + let excludes = { + let filter_ks = self.0.db.filter.clone(); + tokio::task::spawn_blocking(move || { + db_filter::read_set(&filter_ks, db_filter::EXCLUDE_PREFIX) + }) + .await + .into_diagnostic()?? + }; + + let snapshot = FilterSnapshot { + mode: new_filter.mode, + signals: new_filter.signals.iter().map(|s| s.to_string()).collect(), + collections: new_filter + .collections + .iter() + .map(|s| s.to_string()) + .collect(), + excludes, + }; + + filter_handle.store(Arc::new(new_filter)); + Ok(snapshot) + } +} + +// --- repos control --- + +/// information about a tracked or known repository. returned by [`ReposControl`] methods. +#[derive(Debug, Clone, serde::Serialize)] +pub struct RepoInfo { + pub did: String, + pub status: String, + pub tracked: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub rev: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub handle: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub pds: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub signing_key: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub last_updated_at: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub last_message_at: Option>, +} + +/// control over which repositories are tracked and access to their state. +/// +/// in `filter` mode, a repo is only indexed if it either matches a signal or is +/// explicitly tracked via [`ReposControl::track`]. in `full` mode all repos are indexed +/// and tracking is implicit. +/// +/// tracking a DID that hydrant has never seen enqueues an immediate backfill. +/// tracking a DID that hydrant already knows about (but has marked untracked) +/// re-enqueues it for backfill. +#[derive(Clone)] +pub struct ReposControl(Arc); + +impl ReposControl { + /// fetch the current state of a single repository. returns `None` if hydrant + /// has never seen this DID. + pub async fn get(&self, did: &Did<'_>) -> Result> { + let did_key = keys::repo_key(did); + let did_str = did.as_str().to_owned(); + let db = self.0.db.clone(); + + tokio::task::spawn_blocking(move || { + let bytes = db.repos.get(&did_key).into_diagnostic()?; + let state = bytes.as_deref().map(db::deser_repo_state).transpose()?; + Ok(state.map(|s| repo_state_to_info(did_str, s))) + }) + .await + .into_diagnostic()? + } + + /// explicitly track one or more repositories, enqueuing them for backfill if needed. + /// + /// - if a DID is new, a fresh [`RepoState`] is created and backfill is queued. + /// - if a DID is already known but untracked, it is marked tracked and re-enqueued. + /// - if a DID is already tracked, this is a no-op. + pub async fn track(&self, dids: impl IntoIterator>) -> Result<()> { + let dids: Vec> = dids.into_iter().map(|d| d.into_static()).collect(); + let state = self.0.clone(); + + let (new_count, transitions) = tokio::task::spawn_blocking(move || { + let db = &state.db; + let mut batch = db.inner.batch(); + let mut added = 0i64; + let mut transitions: Vec<(GaugeState, GaugeState)> = Vec::new(); + let mut rng = rand::rng(); + + for did in &dids { + let did_key = keys::repo_key(did); + let repo_bytes = db.repos.get(&did_key).into_diagnostic()?; + let existing = repo_bytes + .as_deref() + .map(db::deser_repo_state) + .transpose()?; + + if let Some(mut repo_state) = existing { + if !repo_state.tracked { + let resync = db.resync.get(&did_key).into_diagnostic()?; + let old = db::Db::repo_gauge_state(&repo_state, resync.as_deref()); + repo_state.tracked = true; + batch.insert(&db.repos, &did_key, ser_repo_state(&repo_state)?); + batch.insert( + &db.pending, + keys::pending_key(repo_state.index_id), + &did_key, + ); + batch.remove(&db.resync, &did_key); + transitions.push((old, GaugeState::Pending)); + } + } else { + let repo_state = RepoState::backfilling(rng.next_u64()); + batch.insert(&db.repos, &did_key, ser_repo_state(&repo_state)?); + batch.insert( + &db.pending, + keys::pending_key(repo_state.index_id), + &did_key, + ); + added += 1; + transitions.push((GaugeState::Synced, GaugeState::Pending)); + } + } + + batch.commit().into_diagnostic()?; + Ok::<_, miette::Report>((added, transitions)) + }) + .await + .into_diagnostic()??; + + if new_count > 0 { + self.0.db.update_count_async("repos", new_count).await; + } + for (old, new) in transitions { + self.0.db.update_gauge_diff_async(&old, &new).await; + } + self.0.notify_backfill(); + Ok(()) + } + + /// stop tracking one or more repositories. hydrant will stop processing new events + /// for them and remove them from the pending/resync queues, but existing indexed + /// records are **not** deleted. + pub async fn untrack(&self, dids: impl IntoIterator>) -> Result<()> { + let dids: Vec> = dids.into_iter().map(|d| d.into_static()).collect(); + let state = self.0.clone(); + + let gauge_decrements = tokio::task::spawn_blocking(move || { + let db = &state.db; + let mut batch = db.inner.batch(); + let mut gauge_decrements = Vec::new(); + + for did in &dids { + let did_key = keys::repo_key(did); + let repo_bytes = db.repos.get(&did_key).into_diagnostic()?; + let existing = repo_bytes + .as_deref() + .map(db::deser_repo_state) + .transpose()?; + + if let Some(repo_state) = existing { + if repo_state.tracked { + let resync = db.resync.get(&did_key).into_diagnostic()?; + let old = db::Db::repo_gauge_state(&repo_state, resync.as_deref()); + let mut repo_state = repo_state.into_static(); + repo_state.tracked = false; + batch.insert(&db.repos, &did_key, ser_repo_state(&repo_state)?); + batch.remove(&db.pending, keys::pending_key(repo_state.index_id)); + batch.remove(&db.resync, &did_key); + if old != GaugeState::Synced { + gauge_decrements.push(old); + } + } + } + } + + batch.commit().into_diagnostic()?; + Ok::<_, miette::Report>(gauge_decrements) + }) + .await + .into_diagnostic()??; + + for gauge in gauge_decrements { + self.0 + .db + .update_gauge_diff_async(&gauge, &GaugeState::Synced) + .await; + } + Ok(()) + } +} + +pub fn repo_state_to_info(did: String, s: RepoState<'_>) -> RepoInfo { + RepoInfo { + did, + status: s.status.to_string(), + tracked: s.tracked, + rev: s.rev.as_ref().map(|r| r.to_string()), + handle: s.handle.map(|h| h.to_string()), + pds: s.pds.map(|p| p.to_string()), + signing_key: s.signing_key.map(|k| k.encode()), + last_updated_at: DateTime::from_timestamp_secs(s.last_updated_at), + last_message_at: s.last_message_time.and_then(DateTime::from_timestamp_secs), + } +} + +// --- db control --- + +/// control over database maintenance operations. +/// +/// all methods pause the crawler, firehose, and backfill worker for the duration +/// of the operation and restore their prior state on completion, whether or not +/// the operation succeeds. +#[derive(Clone)] +pub struct DbControl(Arc); + +impl DbControl { + /// trigger a major compaction of all keyspaces in parallel. + /// + /// compaction reclaims disk space from deleted/updated keys and improves + /// read performance. can take several minutes on large datasets. + pub async fn compact(&self) -> Result<()> { + let state = self.0.clone(); + state + .with_ingestion_paused(async || state.db.compact().await) + .await + } + + /// train zstd compression dictionaries for the `repos`, `blocks`, and `events` keyspaces. + /// + /// dictionaries are written to `dict_{name}.bin` files next to the database. + /// a restart is required to apply them. training samples data blocks from the + /// existing database, so the database must have a reasonable amount of data first. + pub async fn train_dicts(&self) -> Result<()> { + let state = self.0.clone(); + state + .with_ingestion_paused(async || { + let train = |name: &'static str| { + let db = state.db.clone(); + tokio::task::spawn_blocking(move || db.train_dict(name)) + .map(|res| res.into_diagnostic().flatten()) + }; + tokio::try_join!(train("repos"), train("blocks"), train("events")).map(|_| ()) + }) + .await + } +} + +// --- stream thread --- + +fn event_stream_thread(state: Arc, tx: mpsc::Sender, cursor: Option) { + let db = &state.db; + let mut event_rx = db.event_tx.subscribe(); + let ks = db.events.clone(); + let mut current_id = match cursor { + Some(c) => c.saturating_sub(1), + None => db.next_event_id.load(Ordering::SeqCst).saturating_sub(1), + }; + + loop { + // catch up from db + loop { + let mut found = false; + for item in ks.range(keys::event_key(current_id + 1)..) { + let (k, v) = match item.into_inner() { + Ok(kv) => kv, + Err(e) => { + error!(err = %e, "failed to read event from db"); + break; + } + }; + + let id = match k.as_ref().try_into().map(u64::from_be_bytes) { + Ok(id) => id, + Err(_) => { + error!("failed to parse event id"); + continue; + } + }; + current_id = id; + + let stored: StoredEvent = match rmp_serde::from_slice(&v) { + Ok(e) => e, + Err(e) => { + error!(err = %e, "failed to deserialize stored event"); + continue; + } + }; + + let Some(evt) = stored_to_event(&state, id, stored) else { + continue; + }; + + if tx.blocking_send(evt).is_err() { + return; // receiver dropped + } + found = true; + } + if !found { + break; + } + } + + // wait for live events + match event_rx.blocking_recv() { + Ok(BroadcastEvent::Persisted(_)) => {} // re-run catch-up + Ok(BroadcastEvent::Ephemeral(evt)) => { + if tx.blocking_send(*evt).is_err() { + return; + } + } + Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {} + Err(tokio::sync::broadcast::error::RecvError::Closed) => break, + } + } +} + +fn stored_to_event(state: &AppState, id: u64, stored: StoredEvent<'_>) -> Option { + let StoredEvent { + live, + did, + rev, + collection, + rkey, + action, + data, + } = stored; + + let record = match data { + StoredData::Ptr(cid) => { + let block = state + .db + .blocks + .get(&keys::block_key(collection.as_str(), &cid.to_bytes())); + match block { + Ok(Some(bytes)) => match serde_ipld_dagcbor::from_slice::(&bytes) { + Ok(val) => Some((cid, serde_json::to_value(val).ok()?)), + Err(e) => { + error!(err = %e, "cant parse block"); + return None; + } + }, + Ok(None) => { + error!("block not found, this is a bug"); + return None; + } + Err(e) => { + error!(err = %e, "cant get block"); + db::check_poisoned(&e); + return None; + } + } + } + StoredData::Block(block) => { + let digest = Sha256::digest(&block); + let hash = + cid::multihash::Multihash::wrap(ATP_CID_HASH, &digest).expect("valid sha256 hash"); + let cid = IpldCid::new_v1(DAG_CBOR_CID_CODEC, hash); + match serde_ipld_dagcbor::from_slice::(&block) { + Ok(val) => Some((cid, serde_json::to_value(val).ok()?)), + Err(e) => { + error!(err = %e, "cant parse block"); + return None; + } + } + } + StoredData::Nothing => None, + }; + + let (cid, record) = record + .map(|(c, r)| (Some(c), Some(r))) + .unwrap_or((None, None)); + + Some(MarshallableEvt { + id, + event_type: "record".into(), + record: Some(RecordEvt { + live, + did: did.to_did(), + rev: CowStr::Owned(rev.to_tid().into()), + collection: CowStr::Owned(collection.as_ref().to_string().into()), + rkey: CowStr::Owned(rkey.to_smolstr().into()), + action: CowStr::Borrowed(action.as_str()), + record, + cid: cid.map(|c| jacquard_common::types::cid::Cid::ipld(c).into()), + }), + identity: None, + account: None, + }) +} diff --git a/src/lib.rs b/src/lib.rs --- a/src/lib.rs +++ b/src/lib.rs @@ -1,12 +1,14 @@ -pub mod api; -pub mod backfill; pub mod config; -pub mod crawler; -pub mod db; +pub mod control; pub mod filter; -pub mod ingest; -pub mod ops; -pub mod resolver; -pub mod state; pub mod types; -pub mod util; + +pub(crate) mod api; +pub(crate) mod backfill; +pub(crate) mod crawler; +pub(crate) mod db; +pub(crate) mod ingest; +pub(crate) mod ops; +pub(crate) mod resolver; +pub(crate) mod state; +pub(crate) mod util; diff --git a/src/main.rs b/src/main.rs --- a/src/main.rs +++ b/src/main.rs @@ -1,15 +1,6 @@ -use futures::{FutureExt, future::BoxFuture}; -use hydrant::config::{Config, SignatureVerification}; -use hydrant::db; -use hydrant::ingest::firehose::FirehoseIngestor; -use hydrant::state::AppState; -use hydrant::{api, backfill::BackfillWorker, ingest::worker::FirehoseWorker}; -use miette::IntoDiagnostic; +use hydrant::config::{AppConfig, Config}; +use hydrant::control::Hydrant; use mimalloc::MiMalloc; -use std::sync::Arc; -use std::sync::atomic::Ordering; -use tokio::{sync::mpsc, task::spawn_blocking}; -use tracing::{debug, error, info}; #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; @@ -21,280 +12,25 @@ .ok(); let cfg = Config::from_env()?; + let app = AppConfig::from_env(); let env_filter = tracing_subscriber::EnvFilter::builder() .with_default_directive(tracing::Level::INFO.into()) .from_env_lossy(); tracing_subscriber::fmt().with_env_filter(env_filter).init(); - info!("{cfg}"); + let hydrant = Hydrant::new(cfg).await?; - let state = AppState::new(&cfg)?; - - if cfg.full_network - || cfg.filter_signals.is_some() - || cfg.filter_collections.is_some() - || cfg.filter_excludes.is_some() - { - let filter_ks = state.db.filter.clone(); - let inner = state.db.inner.clone(); - let full_network = cfg.full_network; - let signals = cfg.filter_signals.clone(); - let collections = cfg.filter_collections.clone(); - let excludes = cfg.filter_excludes.clone(); - - tokio::task::spawn_blocking(move || { - use hydrant::filter::{FilterMode, SetUpdate}; - let mut batch = inner.batch(); - - let mode = if full_network { - Some(FilterMode::Full) - } else { - None - }; - - let signals_update = signals.map(SetUpdate::Set); - let collections_update = collections.map(SetUpdate::Set); - let excludes_update = excludes.map(SetUpdate::Set); - - hydrant::db::filter::apply_patch( - &mut batch, - &filter_ks, - mode, - signals_update, - collections_update, - excludes_update, - )?; - - batch.commit().into_diagnostic() - }) - .await - .into_diagnostic()??; - - let new_filter = hydrant::db::filter::load(&state.db.filter)?; - state.filter.store(new_filter.into()); - } - - let (buffer_tx, buffer_rx) = mpsc::unbounded_channel(); - let state = Arc::new(state); - - if cfg.ephemeral { - let state = state.clone(); - std::thread::Builder::new() - .name("ephemeral-gc".into()) - .spawn(move || db::ephemeral::ephemeral_ttl_worker(state)) - .into_diagnostic()?; - } - - tokio::spawn({ - let state = state.clone(); - let timeout = cfg.repo_fetch_timeout; - BackfillWorker::new( - state.clone(), - buffer_tx.clone(), - timeout, - cfg.backfill_concurrency_limit, - matches!( - cfg.verify_signatures, - SignatureVerification::Full | SignatureVerification::BackfillOnly - ), - cfg.ephemeral, - state.backfill_enabled.subscribe(), - ) - .run() - }); - - if let Err(e) = spawn_blocking({ - let state = state.clone(); - move || hydrant::backfill::manager::queue_gone_backfills(&state) - }) - .await - .into_diagnostic()? - { - error!(err = %e, "failed to queue gone backfills"); - db::check_poisoned_report(&e); - } - - std::thread::spawn({ - let state = state.clone(); - move || hydrant::backfill::manager::retry_worker(state) - }); - - tokio::spawn({ - let state = state.clone(); - let mut last_id = state.db.next_event_id.load(Ordering::Relaxed); - let mut last_time = std::time::Instant::now(); - let mut interval = tokio::time::interval(std::time::Duration::from_secs(60)); - async move { - loop { - interval.tick().await; - - let current_id = state.db.next_event_id.load(Ordering::Relaxed); - let current_time = std::time::Instant::now(); - - let delta = current_id.saturating_sub(last_id); - if delta == 0 { - debug!("no new events in 60s"); - continue; - } - - let elapsed = current_time.duration_since(last_time).as_secs_f64(); - let rate = if elapsed > 0.0 { - delta as f64 / elapsed - } else { - 0.0 - }; - - info!("{rate:.2} events/s ({delta} events in {elapsed:.1}s)"); - - last_id = current_id; - last_time = current_time; - } + if app.enable_debug { + tokio::select! { + r = hydrant.run() => r, + r = hydrant.serve(app.api_port) => r, + r = hydrant.serve_debug(app.debug_port) => r, } - }); - - std::thread::spawn({ - let state = state.clone(); - let persist_interval = cfg.cursor_save_interval; - - move || { - loop { - std::thread::sleep(persist_interval); - - // persist firehose cursors - for (relay, cursor) in &state.relay_cursors { - let seq = cursor.load(Ordering::SeqCst); - if seq > 0 { - if let Err(e) = db::set_firehose_cursor(&state.db, relay, seq) { - error!(relay = %relay, err = %e, "failed to save cursor"); - db::check_poisoned_report(&e); - } - } - } - - // persist counts - // TODO: make this more durable - if let Err(e) = db::persist_counts(&state.db) { - error!(err = %e, "failed to persist counts"); - db::check_poisoned_report(&e); - } - - // persist journal - if let Err(e) = state.db.persist() { - error!(err = %e, "db persist failed"); - db::check_poisoned_report(&e); - } - } + } else { + tokio::select! { + r = hydrant.run() => r, + r = hydrant.serve(app.api_port) => r, } - }); - - let post_patch_crawler = match cfg.enable_crawler { - Some(b) => b, - None => state.filter.load().mode == hydrant::filter::FilterMode::Full, - }; - state.crawler_enabled.send_replace(post_patch_crawler); - - info!( - crawler_enabled = *state.crawler_enabled.borrow(), - firehose_enabled = *state.firehose_enabled.borrow(), - filter_mode = ?state.filter.load().mode, - "starting ingestion" - ); - - let relay_hosts = cfg.relays.clone(); - let crawler_max_pending = cfg.crawler_max_pending_repos; - let crawler_resume_pending = cfg.crawler_resume_pending_repos; - - if !relay_hosts.is_empty() { - let state_for_crawler = state.clone(); - let crawler_rx = state.crawler_enabled.subscribe(); - info!( - relay_count = relay_hosts.len(), - hosts = relay_hosts - .iter() - .map(|h| h.as_str()) - .collect::>() - .join(", "), - enabled = *state.crawler_enabled.borrow(), - "starting crawler(s)" - ); - tokio::spawn(async move { - let crawler = hydrant::crawler::Crawler::new( - state_for_crawler, - relay_hosts, - crawler_max_pending, - crawler_resume_pending, - crawler_rx, - ); - if let Err(e) = crawler.run().await { - error!(err = %e, "crawler error"); - db::check_poisoned_report(&e); - } - }); } - - let firehose_worker = std::thread::spawn({ - let state = state.clone(); - let handle = tokio::runtime::Handle::current(); - move || { - FirehoseWorker::new( - state, - buffer_rx, - matches!(cfg.verify_signatures, SignatureVerification::Full), - cfg.ephemeral, - cfg.firehose_workers, - ) - .run(handle) - } - }); - - let mut tasks: Vec>> = vec![Box::pin( - tokio::task::spawn_blocking(move || { - firehose_worker - .join() - .map_err(|e| miette::miette!("buffer processor died: {e:?}")) - }) - .map(|r| r.into_diagnostic().flatten().flatten()), - )]; - - for relay_url in &cfg.relays { - let ingestor = FirehoseIngestor::new( - state.clone(), - buffer_tx.clone(), - relay_url.clone(), - state.filter.clone(), - state.firehose_enabled.subscribe(), - matches!(cfg.verify_signatures, SignatureVerification::Full), - ); - tasks.push(Box::pin(ingestor.run())); - } - - let state_api = state.clone(); - tasks.push(Box::pin(async move { - api::serve(state_api, cfg.api_port) - .await - .map_err(|e| miette::miette!("API server failed: {e}")) - }) as BoxFuture<_>); - - if cfg.enable_debug { - let state_debug = state.clone(); - tasks.push(Box::pin(async move { - api::serve_debug(state_debug, cfg.debug_port) - .await - .map_err(|e| miette::miette!("debug server failed: {e}")) - }) as BoxFuture<_>); - } - - let res = futures::future::select_all(tasks); - if let (Err(e), _, _) = res.await { - error!(err = %e, "critical worker died"); - db::check_poisoned_report(&e); - } - - if let Err(e) = state.db.persist() { - db::check_poisoned_report(&e); - return Err(e); - } - - Ok(()) } diff --git a/src/api/db.rs b/src/api/db.rs --- a/src/api/db.rs +++ b/src/api/db.rs @@ -1,48 +1,30 @@ -use std::sync::Arc; - -use crate::state::AppState; +use crate::control::Hydrant; use axum::{Router, extract::State, http::StatusCode, routing::post}; -use futures::FutureExt; -use miette::IntoDiagnostic; -pub fn router() -> Router> { +pub fn router() -> Router { Router::new() .route("/db/train", post(handle_train_dict)) .route("/db/compact", post(handle_compact)) } pub async fn handle_train_dict( - State(state): State>, -) -> Result { - state - .with_ingestion_paused(async || { - let train = |name: &'static str| { - let db = state.db.clone(); - tokio::task::spawn_blocking(move || db.train_dict(name)) - .map(|res| res.into_diagnostic().flatten()) - }; - let repos = train("repos"); - let blocks = train("blocks"); - let events = train("events"); - - tokio::try_join!(repos, blocks, events) - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - - Ok(StatusCode::OK) - }) + State(hydrant): State, +) -> Result { + hydrant + .db + .train_dicts() .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + Ok(StatusCode::OK) } -pub async fn handle_compact(State(state): State>) -> Result { - state - .with_ingestion_paused(async || { - state - .db - .compact() - .await - .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; - - Ok(StatusCode::OK) - }) +pub async fn handle_compact( + State(hydrant): State, +) -> Result { + hydrant + .db + .compact() .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + Ok(StatusCode::OK) } diff --git a/src/api/filter.rs b/src/api/filter.rs --- a/src/api/filter.rs +++ b/src/api/filter.rs @@ -1,8 +1,4 @@ -use std::sync::Arc; - -use crate::api::AppState; -use crate::db; -use crate::db::filter::EXCLUDE_PREFIX; +use crate::control::Hydrant; use crate::filter::{FilterMode, SetUpdate}; use axum::{ Json, Router, @@ -10,43 +6,25 @@ http::StatusCode, routing::{get, patch}, }; -use miette::IntoDiagnostic; -use serde::{Deserialize, Serialize}; +use serde::Deserialize; -pub fn router() -> Router> { +type FilterSnapshot = crate::control::FilterSnapshot; + +pub fn router() -> Router { Router::new() .route("/filter", get(handle_get_filter)) .route("/filter", patch(handle_patch_filter)) } -#[derive(Serialize)] -pub struct FilterResponse { - pub mode: FilterMode, - pub signals: Vec, - pub collections: Vec, - pub excludes: Vec, -} - pub async fn handle_get_filter( - State(state): State>, -) -> Result, (StatusCode, String)> { - let filter_ks = state.db.filter.clone(); - let resp = tokio::task::spawn_blocking(move || { - let hot = db::filter::load(&filter_ks).map_err(|e| e.to_string())?; - let excludes = - db::filter::read_set(&filter_ks, EXCLUDE_PREFIX).map_err(|e| e.to_string())?; - Ok::<_, String>(FilterResponse { - mode: hot.mode, - signals: hot.signals.iter().map(|s| s.to_string()).collect(), - collections: hot.collections.iter().map(|s| s.to_string()).collect(), - excludes, - }) - }) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))? - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?; - - Ok(Json(resp)) + State(hydrant): State, +) -> Result, (StatusCode, String)> { + hydrant + .filter + .get() + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())) + .map(Json) } #[derive(Deserialize)] @@ -58,45 +36,13 @@ } pub async fn handle_patch_filter( - State(state): State>, + State(hydrant): State, Json(patch): Json, -) -> Result { - let db = &state.db; - - let filter_ks = db.filter.clone(); - let inner = db.inner.clone(); - - let patch_mode = patch.mode; - let patch_signals = patch.signals; - let patch_collections = patch.collections; - let patch_excludes = patch.excludes; - - let new_filter = tokio::task::spawn_blocking(move || { - let mut batch = inner.batch(); - - db::filter::apply_patch( - &mut batch, - &filter_ks, - patch_mode, - patch_signals, - patch_collections, - patch_excludes, - ) - .map_err(|e| e.to_string())?; - - batch - .commit() - .into_diagnostic() - .map_err(|e| e.to_string())?; - - let new_filter = db::filter::load(&filter_ks).map_err(|e| e.to_string())?; - Ok::<_, String>(new_filter) - }) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))? - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?; - - state.filter.store(Arc::new(new_filter)); - - Ok(StatusCode::OK) +) -> Result, (StatusCode, String)> { + hydrant + .filter + .patch(patch.mode, patch.signals, patch.collections, patch.excludes) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string())) + .map(Json) } diff --git a/src/api/ingestion.rs b/src/api/ingestion.rs --- a/src/api/ingestion.rs +++ b/src/api/ingestion.rs @@ -1,6 +1,4 @@ -use std::sync::Arc; - -use crate::state::AppState; +use crate::control::Hydrant; use axum::{ Json, Router, extract::State, @@ -9,7 +7,7 @@ }; use serde::{Deserialize, Serialize}; -pub fn router() -> Router> { +pub fn router() -> Router { Router::new() .route("/ingestion", get(get_ingestion)) .route("/ingestion", patch(patch_ingestion)) @@ -22,11 +20,11 @@ pub backfill: bool, } -pub async fn get_ingestion(State(state): State>) -> Json { +pub async fn get_ingestion(State(hydrant): State) -> Json { Json(IngestionStatus { - crawler: *state.crawler_enabled.borrow(), - firehose: *state.firehose_enabled.borrow(), - backfill: *state.backfill_enabled.borrow(), + crawler: hydrant.crawler.is_enabled(), + firehose: hydrant.firehose.is_enabled(), + backfill: hydrant.backfill.is_enabled(), }) } @@ -41,17 +39,29 @@ } pub async fn patch_ingestion( - State(state): State>, + State(hydrant): State, Json(body): Json, ) -> StatusCode { if let Some(crawler) = body.crawler { - state.crawler_enabled.send_replace(crawler); + if crawler { + hydrant.crawler.enable(); + } else { + hydrant.crawler.disable(); + } } if let Some(firehose) = body.firehose { - state.firehose_enabled.send_replace(firehose); + if firehose { + hydrant.firehose.enable(); + } else { + hydrant.firehose.disable(); + } } if let Some(backfill) = body.backfill { - state.backfill_enabled.send_replace(backfill); + if backfill { + hydrant.backfill.enable(); + } else { + hydrant.backfill.disable(); + } } StatusCode::OK } diff --git a/src/api/mod.rs b/src/api/mod.rs --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -1,3 +1,4 @@ +use crate::control::Hydrant; use crate::state::AppState; use axum::{Router, routing::get}; use std::{net::SocketAddr, sync::Arc}; @@ -13,7 +14,7 @@ mod stream; mod xrpc; -pub async fn serve(state: Arc, port: u16) -> miette::Result<()> { +pub async fn serve(hydrant: Hydrant, port: u16) -> miette::Result<()> { let app = Router::new() .route("/health", get(|| async { "OK" })) .route("/stats", get(stats::get_stats)) @@ -23,7 +24,7 @@ .merge(repos::router()) .merge(ingestion::router()) .merge(db::router()) - .with_state(state) + .with_state(hydrant) .layer(TraceLayer::new_for_http()) .layer(CorsLayer::permissive()); diff --git a/src/api/repos.rs b/src/api/repos.rs --- a/src/api/repos.rs +++ b/src/api/repos.rs @@ -1,5 +1,5 @@ -use std::sync::Arc; - +use crate::control::{Hydrant, RepoInfo, repo_state_to_info}; +use crate::db::keys; use axum::{ Json, Router, body::Body, @@ -8,17 +8,10 @@ response::{IntoResponse, Response}, routing::{delete, get, put}, }; -use chrono::{DateTime, Utc}; -use jacquard_common::{IntoStatic, types::did::Did}; -use miette::IntoDiagnostic; -use rand::Rng; -use serde::{Deserialize, Serialize}; +use jacquard_common::types::did::Did; +use serde::Deserialize; -use crate::api::AppState; -use crate::db::{keys, ser_repo_state}; -use crate::types::{GaugeState, RepoState}; - -pub fn router() -> Router> { +pub fn router() -> Router { Router::new() .route("/repos", get(handle_get_repos)) .route("/repos/{did}", get(handle_get_repo)) @@ -31,26 +24,6 @@ pub did: String, } -#[derive(Serialize, Debug)] -pub struct RepoResponse { - pub did: String, - pub status: String, - pub tracked: bool, - #[serde(skip_serializing_if = "Option::is_none")] - pub rev: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub handle: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub pds: Option, - // this does not have the did:key: prefix - #[serde(skip_serializing_if = "Option::is_none")] - pub signing_key: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub last_updated_at: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub last_message_at: Option>, -} - #[derive(Deserialize)] pub struct GetReposParams { pub limit: Option, @@ -59,22 +32,22 @@ } pub async fn handle_get_repos( - State(state): State>, + State(hydrant): State, Query(params): Query, ) -> Result { let limit = params.limit.unwrap_or(100).min(1000); let partition = params.partition.unwrap_or_else(|| "all".to_string()); let items = tokio::task::spawn_blocking(move || { - let db = &state.db; + let db = &hydrant.state.db; - let to_response = |k: &[u8], v: &[u8]| -> Result { + let to_info = |k: &[u8], v: &[u8]| -> Result { let repo_state = crate::db::deser_repo_state(v).map_err(internal)?; let did = crate::db::types::TrimmedDid::try_from(k) .map_err(internal)? .to_did(); - Ok(repo_state_to_response(did.to_string(), repo_state)) + Ok(repo_state_to_info(did.to_string(), repo_state)) }; let results = match partition.as_str() { @@ -105,7 +78,7 @@ })? }; - items.push(to_response(&k, &repo_state_bytes)?); + items.push(to_info(&k, &repo_state_bytes)?); } Ok::<_, (StatusCode, String)>(items) } @@ -126,7 +99,7 @@ let (_, did_key) = item.into_inner().map_err(internal)?; if let Ok(Some(v)) = db.repos.get(&did_key) { - items.push(to_response(&did_key, &v)?); + items.push(to_info(&did_key, &v)?); } } Ok(items) @@ -153,194 +126,59 @@ } pub async fn handle_get_repo( - State(state): State>, + State(hydrant): State, Path(did_str): Path, -) -> Result, (StatusCode, String)> { +) -> Result, (StatusCode, String)> { let did = Did::new(&did_str).map_err(bad_request)?; - let did_key = keys::repo_key(&did); - let item = tokio::task::spawn_blocking(move || { - let db = &state.db; - - let repo_bytes = db.repos.get(&did_key).map_err(internal)?; - let repo_state = repo_bytes - .as_deref() - .map(crate::db::deser_repo_state) - .transpose() - .map_err(internal)?; - - Ok(repo_state.map(|s| repo_state_to_response(did_str, s))) - }) - .await - .map_err(internal)??; + let item = hydrant + .repos + .get(&did) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; item.map(Json) .ok_or_else(|| (StatusCode::NOT_FOUND, "repository not found".to_string())) } pub async fn handle_put_repos( - State(state): State>, + State(hydrant): State, req: axum::extract::Request, ) -> Result { let items = parse_body(req).await?; - let state_task = state.clone(); - let (new_repo_count, gauge_transitions) = tokio::task::spawn_blocking(move || { - let db = &state_task.db; - let mut batch = db.inner.batch(); - let mut added = 0i64; - let mut gauge_transitions: Vec<(GaugeState, GaugeState)> = Vec::new(); + let dids: Vec> = items + .into_iter() + .filter_map(|item| Did::new_owned(&item.did).ok()) + .collect(); - let mut rng = rand::rng(); - - for item in items { - let did = Did::new(&item.did).map_err(bad_request)?; - let did_key = keys::repo_key(&did); - - let repo_bytes = db.repos.get(&did_key).map_err(internal)?; - let existing_state = repo_bytes - .as_deref() - .map(crate::db::deser_repo_state) - .transpose() - .map_err(internal)?; - - if let Some(mut repo_state) = existing_state { - if !repo_state.tracked { - let resync_bytes = db.resync.get(&did_key).map_err(internal)?; - let old_gauge = - crate::db::Db::repo_gauge_state(&repo_state, resync_bytes.as_deref()); - - repo_state.tracked = true; - // re-enqueue into pending - batch.insert( - &db.repos, - &did_key, - ser_repo_state(&repo_state).map_err(internal)?, - ); - batch.insert( - &db.pending, - keys::pending_key(repo_state.index_id), - &did_key, - ); - batch.remove(&db.resync, &did_key); - gauge_transitions.push((old_gauge, GaugeState::Pending)); - } - } else { - let repo_state = RepoState::backfilling(rng.next_u64()); - batch.insert( - &db.repos, - &did_key, - ser_repo_state(&repo_state).map_err(internal)?, - ); - batch.insert( - &db.pending, - keys::pending_key(repo_state.index_id), - &did_key, - ); - added += 1; - gauge_transitions.push((GaugeState::Synced, GaugeState::Pending)); // pseudo-transition to just inc pending - } - } - - batch.commit().into_diagnostic().map_err(internal)?; - - Ok::<_, (StatusCode, String)>((added, gauge_transitions)) - }) - .await - .map_err(internal)??; - - if new_repo_count > 0 { - state.db.update_count_async("repos", new_repo_count).await; - } - for (old, new) in gauge_transitions { - state.db.update_gauge_diff_async(&old, &new).await; - } - - // Always notify backfill if anything was added to pending! - state.notify_backfill(); + hydrant + .repos + .track(dids) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; Ok(StatusCode::OK) } pub async fn handle_delete_repos( - State(state): State>, + State(hydrant): State, req: axum::extract::Request, ) -> Result { let items = parse_body(req).await?; - let state_task = state.clone(); - let (deleted_count, gauge_decrements) = tokio::task::spawn_blocking(move || { - let db = &state_task.db; - let mut batch = db.inner.batch(); - // keeping this for later, unused for now - let deleted_count = 0i64; - let mut gauge_decrements = Vec::new(); + let dids: Vec> = items + .into_iter() + .filter_map(|item| Did::new_owned(&item.did).ok()) + .collect(); - for item in items { - let did = Did::new(&item.did).map_err(bad_request)?; - let did_key = keys::repo_key(&did); - - let repo_bytes = db.repos.get(&did_key).map_err(internal)?; - let existing_state = repo_bytes - .as_deref() - .map(crate::db::deser_repo_state) - .transpose() - .map_err(internal)?; - - if let Some(repo_state) = existing_state { - let resync_bytes = db.resync.get(&did_key).map_err(internal)?; - let old_gauge = - crate::db::Db::repo_gauge_state(&repo_state, resync_bytes.as_deref()); - - if repo_state.tracked { - let mut repo_state = repo_state.into_static(); - repo_state.tracked = false; - batch.insert( - &db.repos, - &did_key, - ser_repo_state(&repo_state).map_err(internal)?, - ); - batch.remove(&db.pending, keys::pending_key(repo_state.index_id)); - batch.remove(&db.resync, &did_key); - if old_gauge != GaugeState::Synced { - gauge_decrements.push(old_gauge); - } - } - } - } - - batch.commit().into_diagnostic().map_err(internal)?; - - Ok::<_, (StatusCode, String)>((deleted_count, gauge_decrements)) - }) - .await - .map_err(internal)??; - - if deleted_count > 0 { - state.db.update_count_async("repos", -deleted_count).await; - } - for gauge in gauge_decrements { - state - .db - .update_gauge_diff_async(&gauge, &GaugeState::Synced) - .await; - } + hydrant + .repos + .untrack(dids) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; Ok(StatusCode::OK) -} - -fn repo_state_to_response(did: String, s: RepoState<'_>) -> RepoResponse { - RepoResponse { - did, - status: s.status.to_string(), - tracked: s.tracked, - rev: s.rev.as_ref().map(|r| r.to_string()), - handle: s.handle.map(|h| h.to_string()), - pds: s.pds.map(|p| p.to_string()), - signing_key: s.signing_key.map(|k| k.encode()), - last_updated_at: DateTime::from_timestamp_secs(s.last_updated_at), - last_message_at: s.last_message_time.and_then(DateTime::from_timestamp_secs), - } } async fn parse_body(req: axum::extract::Request) -> Result, (StatusCode, String)> { diff --git a/src/api/stats.rs b/src/api/stats.rs --- a/src/api/stats.rs +++ b/src/api/stats.rs @@ -1,58 +1,9 @@ -use crate::api::AppState; -use axum::{Json, extract::State, response::Result}; -use serde::Serialize; -use std::{collections::BTreeMap, sync::Arc}; +use crate::control::Hydrant; +use axum::{Json, extract::State, http::StatusCode, response::IntoResponse, response::Response}; -#[derive(Serialize)] -pub struct StatsResponse { - pub counts: BTreeMap<&'static str, u64>, - pub size: BTreeMap<&'static str, u64>, -} - -pub async fn get_stats(State(state): State>) -> Result> { - let db = state.db.clone(); - - let mut counts: BTreeMap<&'static str, u64> = futures::future::join_all( - [ - "repos", - "pending", - "resync", - "records", - "blocks", - "error_ratelimited", - "error_transport", - "error_generic", - ] - .into_iter() - .map(|name| { - let db = db.clone(); - async move { (name, db.get_count(name).await) } - }), - ) - .await - .into_iter() - .collect(); - // this should be accurate since we dont remove events - // todo: ...unless in ephemeral mode - counts.insert("events", db.events.approximate_len() as u64); - - let size = tokio::task::spawn_blocking(move || { - let mut size = BTreeMap::new(); - size.insert("repos", db.repos.disk_space()); - size.insert("records", db.records.disk_space()); - size.insert("blocks", db.blocks.disk_space()); - size.insert("cursors", db.cursors.disk_space()); - size.insert("pending", db.pending.disk_space()); - size.insert("resync", db.resync.disk_space()); - size.insert("resync_buffer", db.resync_buffer.disk_space()); - size.insert("events", db.events.disk_space()); - size.insert("counts", db.counts.disk_space()); - size.insert("filter", db.filter.disk_space()); - size.insert("crawler", db.crawler.disk_space()); - size - }) - .await - .map_err(|_| axum::http::StatusCode::INTERNAL_SERVER_ERROR)?; - - Ok(Json(StatsResponse { counts, size })) +pub async fn get_stats(State(hydrant): State) -> Response { + match hydrant.stats().await { + Ok(stats) => Json(stats).into_response(), + Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response(), + } } diff --git a/src/api/stream.rs b/src/api/stream.rs --- a/src/api/stream.rs +++ b/src/api/stream.rs @@ -1,6 +1,4 @@ -use crate::api::AppState; -use crate::db::keys; -use crate::types::{BroadcastEvent, MarshallableEvt, RecordEvt, StoredData, StoredEvent}; +use crate::control::Hydrant; use axum::Router; use axum::routing::get; use axum::{ @@ -10,18 +8,11 @@ }, response::IntoResponse, }; -use cid::multihash::Multihash; -use jacquard_common::types::cid::{ATP_CID_HASH, IpldCid}; -use jacquard_common::{CowStr, RawData}; -use jacquard_repo::DAG_CBOR_CID_CODEC; -use miette::{Context, IntoDiagnostic}; +use futures::StreamExt; use serde::Deserialize; -use sha2::{Digest, Sha256}; -use std::sync::Arc; -use tokio::sync::{broadcast, mpsc, oneshot}; -use tracing::{error, info_span}; +use tracing::error; -pub fn router() -> Router> { +pub fn router() -> Router { Router::new().route("/", get(handle_stream)) } @@ -31,236 +22,25 @@ } pub async fn handle_stream( - State(state): State>, + State(hydrant): State, Query(query): Query, ws: WebSocketUpgrade, ) -> impl IntoResponse { - ws.on_upgrade(move |socket| handle_socket(socket, state, query)) + ws.on_upgrade(move |socket| handle_socket(socket, hydrant, query)) } -async fn handle_socket(mut socket: WebSocket, state: Arc, query: StreamQuery) { - let (tx, mut rx) = mpsc::channel(500); - let (cancel_tx, cancel_rx) = oneshot::channel::<()>(); +async fn handle_socket(mut socket: WebSocket, hydrant: Hydrant, query: StreamQuery) { + let mut stream = hydrant.subscribe(query.cursor); - let runtime = tokio::runtime::Handle::current(); - let id = std::time::SystemTime::UNIX_EPOCH - .elapsed() - .unwrap() - .as_secs(); - - let thread = std::thread::Builder::new() - .name(format!("stream-handler-{id}")) - .spawn(move || { - let _runtime_guard = runtime.enter(); - stream(state, cancel_rx, tx, query, id); - }) - .expect("failed to spawn stream handler thread"); - - while let Some(msg) = rx.recv().await { - if let Err(e) = socket.send(msg).await { - error!(err = %e, "failed to send ws message"); - break; - } - } - - let _ = cancel_tx.send(()); - if let Err(e) = thread.join() { - error!(err = ?e, "stream handler thread panicked"); - } -} - -fn stream( - state: Arc, - mut cancel: oneshot::Receiver<()>, - tx: mpsc::Sender, - query: StreamQuery, - id: u64, -) { - let db = &state.db; - let mut event_rx = db.event_tx.subscribe(); - let ks = db.events.clone(); - let mut current_id = match query.cursor { - Some(cursor) => cursor.saturating_sub(1), - None => { - let max_id = db.next_event_id.load(std::sync::atomic::Ordering::SeqCst); - max_id.saturating_sub(1) - } - }; - let runtime = tokio::runtime::Handle::current(); - - let span = info_span!("stream", id); - let _entered_span = span.enter(); - - loop { - // 1. catch up from DB - loop { - let mut found = false; - for item in ks.range(keys::event_key(current_id + 1)..) { - let (k, v) = match item.into_inner() { - Ok((k, v)) => (k, v), - Err(e) => { - error!(err = %e, "failed to read event from db"); - break; - } - }; - let id = match k - .as_ref() - .try_into() - .into_diagnostic() - .wrap_err("expected event id to be 8 bytes") - .map(u64::from_be_bytes) - { - Ok(id) => id, - Err(e) => { - error!(err = %e, "failed to parse event id"); - continue; - } - }; - current_id = id; - - let StoredEvent { - live, - did, - rev, - collection, - rkey, - action, - data, - } = match rmp_serde::from_slice(&v) { - Ok(e) => e, - Err(e) => { - error!(err = %e, "failed to deserialize stored event"); - continue; - } - }; - - let _entered = info_span!("record", data = ?data).entered(); - - let record = match data { - StoredData::Ptr(cid) => { - let block = db - .blocks - .get(&keys::block_key(collection.as_str(), &cid.to_bytes())); - match block { - Ok(Some(bytes)) => { - match serde_ipld_dagcbor::from_slice::(&bytes) { - Ok(val) => Some(( - cid, - serde_json::to_value(val) - .expect("that cbor raw data is valid json"), - )), - Err(e) => { - error!(err = %e, "cant parse block, must be corrupted?"); - return; - } - } - } - Ok(None) => { - error!("block not found? this is a bug!!"); - continue; - } - Err(e) => { - error!(err = %e, "can't get block"); - crate::db::check_poisoned(&e); - return; - } - } - } - StoredData::Block(block) => { - let digest = Sha256::digest(&block); - let hash = - Multihash::wrap(ATP_CID_HASH, &digest).expect("that its valid sha256"); - let cid = IpldCid::new_v1(DAG_CBOR_CID_CODEC, hash); - match serde_ipld_dagcbor::from_slice::(&block) { - Ok(val) => Some(( - cid, - serde_json::to_value(val) - .expect("that cbor raw data is valid json"), - )), - Err(e) => { - error!(err = %e, "cant parse block, must be corrupted?"); - return; - } - } - } - StoredData::Nothing => None, - }; - - let (cid, record) = record - .map(|(c, r)| (Some(c), Some(r))) - .unwrap_or((None, None)); - let marshallable = MarshallableEvt { - id, - event_type: "record".into(), - record: Some(RecordEvt { - live, - did: did.to_did(), - rev: CowStr::Owned(rev.to_tid().into()), - collection, - rkey: CowStr::Owned(rkey.to_smolstr().into()), - action: CowStr::Borrowed(action.as_str()), - record, - cid: cid.map(|c| jacquard_common::types::cid::Cid::ipld(c).into()), - }), - identity: None, - account: None, - }; - - let json_str = match serde_json::to_string(&marshallable) { - Ok(s) => s, - Err(e) => { - error!(err = %e, "failed to serialize ws event"); - continue; - } - }; - - if let Err(e) = tx.blocking_send(Message::Text(json_str.into())) { - error!(err = %e, "failed to send ws message"); - return; - } - - found = true; - } - if !found { - break; - } - } - - // 2. wait for live events - let next_event = runtime.block_on(async { - tokio::select! { - res = event_rx.recv() => Some(res), - _ = &mut cancel => None, - } - }); - - let Some(next_event) = next_event else { - break; - }; - - match next_event { - Ok(BroadcastEvent::Persisted(_)) => { - // just wake up and run catch-up loop again - } - Ok(BroadcastEvent::Ephemeral(evt)) => { - // send ephemeral event directly - let json_str = match serde_json::to_string(&evt) { - Ok(s) => s, - Err(e) => { - error!(err = %e, "failed to serialize ws event"); - continue; - } - }; - if let Err(e) = tx.blocking_send(Message::Text(json_str.into())) { - error!(err = %e, "failed to send ws message"); - return; + while let Some(evt) = stream.next().await { + match serde_json::to_string(&evt) { + Ok(json) => { + if socket.send(Message::Text(json.into())).await.is_err() { + break; } } - Err(broadcast::error::RecvError::Lagged(_)) => { - // continue to catch up - } - Err(broadcast::error::RecvError::Closed) => { - break; + Err(e) => { + error!(err = %e, "failed to serialize event"); } } } diff --git a/src/api/xrpc.rs b/src/api/xrpc.rs --- a/src/api/xrpc.rs +++ b/src/api/xrpc.rs @@ -1,6 +1,7 @@ -use crate::api::AppState; +use crate::control::Hydrant; use crate::db::types::DbRkey; use crate::db::{self, Db, keys}; +use crate::state::AppState; use axum::extract::FromRequest; use axum::response::IntoResponse; use axum::{Json, Router, extract::State, http::StatusCode}; @@ -27,7 +28,7 @@ use std::{fmt::Display, sync::Arc}; use tokio::task::spawn_blocking; -pub fn router() -> Router> { +pub fn router() -> Router { Router::new() .route( GetRecordRequest::PATH, diff --git a/src/db/keys.rs b/src/db/keys.rs --- a/src/db/keys.rs +++ b/src/db/keys.rs @@ -177,11 +177,3 @@ key.extend_from_slice(cid); key } - -// prefix format: {collection}| -pub fn block_prefix_collection(collection: &str) -> Vec { - let mut prefix = Vec::with_capacity(collection.len() + 1); - prefix.extend_from_slice(collection.as_bytes()); - prefix.push(SEP); - prefix -}