use crate::backfill::admission::BackfillAdmission;
use crate::backfill::client::ThrottledHttpClient;
use crate::backfill::error::BackfillError;
use crate::config::BackfillStrategy;
use crate::db::types::TrimmedDid;
use crate::ingest::indexer::IndexerTx;
use crate::net::public_http_with_timeout;
use crate::state::AppState;
use crate::util::WatchEnabledExt;
use jacquard_common::IntoStatic;
use jacquard_common::types::did::Did;
use miette::Result;
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::{Duration, Instant};
use tokio::sync::Semaphore;
use tracing::{Instrument, debug, error, info};
pub(crate) mod apply;
pub mod process;
pub(crate) mod records;
pub mod task;
pub struct BackfillWorker {
state: Arc<AppState>,
buffer_tx: IndexerTx,
http: ThrottledHttpClient,
semaphore: Arc<Semaphore>,
concurrency_limit: usize,
admission: BackfillAdmission,
verify_signatures: bool,
strategy: BackfillStrategy,
in_flight: Arc<scc::HashSet<Did>>,
deferred: Arc<scc::HashSet<Did>>,
enabled: tokio::sync::watch::Receiver<bool>,
}
impl BackfillWorker {
pub fn new(
state: Arc<AppState>,
buffer_tx: IndexerTx,
timeout: Duration,
concurrency_limit: usize,
verify_signatures: bool,
strategy: BackfillStrategy,
proxies: Vec<url::Url>,
enabled: tokio::sync::watch::Receiver<bool>,
) -> Result<Self> {
// A CONNECT proxy resolves the origin itself, so it cannot carry the
// socket-coupled public-PDS guarantee. Reject it at startup instead of
// silently changing the address policy or pacing budget.
if !proxies.is_empty() {
return Err(miette::miette!(
"backfill proxies are unsupported for public-PDS egress; remove HYDRANT_BACKFILL_PROXIES"
));
}
let client = public_http_with_timeout(timeout);
info!("direct-only backfill client ready");
let admission =
BackfillAdmission::new(concurrency_limit, state.throttler.per_pds_concurrency());
Ok(Self {
state: state.clone(),
buffer_tx,
http: ThrottledHttpClient::new_with_body_limit(
client,
state.throttler.clone(),
state.max_car_body_bytes,
),
semaphore: Arc::new(Semaphore::new(concurrency_limit)),
concurrency_limit,
admission,
verify_signatures,
strategy,
in_flight: Arc::new(scc::HashSet::new()),
deferred: Arc::new(scc::HashSet::new()),
enabled,
})
}
}
struct Progress {
pending: u64,
busy: usize,
workers: usize,
finished: usize,
elapsed: Duration,
crawler_throttled: bool,
}
impl Progress {
fn log(&self) {
if self.pending == 0 && self.busy == 0 && self.finished == 0 {
return;
}
let secs = self.elapsed.as_secs_f64();
let rate = if secs > 0.0 {
self.finished as f64 / secs
} else {
0.0
};
info!(
pending = self.pending,
busy = format_args!("{}/{}", self.busy, self.workers),
finished = self.finished,
rate = format_args!("{rate:.2}/s"),
crawler_throttled = self.crawler_throttled,
"backfill progress"
);
}
}
struct InFlightGuard {
did: Did,
set: Arc<scc::HashSet<Did>>,
}
impl Drop for InFlightGuard {
fn drop(&mut self) {
let _ = self.set.remove_sync(&self.did);
}
}
impl BackfillWorker {
/// logs the queue, busy workers and finished repos once a minute, so a crawler sitting
/// throttled on a full queue for hours can be told apart from a worker that stopped
pub fn progress(&self) -> impl Future<Output = ()> + Send + 'static {
let state = self.state.clone();
let semaphore = self.semaphore.clone();
let workers = self.concurrency_limit;
async move {
let mut last_time = Instant::now();
let mut interval = tokio::time::interval(Duration::from_secs(60));
loop {
interval.tick().await;
let now = Instant::now();
Progress {
// exact like /stats, since an estimate goes stale once the queue empties
pending: state.db.get_count("pending").await,
// every task holds its permit until it ends
busy: workers.saturating_sub(semaphore.available_permits()),
workers,
finished: state.backfill_drained.swap(0, Ordering::Relaxed),
elapsed: now.duration_since(last_time),
crawler_throttled: state.crawler_throttled.load(Ordering::Relaxed),
}
.log();
last_time = now;
}
}
}
pub async fn run(mut self) {
info!("backfill worker started");
loop {
self.enabled.wait_enabled("backfill").await;
let mut spawned = 0;
for guard in self.state.db.indexer.pending_iter() {
let (key, value) = match guard.into_inner() {
Ok(kv) => kv,
Err(e) => {
error!(err = %e, "failed to read pending entry");
self.state.db.poison.check(&e);
continue;
}
};
let did = match TrimmedDid::try_from(value.as_ref()) {
Ok(d) => d.to_did(),
Err(e) => {
error!(err = %e, "invalid did in pending value");
continue;
}
};
// a task that ends in a stop leaves its key queued for the next start, so
// spawning on would only fetch the same repos again
if self.state.stop.is_set() {
return;
}
// check before trying to acquire a permit so we dont acquire a permit
// for no reason, the read will be cheap anyhow
if self.in_flight.contains_sync(&did) || self.deferred.contains_sync(&did) {
continue;
}
let permit = match self.semaphore.clone().try_acquire_owned() {
Ok(p) => p,
Err(_) => break,
};
// only mark as in flight if we can acquire a permit
if self
.in_flight
.insert_sync(did.clone().into_static())
.is_err()
{
// a task is already running, weh
// so we don't need this one anymore...
break;
}
let guard = InFlightGuard {
did: did.clone().into_static(),
set: self.in_flight.clone(),
};
let state = self.state.clone();
let http = self.http.clone();
let did = did.clone();
let buffer_tx = self.buffer_tx.clone();
let verify = self.verify_signatures;
let strategy = self.strategy;
let admission = self.admission.clone();
let deferred = self.deferred.clone();
let span = tracing::info_span!("backfill", did = %did);
tokio::spawn(
async move {
let _guard = guard;
let res = task::did_task(
&state, http, buffer_tx, &did, key, admission, permit, verify, strategy,
)
.await;
match &res {
Ok(task::TaskDisposition::Deferred) => {
let _ = deferred.insert_sync(did.clone().into_static());
}
Ok(task::TaskDisposition::Finished) | Err(_) => {
deferred.clear_sync();
}
}
if let Err(e) = res {
match &e {
BackfillError::Ratelimited
| BackfillError::PreemptivelyThrottled => {
debug!(err = %e, "process failed");
}
_ => {
error!(err = %e, "process failed");
}
}
if let BackfillError::Generic(report) = &e {
state.db.poison.check_report(report);
}
}
// wake worker to pick up more (in case we were sleeping at limit)
state.backfill_notify.notify_one();
}
.instrument(span),
);
spawned += 1;
}
if spawned == 0 {
// wait for new tasks
self.state.backfill_notify.notified().await;
}
// loop immediately since we might have more tasks
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::Config;
use crate::db::keys;
use crate::ingest::indexer::FirehoseWorker;
use miette::IntoDiagnostic;
fn progress_logs(progress: Progress) -> String {
crate::util::capture_logs("info", || progress.log())
}
#[test]
fn progress_tells_a_full_queue_from_a_stalled_worker() {
let progress = |busy, finished| Progress {
pending: 5000,
busy,
workers: 8,
finished,
elapsed: Duration::from_secs(60),
crawler_throttled: true,
};
let draining = progress_logs(progress(8, 120));
assert!(draining.contains("pending=5000"), "{draining}");
assert!(draining.contains("busy=8/8"), "{draining}");
assert!(draining.contains("finished=120"), "{draining}");
assert!(draining.contains("rate=2.00/s"), "{draining}");
assert!(draining.contains("crawler_throttled=true"), "{draining}");
let stalled = progress_logs(progress(0, 0));
assert!(stalled.contains("busy=0/8"), "{stalled}");
assert!(stalled.contains("finished=0"), "{stalled}");
}
#[test]
fn progress_stays_quiet_with_nothing_queued_or_running() {
let idle = Progress {
pending: 0,
busy: 0,
workers: 8,
finished: 0,
elapsed: Duration::from_secs(60),
crawler_throttled: false,
};
let last_task = Progress { busy: 1, ..idle };
assert_eq!(progress_logs(idle), "");
assert!(progress_logs(last_task).contains("busy=1/8"));
}
#[tokio::test]
async fn a_stopped_worker_leaves_the_queue_alone() -> Result<()> {
let tmp = tempfile::tempdir().into_diagnostic()?;
let state = Arc::new(AppState::new(&Config {
database_path: tmp.path().to_path_buf(),
..Default::default()
})?);
let did = Did::new_static("did:web:stopped.example").into_diagnostic()?;
let mut batch = state.db.inner.batch();
state
.db
.indexer
.stage_pending_insert(&mut batch, &did, &keys::pending_key(7));
batch.commit().into_diagnostic()?;
state.stop.trigger();
let worker = BackfillWorker::new(
state.clone(),
FirehoseWorker::new(state.clone(), 1).0,
Duration::from_secs(1),
1,
false,
BackfillStrategy::Full,
Vec::new(),
state.backfill_enabled.subscribe(),
)?;
// a worker that kept going would hand the key to a task and never return
tokio::time::timeout(Duration::from_secs(5), worker.run())
.await
.into_diagnostic()?;
assert!(
state
.db
.indexer
.is_pending(&keys::pending_key(7))
.into_diagnostic()?
);
Ok(())
}
}