use crate::filter::{FilterHandle, FilterMode};
use crate::ingest::cursor::{self, SeqTracker};
use crate::ingest::stream::{
FirehoseError, FirehoseStream, InfoName, SubscribeReposMessage, decode_frame,
};
use crate::ingest::{BufferTx, IngestMessage};
use crate::pds_meta::HostStatus;
use crate::state::AppState;
use crate::util::WatchEnabledExt;
use crate::util::throttle::ThrottleHandle;
use jacquard_common::IntoStatic;
use jacquard_common::types::did::Did;
use miette::{IntoDiagnostic, Result};
use rand::RngExt;
use rand::rngs::SmallRng;
use std::borrow::Cow;
use std::io::ErrorKind;
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::{Duration, Instant};
use tokio::sync::watch;
use tracing::{Span, debug, error, info, trace, warn};
use url::Url;
// these match ref relay
const MAX_BACKOFF: Duration = Duration::from_secs(60);
const MAX_OFFLINE_BACKOFF: Duration = Duration::from_secs(3600);
#[derive(Debug, Clone)]
struct FirehoseFailure {
kind: &'static str,
detail: String,
}
impl FirehoseFailure {
fn new(kind: &'static str, detail: impl Into<String>) -> Self {
Self {
kind,
detail: detail.into(),
}
}
}
fn classify_firehose_error(err: &FirehoseError) -> FirehoseFailure {
match err {
FirehoseError::WebSocket(err) => classify_websocket_error(err),
FirehoseError::UnknownScheme(scheme) => {
FirehoseFailure::new("config", format!("unsupported URL scheme `{scheme}`"))
}
FirehoseError::InvalidUri(err) => {
FirehoseFailure::new("config", format!("invalid websocket URI: {err}"))
}
FirehoseError::Decode(err) => {
FirehoseFailure::new("decode", format!("failed to decode firehose frame: {err}"))
}
FirehoseError::EmptyFrame => FirehoseFailure::new("protocol", "received empty frame"),
FirehoseError::RelayError { error, message } => FirehoseFailure::new(
"relay_error",
message
.as_deref()
.map(|message| format!("{error}: {message}"))
.unwrap_or_else(|| error.clone()),
),
FirehoseError::UnknownOp(op) => {
FirehoseFailure::new("protocol", format!("unknown frame op {op}"))
}
FirehoseError::MissingType => FirehoseFailure::new("protocol", "missing frame type header"),
FirehoseError::UnknownType(ty) => {
FirehoseFailure::new("protocol", format!("unknown frame type `{ty}`"))
}
FirehoseError::Cbor(err) => {
FirehoseFailure::new("decode", format!("cbor decode error: {err}"))
}
FirehoseError::StreamClosed { code, reason } => {
FirehoseFailure::new("stream_closed", format!("close code {code}: {reason}"))
}
FirehoseError::TcpDropped => FirehoseFailure::new("tcp_dropped", "tcp layer dropped"),
FirehoseError::ConnectTimeout => {
FirehoseFailure::new("connect_timeout", "connect timed out")
}
FirehoseError::Silent(silent) => {
FirehoseFailure::new("silent", format!("no frame or pong for {silent:?}"))
}
FirehoseError::FutureCursor { message } => FirehoseFailure::new(
"future_cursor",
message
.as_deref()
.map_or_else(|| "future cursor".to_owned(), |message| message.to_owned()),
),
}
}
fn classify_websocket_error(err: &tokio_websockets::Error) -> FirehoseFailure {
match err {
tokio_websockets::Error::CannotResolveHost => {
FirehoseFailure::new("dns", "host could not be resolved")
}
tokio_websockets::Error::Io(err) => {
let kind = match err.kind() {
ErrorKind::ConnectionRefused => "tcp_refused",
ErrorKind::ConnectionReset => "tcp_reset",
ErrorKind::ConnectionAborted => "tcp_aborted",
ErrorKind::TimedOut => "tcp_timeout",
ErrorKind::UnexpectedEof => "tcp_eof",
ErrorKind::AddrInUse => "tcp_addr_in_use",
ErrorKind::AddrNotAvailable => "tcp_addr_not_available",
ErrorKind::PermissionDenied => "tcp_permission",
_ => "io",
};
FirehoseFailure::new(kind, format!("{err}"))
}
tokio_websockets::Error::InvalidDNSName(_) => {
FirehoseFailure::new("tls_invalid_dns_name", format!("{err}"))
}
tokio_websockets::Error::Rustls(_) => FirehoseFailure::new("tls", format!("{err}")),
tokio_websockets::Error::Upgrade(upgrade) => {
let kind = match upgrade {
tokio_websockets::upgrade::Error::DidNotSwitchProtocols(_) => "http_upgrade",
_ => "websocket_upgrade",
};
FirehoseFailure::new(kind, format!("{upgrade}"))
}
tokio_websockets::Error::UnsupportedScheme => {
FirehoseFailure::new("config", "unsupported websocket URL scheme")
}
tokio_websockets::Error::Protocol(err) => {
FirehoseFailure::new("websocket_protocol", format!("{err}"))
}
tokio_websockets::Error::PayloadTooLong { len, max_len } => FirehoseFailure::new(
"websocket_payload",
format!("payload length {len} > {max_len}"),
),
_ => FirehoseFailure::new("websocket", format!("{err}")),
}
}
fn message_stats(msg: &SubscribeReposMessage<'_>) -> (&'static str, Option<i64>) {
match msg {
SubscribeReposMessage::Commit(commit) => ("commit", Some(commit.seq)),
SubscribeReposMessage::Sync(sync) => ("sync", Some(sync.seq)),
SubscribeReposMessage::Identity(identity) => ("identity", Some(identity.seq)),
SubscribeReposMessage::Account(account) => ("account", Some(account.seq)),
SubscribeReposMessage::Info(_) => ("info", None),
}
}
trait AddJitter: rand::Rng {
fn add_jitter(&mut self, timeout: Duration) -> Duration {
let timeout_secs = timeout.as_secs_f32();
let amt = timeout_secs * self.random_range(-0.12..0.12);
Duration::from_secs_f32(timeout_secs + amt)
}
}
impl<R: rand::Rng> AddJitter for R {}
/// resume from our own cursor, otherwise live tail like the reference relay.
/// starting a PDS we have no cursor for at 0 would replay its whole retained
/// window, which is days of events on a large host.
fn select_start_cursor(local_cursor: Option<i64>) -> Option<i64> {
local_cursor.filter(|cursor| *cursor > 0)
}
pub struct FirehoseIngestor {
state: Arc<AppState>,
buffer_tx: BufferTx,
relay_host: Url,
is_pds: bool,
public: bool,
filter: FilterHandle,
enabled: watch::Receiver<bool>,
_verify_signatures: bool,
throttle: ThrottleHandle,
max_failures: usize,
stats: Arc<crate::ingest::firehose_stats::FirehoseSourceStats>,
/// shared with every ingestor this source ever gets, so frames a reconnect or a respawn
/// reads again queue up behind the ones still in flight from before
tracker: Arc<SeqTracker>,
}
impl FirehoseIngestor {
pub async fn new(
state: Arc<AppState>,
buffer_tx: BufferTx,
relay_host: Url,
is_pds: bool,
public: bool,
filter: FilterHandle,
enabled: watch::Receiver<bool>,
verify_signatures: bool,
max_failures: usize,
) -> Self {
let throttle = state.throttler.get_handle(&relay_host).await;
let stats = state.firehose_stats.handle(&relay_host).await;
let tracker = state.firehose_tracker(&relay_host).await;
Self {
state,
buffer_tx,
relay_host,
is_pds,
public,
filter,
enabled,
_verify_signatures: verify_signatures,
throttle,
max_failures,
stats,
tracker,
}
}
#[tracing::instrument(skip(self), fields(host = %self.relay_host))]
pub async fn run(mut self) -> Result<()> {
let host = self.relay_host.host_str().unwrap_or("");
let count_key = crate::db::keys::pds_account_count_key(host);
let mut rng: SmallRng = rand::make_rng();
loop {
if self.state.pds_meta.load().is_banned(host) {
break Ok(());
}
self.enabled.wait_enabled("firehose").await;
// get cursor
let local_cursor = self
.state
.firehose_cursors
.peek_with(&self.relay_host, |_, c| c.load(Ordering::SeqCst));
let start_cursor = select_start_cursor(local_cursor);
match start_cursor {
Some(c) => info!(cursor = %c, "resuming from cursor"),
None => info!("no cursor found, live tailing"),
}
self.stats.record_connect_attempt(start_cursor);
let host_status = self.is_pds.then(|| {
let meta = self.state.pds_meta.load();
meta.status(host).as_str()
});
debug!(
is_pds = self.is_pds,
cursor = ?start_cursor,
host_status = ?host_status,
consecutive_failures = self.throttle.consecutive_failures(),
throttled_until = self.throttle.throttled_until(),
"connecting to firehose"
);
let connect_started = Instant::now();
let mut stream =
match FirehoseStream::connect(self.relay_host.clone(), start_cursor, self.public)
.await
{
Ok(s) => s,
Err(e) => {
let failure = classify_firehose_error(&e);
self.stats.record_connect_error(failure.kind);
let timeout = self.on_failure(&failure, &mut rng).await;
let fmt = crate::util::format_duration(timeout);
error!(
err = %e,
failure_kind = failure.kind,
failure = %failure.detail,
failures = self.throttle.consecutive_failures(),
max_failures = self.max_failures,
in = %fmt,
"failed to connect to firehose, retrying later"
);
tokio::time::sleep(timeout).await;
continue;
}
};
info!(
elapsed_ms = connect_started.elapsed().as_millis(),
"firehose connected"
);
self.stats.record_connected(connect_started.elapsed());
let mut marked_active = false;
let active_sleep_secs = if cfg!(debug_assertions) { 1 } else { 60 };
let mut active_sleep =
std::pin::pin!(tokio::time::sleep(Duration::from_secs(active_sleep_secs)));
let mut reported_outdated_cursor = false;
let res = loop {
tokio::select! {
msg = stream.next() => {
let bytes = match msg {
Ok(b) => b,
Err(e) => break Err(e),
};
self.stats.record_frame(bytes.len());
match decode_frame(&bytes) {
Ok(msg) => {
let (kind, seq) = message_stats(&msg);
self.stats.record_decoded(kind, seq);
if let SubscribeReposMessage::Info(info) = &msg {
if info.name == InfoName::OutdatedCursor
&& !reported_outdated_cursor
{
reported_outdated_cursor = true;
warn!(
relay = %self.relay_host,
requested_cursor = ?start_cursor,
message = ?info.message,
"firehose cursor predates relay retention; consuming the oldest retained events"
);
}
continue;
}
if self.is_pds {
let tier = {
let meta = self.state.pds_meta.load();
let banned = meta.is_banned(host);
if banned {
break Ok(());
}
let override_name = meta.hosts.get(host).and_then(|h| h.tier.as_ref());
self.state.tier_policy.resolve(host, override_name)
};
let accounts = self.state.db.get_count(&count_key).await;
let throttle_started = crate::ingest::firehose_stats::StatsInstant::now();
let mut disabled = false;
loop {
tokio::select! {
_ = self.throttle.wait_for_allow(accounts, &tier) => {
self.stats.record_throttle_wait(throttle_started.elapsed());
break;
}
res = self.enabled.changed() => {
if res.is_err() || !*self.enabled.borrow() {
disabled = true;
break;
}
}
}
}
if disabled {
info!("firehose disabled, disconnecting");
break Ok(());
}
}
self.handle_message(msg).await;
},
Err(e) => match e {
// don't disconnect on unknown op or type
FirehoseError::UnknownOp(op) => {
warn!(op = %op, "unknown frame op");
continue;
},
FirehoseError::UnknownType(t) => {
warn!(ty = %t, "unknown frame type");
continue;
},
// everything else is a hard error
e => break Err(e),
},
}
}
_ = &mut active_sleep, if !marked_active => {
marked_active = true;
// only reset failure state once the stream has been healthy for a bit
// so we dont get in a "connects successfully, sends garbage" situation
self.throttle.record_success();
if !self.is_pds {
self.state.track_dispatched(&self.relay_host).await;
} else {
let (current_status, tier) = {
let meta = self.state.pds_meta.load();
let override_name = meta.hosts.get(host).and_then(|h| h.tier.as_ref());
(meta.status(host), self.state.tier_policy.resolve(host, override_name))
};
if current_status == HostStatus::Banned {
break Ok(());
}
let count = self.state.db.get_count_sync(&count_key);
let new_status = tier.account_limit.is_some_and(|l| count >= l)
.then_some(HostStatus::Throttled)
.unwrap_or(HostStatus::Active);
debug!(
host,
?current_status,
account_limit = ?tier.account_limit,
count,
?new_status,
"active_sleep: computed status transition"
);
if current_status != new_status
&& let Err(e) = self.set_host_status(new_status)
{
error!(err = %e, "failed to update host status");
}
}
}
_ = self.enabled.changed() => {
if !*self.enabled.borrow() {
info!("firehose disabled, disconnecting");
break Ok(());
}
}
}
};
match res {
Ok(()) => {}
// 1001 "going away" means the server shut down cleanly (restart, deploy, etc.).
// this is not a real failure so we intentionally skip on_failure / backoff and
// reconnect quickly. a malicious source could abuse this but they could equally
// just drop the TCP connection, which would hit the normal backoff path anyway.
Err(FirehoseError::StreamClosed { code: 1001, reason }) => {
self.stats.record_stream_error("stream_closed");
debug!(reason = %reason, "host gone away");
tokio::time::sleep(Duration::from_secs(1)).await;
}
Err(FirehoseError::FutureCursor { message }) => {
let failure = FirehoseFailure::new(
"future_cursor",
message
.as_deref()
.map_or_else(|| "future cursor".to_owned(), str::to_owned),
);
self.stats.record_stream_error(failure.kind);
if start_cursor.is_some() {
if self.is_pds
&& let Err(e) = self.set_host_status(HostStatus::Idle)
{
error!(err = %e, "failed to update host status to idle");
}
if let Err(e) = self.clear_stale_cursor() {
error!(err = %e, "failed to clear future cursor");
}
warn!(
relay = %self.relay_host,
requested_cursor = ?start_cursor,
message = ?message,
"relay rejected a future cursor; cleared stored cursor and retrying without it"
);
tokio::time::sleep(Duration::from_secs(1)).await;
continue;
}
let timeout = self.on_failure(&failure, &mut rng).await;
let fmt = crate::util::format_duration(timeout);
error!(
failure_kind = failure.kind,
failure = %failure.detail,
failures = self.throttle.consecutive_failures(),
max_failures = self.max_failures,
in = %fmt,
"relay reported a future cursor without a requested cursor, reconnecting later"
);
tokio::time::sleep(timeout).await;
}
Err(FirehoseError::RelayError { error, message }) => {
let message = message
.as_deref()
.map_or(Cow::Borrowed("<no message>"), Cow::Borrowed);
let failure =
FirehoseFailure::new("relay_error", format!("{error}: {message}"));
self.stats.record_stream_error(failure.kind);
error!(err = %error, "relay sent error: {message}");
let timeout = self.on_failure(&failure, &mut rng).await;
let fmt = crate::util::format_duration(timeout);
error!(
failure_kind = failure.kind,
failure = %failure.detail,
failures = self.throttle.consecutive_failures(),
max_failures = self.max_failures,
in = %fmt,
"firehose disconnected, reconnecting later"
);
tokio::time::sleep(timeout).await;
}
Err(e) => {
let failure = classify_firehose_error(&e);
self.stats.record_stream_error(failure.kind);
let timeout = self.on_failure(&failure, &mut rng).await;
let fmt = crate::util::format_duration(timeout);
error!(
err = %e,
failure_kind = failure.kind,
failure = %failure.detail,
failures = self.throttle.consecutive_failures(),
max_failures = self.max_failures,
in = %fmt,
"firehose stream error, reconnecting later"
);
tokio::time::sleep(timeout).await;
}
}
}
}
async fn on_failure(&self, failure: &FirehoseFailure, rng: &mut SmallRng) -> Duration {
let secs = self
.throttle
.record_failure_detail(failure.kind, failure.detail.clone())
.unwrap_or_else(|| {
let until = self.throttle.throttled_until();
0.max(until - chrono::Utc::now().timestamp()) as u64
});
let failures = self.throttle.consecutive_failures();
if failures == self.max_failures {
warn!(
failures,
max_failures = self.max_failures,
failure_kind = failure.kind,
failure = %failure.detail,
"too many consecutive failures, marking host offline"
);
if !self.is_pds {
self.state.forget_dispatched(&self.relay_host).await;
} else if let Err(e) = self.set_host_status(HostStatus::Offline) {
error!(err = %e, "failed to update host status to offline");
}
}
let Some(past) = failures.checked_sub(self.max_failures) else {
return rng.add_jitter(Duration::from_secs(secs).min(MAX_BACKOFF));
};
let backoff = MAX_BACKOFF
.saturating_mul(2u32.pow(past.min(6) as u32 + 1))
.min(MAX_OFFLINE_BACKOFF);
// hosts that went down together would otherwise all retry at once
rng.random_range(backoff / 2..=backoff)
}
fn set_host_status(&self, status: HostStatus) -> Result<()> {
let Some(host) = self.relay_host.host_str() else {
return Ok(());
};
debug!(host = %host, status = ?status, "updating host status");
let _status_write = self.state.pds_status_write.lock();
let mut batch = self.state.db.inner.batch();
crate::db::pds_meta::set_status(&mut batch, &self.state.db.filter, host, status)?;
batch.commit().into_diagnostic()?;
crate::pds_meta::PdsMeta::update_host(&self.state.pds_meta, host, |h| h.status = status);
Ok(())
}
fn clear_stale_cursor(&self) -> Result<()> {
crate::db::set_firehose_cursor(&self.state.db, &self.relay_host, -1)?;
self.state
.firehose_cursors
.peek_with(&self.relay_host, |_, cursor| {
cursor.store(-1, Ordering::SeqCst)
});
Ok(())
}
async fn handle_message(&self, msg: SubscribeReposMessage<'_>) {
let (Some(did), Some((seq, time))) = (msg.did(), msg.position()) else {
return;
};
let ticket = self.tracker.ticket(seq, time.0.timestamp_millis());
let should_process_started = crate::ingest::firehose_stats::StatsInstant::now();
let process = self
.should_process(did)
.await
.inspect_err(|e| error!(did = %did, err = %e, "failed to check if we should process"))
.unwrap_or(false);
self.stats
.record_should_process(should_process_started.elapsed());
if !process {
self.stats.record_skipped();
trace!(did = %did, "skipping: not in filter");
cursor::done(&self.state, &self.relay_host, ticket);
return;
}
trace!(did = %did, "forwarding message to ingest buffer");
let send_started = crate::ingest::firehose_stats::StatsInstant::now();
let res = self
.buffer_tx
.send(IngestMessage::Firehose {
url: self.relay_host.clone(),
is_pds: self.is_pds,
msg: msg.into_static(),
ticket,
})
.await;
{
let elapsed = send_started.elapsed();
if res.is_ok() {
self.stats.record_forwarded(elapsed);
} else {
self.stats.record_forward_error(elapsed);
}
}
if let Err(e) = res {
error!(err = %e, "failed to send message to buffer processor");
}
}
async fn should_process(&self, did: &Did) -> Result<bool> {
let filter = self.filter.load();
let state = self.state.clone();
let did = did.clone().into_static();
let span = Span::current();
tokio::task::spawn_blocking(move || {
let _entered = span.entered();
let _entered = tracing::info_span!("should_process", repo = %did).entered();
let excl_key = crate::db::filter::exclude_key(did.as_str())?;
if state.db.filter.contains_key(&excl_key).into_diagnostic()? {
return Ok(false);
}
match filter.mode {
FilterMode::Full => Ok(true),
FilterMode::Filter => {
let metadata_key = crate::db::keys::repo_metadata_key(&did);
if let Some(bytes) = state
.db
.repo_metadata
.get(&metadata_key)
.into_diagnostic()?
{
let metadata = crate::db::deser_repo_meta(bytes.as_ref())?;
if metadata.tracked {
trace!(did = %did, "tracked repo, processing");
return Ok(true);
} else {
debug!(did = %did, "known but explicitly untracked, skipping");
return Ok(false);
}
}
if !filter.signals.is_empty() {
trace!(did = %did, "unknown, passing to worker for signal check");
Ok(true)
} else {
trace!(did = %did, "unknown and no signals configured, skipping");
Ok(false)
}
}
}
})
.await
.into_diagnostic()
.flatten()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::Config;
use crate::control::Hydrant;
use crate::filter::FilterConfig;
use crate::ingest::stream::{Datetime, Identity, Info};
use futures::{SinkExt as _, StreamExt as _};
use serde::Serialize;
use std::sync::atomic::AtomicI64;
use tokio_websockets::{Message as WsMessage, ServerBuilder};
#[derive(Serialize)]
struct TestFrameHeader<'a> {
op: i64,
#[serde(skip_serializing_if = "Option::is_none")]
t: Option<&'a str>,
}
fn test_frame<T: Serialize>(op: i64, ty: Option<&str>, body: &T) -> Vec<u8> {
let mut bytes = serde_ipld_dagcbor::to_vec(&TestFrameHeader { op, t: ty }).unwrap();
bytes.extend_from_slice(&serde_ipld_dagcbor::to_vec(body).unwrap());
bytes
}
fn test_identity(seq: i64) -> Identity {
Identity {
did: Did::new_static("did:plc:ewvi7nxzyoun6zhxrhs64oiz").unwrap(),
handle: None,
seq,
time: serde_json::from_str::<Datetime>(r#""2026-08-10T12:34:56Z""#).unwrap(),
}
}
fn outdated_cursor_frame() -> Vec<u8> {
test_frame(
1,
Some("#info"),
&Info {
message: Some("cursor is outside the retained window".into()),
name: InfoName::OutdatedCursor,
},
)
}
#[tokio::test]
async fn a_host_past_max_failures_goes_offline_and_is_retried_up_to_hourly() -> Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let hydrant = Hydrant::new(Config {
database_path: tmp.path().to_path_buf(),
..Config::default()
})
.await?;
let host = "flaky.example";
let ingestor = FirehoseIngestor::new(
hydrant.state.clone(),
BufferTx::channel(1).0,
Url::parse("wss://flaky.example/").into_diagnostic()?,
true,
false,
hydrant.state.filter.clone(),
hydrant.state.firehose_enabled.subscribe(),
false,
2,
)
.await;
let failure = FirehoseFailure::new("test", "down");
let mut rng: SmallRng = rand::make_rng();
let mut waits = Vec::new();
for _ in 0..20 {
waits.push(ingestor.on_failure(&failure, &mut rng).await);
}
assert!(
waits[0] <= MAX_BACKOFF,
"before max_failures the ref relay backoff applies"
);
assert_eq!(
hydrant.state.pds_meta.load().status(host),
HostStatus::Offline
);
assert!((MAX_BACKOFF..=2 * MAX_BACKOFF).contains(&waits[1]));
assert!(waits.iter().all(|w| *w <= MAX_OFFLINE_BACKOFF));
assert!(
waits[19] >= MAX_OFFLINE_BACKOFF / 2,
"and it keeps retrying"
);
Ok(())
}
#[tokio::test]
async fn operational_transition_does_not_clear_an_operator_ban() -> Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let cfg = Config {
database_path: tmp.path().to_path_buf(),
..Config::default()
};
let hydrant = Hydrant::new(cfg.clone()).await?;
hydrant.firehose.startable_for_test();
let host = "banned-transition.example";
let relay_host = Url::parse("wss://banned-transition.example/").into_diagnostic()?;
hydrant.pds.ban(host).await?;
let ingestor = FirehoseIngestor::new(
hydrant.state.clone(),
BufferTx::channel(1).0,
relay_host,
true,
false,
hydrant.state.filter.clone(),
hydrant.state.firehose_enabled.subscribe(),
false,
1,
)
.await;
ingestor.set_host_status(HostStatus::Offline)?;
ingestor.set_host_status(HostStatus::Active)?;
assert_eq!(
hydrant.state.pds_meta.load().status(host),
HostStatus::Active
);
assert!(hydrant.state.pds_meta.load().is_banned(host));
drop(ingestor);
drop(hydrant);
let reopened = Hydrant::new(cfg).await?;
assert_eq!(
reopened.state.pds_meta.load().status(host),
HostStatus::Active
);
assert!(reopened.state.pds_meta.load().is_banned(host));
assert_eq!(
reopened
.firehose
.add_source_validated(
Url::parse("wss://banned-transition.example/").into_diagnostic()?
)
.await?,
Err(crate::control::firehose::PdsAdmissionRejection::HostBanned)
);
Ok(())
}
/// a source whose ingestor is replaced (restart, remove and add again) while frames the old
/// one read are still in the shards must not get a second watermark for the same cursor
#[tokio::test]
async fn a_replaced_ingestor_waits_for_the_frames_of_the_old_one() -> Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let cfg = Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
};
let state = Arc::new(AppState::new(&cfg)?);
let relay = Url::parse("wss://relay.example/").into_diagnostic()?;
let ingestor = || {
FirehoseIngestor::new(
state.clone(),
BufferTx::channel(1).0,
relay.clone(),
false,
false,
state.filter.clone(),
state.firehose_enabled.subscribe(),
false,
1,
)
};
let first = ingestor().await;
let in_flight = first.tracker.ticket(1, 0);
drop(first);
let second = ingestor().await;
let mut moved = None;
second
.tracker
.ticket(2, 0)
.done(|watermark| moved = Some(watermark.seq));
assert_eq!(
moved, None,
"the cursor passed a frame the first ingestor read"
);
assert_eq!(second.tracker.pending(), 2);
in_flight.done(|_| {});
assert_eq!(second.tracker.pending(), 0);
Ok(())
}
#[tokio::test]
async fn clear_stale_cursor_resets_db_and_memory() -> Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let cfg = Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
};
let state = Arc::new(AppState::new(&cfg)?);
let relay_host = Url::parse("wss://example.com/").into_diagnostic()?;
crate::db::set_firehose_cursor(&state.db, &relay_host, 1234)?;
let _ = state
.firehose_cursors
.insert_async(relay_host.clone(), AtomicI64::new(1234))
.await;
let ingestor = FirehoseIngestor::new(
state.clone(),
BufferTx::channel(1).0,
relay_host.clone(),
true,
false,
state.filter.clone(),
state.firehose_enabled.subscribe(),
false,
1,
)
.await;
ingestor.clear_stale_cursor()?;
assert_eq!(
crate::db::get_firehose_cursor(&state.db, &relay_host).await?,
Some(-1)
);
let in_memory = state
.firehose_cursors
.peek_with(&relay_host, |_, cursor| cursor.load(Ordering::SeqCst))
.unwrap_or(0);
assert_eq!(in_memory, -1);
Ok(())
}
#[test]
fn select_start_cursor_cases() {
assert_eq!(select_start_cursor(Some(1234)), Some(1234));
assert_eq!(select_start_cursor(Some(0)), None);
assert_eq!(select_start_cursor(None), None);
assert_eq!(select_start_cursor(Some(-1)), None);
}
#[tokio::test]
async fn new_pds_source_connects_at_live_head() -> Result<()> {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.into_diagnostic()?;
let addr = listener.local_addr().into_diagnostic()?;
let (query_tx, query_rx) = tokio::sync::oneshot::channel();
let server = tokio::spawn(async move {
let (tcp, _) = listener.accept().await.unwrap();
let (request, mut ws) = ServerBuilder::new().accept(tcp).await.unwrap();
let _ = query_tx.send(request.uri().query().map(str::to_owned));
let _ = ws.next().await;
});
let tmp = tempfile::tempdir().into_diagnostic()?;
let cfg = Config {
database_path: tmp.path().to_path_buf(),
..Config::default()
};
let state = Arc::new(AppState::new(&cfg)?);
let pds = Url::parse(&format!("http://{addr}/")).into_diagnostic()?;
let ingestor = FirehoseIngestor::new(
state.clone(),
BufferTx::channel(1).0,
pds,
true,
false,
state.filter.clone(),
state.firehose_enabled.subscribe(),
false,
2,
)
.await;
let run = tokio::spawn(ingestor.run());
let query = tokio::time::timeout(Duration::from_secs(2), query_rx)
.await
.into_diagnostic()?
.into_diagnostic()?;
assert_eq!(
query, None,
"a PDS we have no cursor for must not replay its retained history"
);
run.abort();
let _ = run.await;
server.abort();
Ok(())
}
#[tokio::test]
async fn in_window_resume_requests_cursor_and_receives_replay() -> Result<()> {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.into_diagnostic()?;
let addr = listener.local_addr().into_diagnostic()?;
let server = tokio::spawn(async move {
let (tcp, _) = listener.accept().await.unwrap();
let (request, mut ws) = ServerBuilder::new().accept(tcp).await.unwrap();
assert_eq!(
request.uri().path(),
"/xrpc/com.atproto.sync.subscribeRepos"
);
assert_eq!(request.uri().query(), Some("cursor=42"));
ws.send(WsMessage::binary(test_frame(
1,
Some("#identity"),
&test_identity(42),
)))
.await
.unwrap();
});
let relay = Url::parse(&format!("http://{addr}/ignored")).into_diagnostic()?;
let mut stream = FirehoseStream::connect(relay, Some(42), false).await?;
let bytes = tokio::time::timeout(Duration::from_secs(2), stream.next())
.await
.into_diagnostic()??;
let SubscribeReposMessage::Identity(identity) = decode_frame(&bytes)? else {
panic!("expected identity replay");
};
assert_eq!(identity.seq, 42);
server.await.into_diagnostic()?;
Ok(())
}
#[tokio::test]
async fn outdated_cursor_info_keeps_stream_open_and_forwards_replay() -> Result<()> {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.into_diagnostic()?;
let addr = listener.local_addr().into_diagnostic()?;
let server = tokio::spawn(async move {
let (tcp, _) = listener.accept().await.unwrap();
let (request, mut ws) = ServerBuilder::new().accept(tcp).await.unwrap();
assert_eq!(request.uri().query(), Some("cursor=42"));
ws.send(WsMessage::binary(outdated_cursor_frame()))
.await
.unwrap();
ws.send(WsMessage::binary(test_frame(
1,
Some("#identity"),
&test_identity(100),
)))
.await
.unwrap();
let _ = ws.next().await;
});
let tmp = tempfile::tempdir().into_diagnostic()?;
let cfg = Config {
database_path: tmp.path().to_path_buf(),
..Config::default()
};
let state = Arc::new(AppState::new(&cfg)?);
state
.filter
.store(Arc::new(FilterConfig::new(FilterMode::Full)));
let relay = Url::parse(&format!("http://{addr}/")).into_diagnostic()?;
let _ = state
.firehose_cursors
.insert_async(relay.clone(), AtomicI64::new(42))
.await;
let (buffer_tx, mut rxs) = BufferTx::channel(1);
let ingestor = FirehoseIngestor::new(
state.clone(),
buffer_tx,
relay.clone(),
false,
false,
state.filter.clone(),
state.firehose_enabled.subscribe(),
false,
2,
)
.await;
let run = tokio::spawn(ingestor.run());
let replay = tokio::time::timeout(Duration::from_secs(2), rxs[0].recv())
.await
.into_diagnostic()?
.ok_or_else(|| miette::miette!("firehose buffer closed before replay arrived"))?;
let IngestMessage::Firehose {
msg: SubscribeReposMessage::Identity(identity),
..
} = replay
else {
panic!("expected retained identity event after OutdatedCursor info");
};
assert_eq!(identity.seq, 100);
assert_eq!(
state
.firehose_cursors
.peek_with(&relay, |_, cursor| cursor.load(Ordering::SeqCst)),
Some(42),
"receiving a frame must not bypass the downstream durable checkpoint"
);
run.abort();
let _ = run.await;
tokio::time::timeout(Duration::from_secs(2), server)
.await
.into_diagnostic()?
.into_diagnostic()?;
Ok(())
}
}