use chrono::{DateTime, TimeDelta, Utc};
use fjall::OwnedWriteBatch;
use futures::{Future, FutureExt};
use jacquard_common::{IntoStatic, types::{nsid::Nsid, string::Did}};
#[derive(Deserialize)]
struct DescribeRepoOutput {
collections: Vec<Nsid>,
}
use miette::{IntoDiagnostic, Result};
use reqwest::StatusCode;
use serde::{Deserialize, Serialize};
use smol_str::{SmolStr, ToSmolStr};
use std::collections::{HashMap, HashSet};
use std::ops::{Add, Sub};
use std::sync::Arc;
use std::time::Duration;
use tracing::{Instrument, error, info_span, trace, warn};
use url::Url;
use super::super::InFlightGuard;
use crate::db::{Db, keys};
use crate::net::{PublicEndpoint, PublicHttpClient};
use crate::state::AppState;
use crate::util::throttle::{OrFailure, ThrottleHandle, Throttler};
use crate::util::{
ErrorForStatus, RetryOutcome, RetryWithBackoff, is_io_error_their_fault, is_status_their_fault,
is_tls_cert_error, parse_retry_after,
};
pub(super) const MAX_RETRY_ATTEMPTS: u32 = 5;
#[derive(Debug, Serialize, Deserialize, Clone)]
pub(crate) struct RetryState {
pub(crate) after: DateTime<Utc>,
pub(crate) duration: TimeDelta,
pub(crate) attempts: u32,
#[serde(serialize_with = "crate::util::ser_status_code")]
#[serde(deserialize_with = "crate::util::deser_status_code")]
pub(crate) status: Option<StatusCode>,
}
impl RetryState {
pub(crate) fn new(secs: i64) -> Self {
let duration = TimeDelta::seconds(secs);
Self {
duration,
after: Utc::now().add(duration),
attempts: 0,
status: None,
}
}
/// returns the next retry state or `None` if the attempt count would reach the cap.
pub(crate) fn next_attempt(self) -> Option<Self> {
let attempts = self.attempts + 1;
if attempts >= MAX_RETRY_ATTEMPTS {
return None;
}
let duration = self.duration * 2;
Some(Self {
after: Utc::now().add(duration),
duration,
attempts,
status: None,
})
}
pub(crate) fn with_status(mut self, code: StatusCode) -> Self {
self.status = Some(code);
self
}
}
trait ToRetryState {
fn to_retry_state(&self) -> RetryState;
}
impl ToRetryState for ThrottleHandle {
fn to_retry_state(&self) -> RetryState {
let after = chrono::DateTime::from_timestamp_secs(self.throttled_until()).unwrap();
RetryState {
duration: after.sub(Utc::now()),
after,
attempts: 0,
status: None,
}
}
}
pub(super) enum CrawlCheckResult {
Signal,
NoSignal,
Retry(RetryState),
}
impl From<RetryState> for CrawlCheckResult {
fn from(value: RetryState) -> Self {
Self::Retry(value)
}
}
fn is_throttle_worthy(e: &reqwest::Error) -> bool {
use std::error::Error;
if e.is_timeout() {
return true;
}
let mut src = e.source();
while let Some(s) = src {
if let Some(e) = s.downcast_ref::<std::io::Error>()
&& (is_io_error_their_fault(e) || is_tls_cert_error(e))
{
return true;
}
src = s.source();
}
e.status().is_some_and(|s| {
// lets not consider internal server errors for throttling
// a server might be having a hard time on one request, but not on the rest
s.as_u16() != 500 && s.as_u16() != 404 && is_status_their_fault(s.as_u16())
})
}
/// shared describeRepo signal-checking logic used by both relay and retry producers.
#[derive(Clone)]
pub(crate) struct SignalChecker {
pub(crate) http: PublicHttpClient,
pub(crate) state: Arc<AppState>,
pub(crate) throttler: Throttler,
}
impl SignalChecker {
fn check_repo_signals(
&self,
filter: Arc<crate::filter::FilterConfig>,
did: Did,
) -> impl Future<Output = (Did, CrawlCheckResult)> + Send + 'static {
let resolver = self.state.resolver.clone();
let http = self.http.clone();
let throttler = self.throttler.clone();
async move {
const MAX_RETRIES: u32 = 5;
let pds_url = (|| resolver.resolve_identity_info(&did))
.retry(MAX_RETRIES, |e, attempt| {
matches!(e, crate::resolver::ResolverError::Ratelimited)
.then(|| Duration::from_secs(1 << attempt.min(5)))
})
.await;
let endpoint = match pds_url {
Ok((url, _)) => match PublicEndpoint::parse_http(&url) {
Ok(endpoint) => endpoint,
Err(error) => {
warn!(%did, %error, "resolver returned an unsafe PDS endpoint");
return (did, CrawlCheckResult::NoSignal);
}
},
Err(RetryOutcome::Ratelimited) => {
error!(
retries = MAX_RETRIES,
"rate limited resolving identity, giving up"
);
return (did, RetryState::new(60).into());
}
Err(RetryOutcome::Failed(e)) => {
error!(err = %e, "failed to resolve identity");
return (did, RetryState::new(60).into());
}
};
let pds_url = endpoint.url().clone();
let throttle = throttler.get_handle(&pds_url).await;
if throttle.is_throttled() {
trace!(host = pds_url.host_str(), "skipping throttled pds");
return (did, throttle.to_retry_state().into());
}
let _permit = throttle.acquire().unit_error().or_failure(&throttle, || ());
let Ok(_permit) = _permit.await else {
trace!(
host = pds_url.host_str(),
"pds failed while waiting for permit"
);
return (did, throttle.to_retry_state().into());
};
enum RequestError {
Reqwest(reqwest::Error),
RateLimited(Option<u64>),
Throttled,
}
let mut describe_endpoint = endpoint.with_path("/xrpc/com.atproto.repo.describeRepo");
describe_endpoint.append_query_pair("repo", did.as_str());
let resp = async {
let resp = http
.get(&describe_endpoint)
.timeout(throttle.timeout())
.send()
.await
.map_err(RequestError::Reqwest)?;
if resp.status() == StatusCode::TOO_MANY_REQUESTS {
return Err(RequestError::RateLimited(parse_retry_after(&resp)));
}
resp.error_for_status().map_err(RequestError::Reqwest)
}
.or_failure(&throttle, || RequestError::Throttled)
.await;
let resp = match resp {
Ok(r) => {
throttle.record_success();
r
}
Err(RequestError::RateLimited(secs)) => {
throttle.record_ratelimit(secs);
return (
did,
throttle
.to_retry_state()
.with_status(StatusCode::TOO_MANY_REQUESTS)
.into(),
);
}
Err(RequestError::Throttled) => {
return (did, throttle.to_retry_state().into());
}
Err(RequestError::Reqwest(e)) => {
if e.is_timeout() && !throttle.record_timeout() {
let mut retry_state = RetryState::new(60);
retry_state.status = e.status();
return (did, retry_state.into());
}
if is_throttle_worthy(&e) {
if let Some(secs) = throttle.record_failure() {
warn!(url = %pds_url, secs, "throttling pds due to hard failure");
}
let mut retry_state = throttle.to_retry_state();
retry_state.status = e.status();
return (did, retry_state.into());
}
match e.status() {
Some(StatusCode::NOT_FOUND | StatusCode::GONE) => {
trace!("repo not found");
return (did, CrawlCheckResult::NoSignal);
}
Some(s) if s.is_client_error() => {
error!(status = %s, "repo unavailable");
return (did, CrawlCheckResult::NoSignal);
}
_ => {
error!(err = %e, "repo errored");
let mut retry_state = RetryState::new(60 * 15);
retry_state.status = e.status();
return (did, retry_state.into());
}
}
}
};
let bytes = match crate::util::read_limited_bytes(resp, 2 * 1024 * 1024).await {
Ok(b) => b,
Err(e) => {
error!(err = %e, "failed to read describeRepo response");
return (did, RetryState::new(60 * 5).into());
}
};
let out = match serde_json::from_slice::<DescribeRepoOutput>(&bytes) {
Ok(out) => out,
Err(e) => {
error!(err = %e, "failed to parse describeRepo response");
return (did, RetryState::new(60 * 10).into());
}
};
let found_signal = out
.collections
.iter()
.any(|col| filter.matches_signal(col.as_str()));
if !found_signal {
trace!("no signal-matching collections found");
}
(
did,
found_signal
.then_some(CrawlCheckResult::Signal)
.unwrap_or(CrawlCheckResult::NoSignal),
)
}
}
/// checks signals for each DID in `in_flight`, writes retry/no-signal state changes
/// into `batch`, and returns a vec containing only the guards for confirmed DIDs.
///
/// guards for non-confirmed DIDs are dropped here, releasing their in-flight slots.
/// confirmed DIDs' retry entries are NOT removed here. the worker removes them
/// atomically with the pending insert.
pub(crate) async fn check_signals_batch(
&self,
in_flight: Vec<InFlightGuard>,
filter: &Arc<crate::filter::FilterConfig>,
batch: &mut OwnedWriteBatch,
existing: &HashMap<Did, RetryState>,
) -> Result<Vec<InFlightGuard>> {
let db = &self.state.db;
let mut valid: Vec<Did> = Vec::new();
let mut set = tokio::task::JoinSet::new();
for guard in &in_flight {
let filter = filter.clone();
let did = guard.did.clone();
let span = tracing::info_span!("signals", did = %did);
set.spawn(self.check_repo_signals(filter, did).instrument(span));
}
while let Some(res) = tokio::time::timeout(Duration::from_secs(60), set.join_next())
.await
.into_diagnostic()
.map_err(|_| {
error!("signal check task timed out after 60s");
miette::miette!("signal check task timed out")
})?
{
let (did, result) = match res {
Ok(inner) => inner,
Err(e) => {
error!(err = ?e, "signal check panicked");
continue;
}
};
match result {
CrawlCheckResult::Signal => {
valid.push(did);
}
CrawlCheckResult::NoSignal => {
batch.remove(&db.crawler, keys::crawler_retry_key(&did));
}
CrawlCheckResult::Retry(state) => {
let prev_attempts = existing.get(&did).map(|s| s.attempts).unwrap_or(0);
let carried = RetryState {
attempts: prev_attempts,
..state
};
let next = match carried.next_attempt() {
Some(next) => next,
None => RetryState {
attempts: MAX_RETRY_ATTEMPTS,
..state
},
};
batch.insert(
&db.crawler,
keys::crawler_retry_key(&did),
rmp_serde::to_vec(&next).into_diagnostic()?,
);
}
}
}
let valid_set: HashSet<&Did> = valid.iter().collect();
Ok(in_flight
.into_iter()
.filter(|g| valid_set.contains(g.as_did()))
.collect())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn retry_state_persisted_golden_fixture() {
let state = RetryState {
after: DateTime::parse_from_rfc3339("2026-08-10T12:34:56Z")
.unwrap()
.with_timezone(&Utc),
duration: TimeDelta::seconds(90),
attempts: 3,
status: Some(StatusCode::TOO_MANY_REQUESTS),
};
let encoded = rmp_serde::to_vec(&state).unwrap();
assert_eq!(
crate::util::hex::encode(&encoded),
"94b4323032362d30382d31305431323a33343a35365a925a0003cd01ad"
);
let decoded: RetryState = rmp_serde::from_slice(&encoded).unwrap();
assert_eq!(decoded.after, state.after);
assert_eq!(decoded.duration, state.duration);
assert_eq!(decoded.attempts, 3);
assert_eq!(decoded.status, Some(StatusCode::TOO_MANY_REQUESTS));
}
}