use futures::FutureExt;
use miette::{IntoDiagnostic, Result};
use std::future::Future;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};
use tokio::sync::watch;
use tracing::{debug, error, info, warn};
use url::Url;
use crate::config::{Config, SignatureVerification};
use crate::db::{self, filter as db_filter};
use crate::filter::FilterMode;
use crate::state::AppState;
use super::firehose::{FirehoseHandle, FirehoseShared};
use super::{DbControl, FilterControl, ReposControl, pds::PdsControl};
#[cfg(feature = "indexer")]
use super::{BackfillHandle, crawler};
#[cfg(feature = "indexer")]
use crate::backfill::BackfillWorker;
#[cfg(feature = "indexer")]
use crate::db::load_persisted_crawler_sources;
#[cfg(feature = "indexer")]
use crate::ingest::indexer::FirehoseWorker;
#[cfg(feature = "backlinks")]
use crate::backlinks::BacklinksControl;
use super::seed;
pub(crate) mod run;
/// the file fjall locks while a database is open, `fjall::file::LOCK_FILE` isn't public
const FJALL_LOCK_FILE: &str = "lock";
/// [`Hydrant::run`] fails with this once fjall refuses writes because a disk write
/// failed. every later write fails the same way, so the database has to be closed
/// and opened again, and looked at if it keeps happening.
#[derive(Debug, miette::Diagnostic, thiserror::Error)]
#[error("the database is poisoned: a disk write failed, so fjall refuses all writes")]
pub struct DatabasePoisoned;
/// 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.
#[derive(Clone)]
pub struct Hydrant {
#[cfg(feature = "indexer")]
pub crawler: crawler::CrawlerHandle,
pub firehose: FirehoseHandle,
#[cfg(feature = "indexer")]
pub backfill: BackfillHandle,
pub filter: FilterControl,
pub pds: PdsControl,
pub repos: ReposControl,
pub db: DbControl,
#[cfg(feature = "backlinks")]
pub backlinks: crate::backlinks::BacklinksControl,
pub(crate) state: Arc<AppState>,
pub(crate) config: Arc<Config>,
pub(crate) started: Arc<AtomicBool>,
_priv: (),
}
impl Hydrant {
/// Returns a reference to the internal DID resolver.
pub fn resolver(&self) -> &crate::resolver::Resolver {
&self.state.resolver
}
#[cfg(any(feature = "indexer_stream", feature = "relay", feature = "jetstream"))]
pub(crate) fn stream_send_timeout(&self) -> std::time::Duration {
self.config.stream_send_timeout
}
/// 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<Self> {
Self::new_inner(config, None).await
}
#[cfg(all(test, feature = "indexer_stream"))]
pub(crate) async fn new_with_resolver(
config: Config,
resolver: crate::resolver::Resolver,
) -> Result<Self> {
Self::new_inner(config, Some(resolver)).await
}
async fn new_inner(
config: Config,
resolver: Option<crate::resolver::Resolver>,
) -> Result<Self> {
info!("{config}");
crate::net::allow_private_networks(config.allow_private_networks);
#[cfg(feature = "relay")]
if config.only_index_links {
miette::bail!("HYDRANT_ONLY_INDEX_LINKS is not supported in relay mode");
}
// 1. open database and construct AppState
let mut state = AppState::new(&config)?;
if let Some(resolver) = resolver {
state.resolver = resolver;
}
// 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 mode = config.full_network.then_some(FilterMode::Full);
let signals = config
.filter_signals
.clone()
.map(crate::patch::SetUpdate::Set);
let collections = config
.filter_collections
.clone()
.map(crate::patch::SetUpdate::Set);
let excludes = config
.filter_excludes
.clone()
.map(crate::patch::SetUpdate::Set);
state
.db
.run(move |db| {
let mut batch = db.inner.batch();
db_filter::apply_patch(
&mut batch,
&db.filter,
mode,
signals,
collections,
excludes,
)?;
batch.commit().into_diagnostic()
})
.await?;
// 3. reload the live filter into the hot-path arc-swap
let new_filter = state.db.run(move |db| db_filter::load(&db.filter)).await?;
state.filter.store(Arc::new(new_filter));
}
#[cfg(feature = "indexer")]
{
// 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
|| !config.crawler_sources.is_empty()
}
};
state.crawler_enabled.send_replace(post_patch_crawler);
}
let state = Arc::new(state);
Ok(Self {
firehose: FirehoseHandle::new(state.clone()),
filter: FilterControl(state.clone()),
pds: PdsControl(state.clone()),
repos: ReposControl(state.clone()),
db: DbControl(state.clone()),
#[cfg(feature = "indexer")]
crawler: crawler::CrawlerHandle {
state: state.clone(),
shared: Arc::new(std::sync::OnceLock::new()),
tasks: Arc::new(scc::HashMap::new()),
persisted: Arc::new(scc::HashSet::new()),
},
#[cfg(feature = "indexer")]
backfill: BackfillHandle::new(state.clone()),
#[cfg(feature = "backlinks")]
backlinks: crate::backlinks::BacklinksControl(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> {
Self::new(Config::from_env()?).await
}
/// returns a future that runs the HTTP management API server on the given bind addresses.
///
/// 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, binds: super::ApiBinds) -> impl Future<Output = Result<()>> {
let hydrant = self.clone();
async move { crate::api::serve(hydrant, binds).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.).
/// binds only to loopback.
pub fn serve_debug(&self, port: u16) -> impl Future<Output = Result<()>> {
let state = self.state.clone();
async move { crate::api::serve_debug(state, port).await }
}
/// call after the api stops accepting and before [`shutdown`](Self::shutdown).
pub async fn drain_streams(&self, timeout: Duration) {
self.state.drain.drain(timeout).await;
}
/// stops hydrant's own threads and waits until its database is closed, which
/// gives its memory back and lets the same path be opened again.
///
/// the tasks [`run`](Self::run) and [`serve`](Self::serve) spawned live on your
/// tokio runtime and can't be stopped from here, so shut that runtime down first
/// (`shutdown_background`, since dropping it would wait on threads this stops),
/// and drop every other clone of this handle and any open stream. whatever they
/// were doing is cut off like a crash, which hydrant recovers from on the next
/// start.
///
/// fails if the database is still open after `timeout`, which means something
/// still holds on to it.
pub fn shutdown(self, timeout: Duration) -> Result<()> {
self.state.stop.trigger();
// the shards stop once nothing can send them anything
#[cfg(feature = "indexer")]
self.state.indexer.lock().take();
let db = Arc::downgrade(&self.state.db.inner);
let lock_path = self.state.db.path.join(FJALL_LOCK_FILE);
drop(self);
let deadline = Instant::now() + timeout;
while db.strong_count() > 0 {
if Instant::now() >= deadline {
miette::bail!(
"the database is still open after {timeout:?}, something still holds a hydrant handle"
);
}
std::thread::sleep(Duration::from_millis(10));
}
// the last handle going away isn't the database closing yet: fjall's drop still
// waits for its own threads, and keyspace handles keep its lock file locked until
// they're dropped too. a restart needs that lock, so wait for it
let lock = std::fs::File::open(&lock_path).into_diagnostic()?;
loop {
match lock.try_lock() {
Ok(()) => return Ok(()),
Err(std::fs::TryLockError::WouldBlock) => {}
Err(std::fs::TryLockError::Error(e)) => return Err(e).into_diagnostic(),
}
if Instant::now() >= deadline {
miette::bail!(
"the database is still locked after {timeout:?}, something still holds one of its keyspaces"
);
}
std::thread::sleep(Duration::from_millis(10));
}
}
}
impl axum::extract::FromRef<Hydrant> for Arc<AppState> {
fn from_ref(h: &Hydrant) -> Self {
h.state.clone()
}
}
#[cfg(all(test, feature = "indexer"))]
mod tests {
use super::*;
use crate::config::Config;
use tempfile::tempdir;
fn test_config(path: &std::path::Path) -> Config {
Config {
database_path: path.to_path_buf(),
..Default::default()
}
}
/// starts a hydrant that never touches the network on its own runtime.
fn start_offline(
path: &std::path::Path,
) -> (
tokio::runtime::Runtime,
Hydrant,
tokio::task::JoinHandle<Result<()>>,
) {
let runtime = tokio::runtime::Runtime::new().unwrap();
let config = Config {
relays: Vec::new(),
seed_hosts: Vec::new(),
crawler_sources: Vec::new(),
enable_crawler: Some(false),
..test_config(path)
};
let hydrant = runtime.block_on(Hydrant::new(config)).unwrap();
let run = runtime.spawn({
let hydrant = hydrant.clone();
async move { hydrant.run()?.await }
});
// long enough for run to have spawned its threads
std::thread::sleep(Duration::from_millis(300));
(runtime, hydrant, run)
}
#[test]
fn shutdown_closes_the_database_for_the_next_start() {
let tmp = tempdir().unwrap();
let (runtime, hydrant, _) = start_offline(tmp.path());
runtime.shutdown_background();
hydrant.shutdown(Duration::from_secs(10)).unwrap();
// fjall locks the database, so this only opens if the first one let go
let (runtime, hydrant, _) = start_offline(tmp.path());
runtime.shutdown_background();
hydrant.shutdown(Duration::from_secs(10)).unwrap();
}
#[test]
fn shutdown_waits_until_the_database_is_unlocked() {
let tmp = tempdir().unwrap();
let (runtime, hydrant, _) = start_offline(tmp.path());
runtime.shutdown_background();
// a keyspace handle outliving every hydrant handle keeps fjall's lock, like fjall's
// own drop does while it waits for its threads, only longer
let keyspace = hydrant.state.db.repos.keyspace();
let release = std::thread::spawn(move || {
std::thread::sleep(Duration::from_millis(300));
drop(keyspace);
});
hydrant.shutdown(Duration::from_secs(10)).unwrap();
// opening again would hide the race, since fjall retries its lock for 200ms
let lock = std::fs::File::open(tmp.path().join(FJALL_LOCK_FILE)).unwrap();
lock.try_lock()
.expect("shutdown returned while the database was still locked");
release.join().unwrap();
}
#[test]
fn shutdown_fails_while_a_clone_holds_the_database() {
let tmp = tempdir().unwrap();
let (runtime, hydrant, _) = start_offline(tmp.path());
runtime.shutdown_background();
let forgotten = hydrant.clone();
let err = hydrant.shutdown(Duration::from_millis(200)).unwrap_err();
assert!(err.to_string().contains("still open"), "{err}");
// the clone still works and is the last one, so it can finish the job
forgotten.shutdown(Duration::from_secs(10)).unwrap();
}
#[test]
fn run_fails_once_the_database_is_poisoned() {
let tmp = tempdir().unwrap();
let (runtime, hydrant, run) = start_offline(tmp.path());
hydrant.state.db.poison.check(&fjall::Error::Poisoned);
let err = runtime
.block_on(async { tokio::time::timeout(Duration::from_secs(10), run).await })
.expect("run kept going on a poisoned database")
.unwrap()
.unwrap_err();
assert!(err.downcast_ref::<DatabasePoisoned>().is_some(), "{err}");
// and it can still be shut down, so the database can be reopened
runtime.shutdown_background();
hydrant.shutdown(Duration::from_secs(10)).unwrap();
}
}