//! the network lookups a relay shard needs before it can run a did's messages. a shard hands
//! them to this pool and keeps going with every other did, so a host that hangs only holds up
//! the dids it serves
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::time::{Duration, Instant};
use jacquard_common::types::string::Did;
use parking_lot::Mutex;
use scc::HashCache;
use smol_str::SmolStr;
use tokio::sync::{Notify, Semaphore, mpsc};
use tracing::debug;
use url::Url;
use super::StatusProbe;
use crate::net::PublicHttpClient;
use crate::resolver::{MiniDoc, ResolveFailure, Resolver};
use crate::types::RepoState;
/// a job that hasn't answered by now answers with [`LookupError::TimedOut`]
pub(crate) const JOB_DEADLINE: Duration = Duration::from_secs(12);
/// the pds status probe of a new account, inside the job's deadline
const PROBE_TIMEOUT: Duration = Duration::from_secs(5);
const PROBE_MARGIN: Duration = Duration::from_millis(250);
const GLOBAL_PERMITS: usize = 32;
/// one burst of new accounts mustn't get us rate limited by plc.directory
const PLC_PERMITS: usize = 16;
/// for each did:web host, and each pds a new account's probe goes to
const HOST_PERMITS: usize = 2;
/// ttl refreshes never take more than this many of the global permits, so they can't starve a
/// did that's waiting
const BACKGROUND_PERMITS: usize = 4;
/// jobs submitted and not answered yet. past this a submit is refused at once
pub(crate) const LOOKUP_QUEUE: usize = 4096;
const NEGATIVE_CACHE_SIZE: usize = 65_536;
/// messages every shard holds together while their dids' lookups run
const PARKED_BYTES: usize = 16 * 1024 * 1024;
/// dids held at once whose lookup goes to one did:web host, so a host with many dids that all
/// hang can't fill the shards on its own
const PARKED_PER_HOST: usize = 64;
/// hosts we keep breakers for. past it, the healthy ones are forgotten
const HOSTS_CAP: usize = 4096;
/// failures in a row before a host's breaker opens
const BREAKER_TRIP: u32 = 3;
const BREAKER_OPEN: Duration = Duration::from_secs(30);
const BREAKER_OPEN_MAX: Duration = Duration::from_secs(600);
const FAILED_FIRST: Duration = Duration::from_secs(60);
const FAILED_MAX: Duration = Duration::from_secs(3600);
/// for answers that say the did itself is broken, which won't change soon
const FAILED_FOR_GOOD: Duration = Duration::from_secs(600);
/// a plc mirror can be a few seconds behind on an account made just now, so a did it doesn't
/// know yet is asked again soon, then less and less often up to [`FAILED_FOR_GOOD`]
const NOT_FOUND_FIRST: Duration = Duration::from_secs(5);
/// how often an #identity may clear a did's backoff, so one spammed with them still costs at most
/// a lookup a minute. it's the doubling backoff's first step, the longest of the first steps
const IDENTITY_CLEAR_COOLDOWN: Duration = FAILED_FIRST;
/// how often a did's bad signatures may refresh its doc. a forged commit would cost a plc lookup
/// otherwise, and a real rotation only needs the first one
const SIG_REFRESH_COOLDOWN: Duration = Duration::from_secs(60);
/// plc's breaker stops every did:plc lookup at once, not just one host's dids, so it takes more
/// to open and never stays open for long
const PLC_BREAKER_TRIP: u32 = 10;
const PLC_BREAKER_OPEN: Duration = Duration::from_secs(30);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum LookupKind {
/// a did with no repo row yet: its doc, then its pds's word on the account
NewAccount,
/// a known did whose row has no signing key
Key,
/// a fresh doc for a known did
Refresh,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Priority {
/// a shard holds messages until this answers
Park,
/// nobody waits on it
Background,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
pub(crate) enum LookupError {
#[error("identity lookup failed: {0:?}")]
Resolve(ResolveFailure),
#[error("identity host is failing, not asked")]
HostDown,
#[error("identity lookup timed out")]
TimedOut,
/// the job ended without an answer, which only a panic does
#[error("identity lookup was lost")]
Lost,
/// the shard let go of the did's held messages before it answered
#[error("identity lookup wasn't waited for")]
Abandoned,
}
#[derive(Debug, Clone)]
pub(crate) struct Resolved {
pub(crate) doc: MiniDoc,
/// the pds's answer about a new account, `None` when it didn't give one
pub(crate) status: Option<RepoState<'static>>,
/// how long the doc fetch took, and the probe if there was one
pub(crate) doc_elapsed: Duration,
pub(crate) probe_elapsed: Option<Duration>,
}
#[derive(Debug)]
pub(crate) struct Answer {
pub(crate) did: Did,
pub(crate) kind: LookupKind,
pub(crate) result: Result<Resolved, LookupError>,
/// unix secs of the resolve, what the row stamps as `identity_resolved_at`
pub(crate) resolved_at: i64,
}
pub(crate) type AnswerTx = mpsc::UnboundedSender<Answer>;
pub(crate) struct Request {
pub(crate) did: Did,
pub(crate) kind: LookupKind,
pub(crate) priority: Priority,
pub(crate) resolver: Resolver,
/// the client for the status probe, which only indexer builds ask for
pub(crate) probe: Option<PublicHttpClient>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Submitted {
Started,
/// a job for the did is already running, and its answer covers this one
Joined,
/// the did failed recently and is in its backoff
Failing,
/// [`LOOKUP_QUEUE`] jobs are waiting already
Full,
}
struct InFlight {
/// a background job someone now waits on skips the background limit
urgent: AtomicBool,
woken: Notify,
}
#[derive(Default)]
struct Breaker {
failures: u32,
open_until: Option<Instant>,
backoff: Option<Duration>,
/// open time ran out and the request started then is checking the host again. a job can
/// be cut off before it reports, so a check older than a job's deadline doesn't count
probing: Option<Instant>,
}
/// whose breaker and permits a request goes through. plc is a key of its own, since a did:web
/// host can be named anything, `plc` included
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
enum HostKey {
Plc,
Named(SmolStr),
}
struct Host {
permits: Arc<Semaphore>,
size: usize,
breaker: Breaker,
}
impl Host {
fn new(key: &HostKey) -> Self {
let size = match key {
HostKey::Plc => PLC_PERMITS,
HostKey::Named(_) => HOST_PERMITS,
};
Self {
permits: Arc::new(Semaphore::new(size)),
size,
breaker: Breaker::default(),
}
}
fn busy(&self) -> bool {
self.permits.available_permits() < self.size
}
}
#[derive(Clone, Copy, Default)]
struct Negative {
failed_until: Option<Instant>,
backoff: Option<Duration>,
/// a sig failure refreshed the doc, and the next one before this gets no lookup
sig_refresh_until: Option<Instant>,
/// a background refresh failed, and the next one waits until this. a held message can
/// still ask, since the did itself isn't backed off
background_until: Option<Instant>,
/// an #identity cleared the backoff, and the next one before this doesn't
identity_clear_until: Option<Instant>,
}
/// why a shard let go of messages it held or was about to hold
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Dropped {
/// a park cap was hit, or the pool refused the lookup
Cap,
/// held past [`super::parked::HOLD_MAX`]
Deadline,
}
#[cfg(test)]
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub(crate) struct LookupStats {
pub(crate) started: u64,
pub(crate) joined: u64,
pub(crate) refused_failing: u64,
pub(crate) refused_full: u64,
pub(crate) failed: u64,
pub(crate) host_down: u64,
pub(crate) parked: u64,
pub(crate) parked_dropped_cap: u64,
pub(crate) parked_dropped_deadline: u64,
}
/// only the tests read these, so outside them every count compiles to nothing
#[derive(Default)]
struct Counters {
started: Counter,
joined: Counter,
refused_failing: Counter,
refused_full: Counter,
failed: Counter,
host_down: Counter,
parked: Counter,
parked_dropped_cap: Counter,
parked_dropped_deadline: Counter,
}
#[cfg(test)]
type Counter = std::sync::atomic::AtomicU64;
#[cfg(not(test))]
#[derive(Default)]
struct Counter;
#[cfg(not(test))]
impl Counter {
#[inline(always)]
fn fetch_add(&self, _: u64, _: Ordering) {}
}
#[cfg(test)]
type Stub = Arc<
dyn Fn(Did, LookupKind) -> futures::future::BoxFuture<'static, Result<Resolved, LookupError>>
+ Send
+ Sync,
>;
struct Inner {
in_flight: Mutex<HashMap<Did, Arc<InFlight>>>,
hosts: Mutex<HashMap<HostKey, Host>>,
parked_bytes: AtomicUsize,
parked_hosts: Mutex<HashMap<SmolStr, usize>>,
negative: HashCache<Did, Negative>,
global: Semaphore,
background: Semaphore,
counters: Counters,
deadline: Duration,
#[cfg(test)]
stub: Mutex<Option<Stub>>,
}
/// the pool on `AppState`. its tasks hold this and their request's clients, never the
/// `AppState`, so one that is still running can't keep the database open
#[derive(Clone)]
pub(crate) struct Lookups {
inner: Arc<Inner>,
}
impl Default for Lookups {
fn default() -> Self {
Self {
inner: Arc::new(Inner {
in_flight: Mutex::default(),
hosts: Mutex::default(),
parked_bytes: AtomicUsize::new(0),
parked_hosts: Mutex::default(),
negative: HashCache::with_capacity(1024, NEGATIVE_CACHE_SIZE),
global: Semaphore::new(GLOBAL_PERMITS),
background: Semaphore::new(BACKGROUND_PERMITS),
counters: Counters::default(),
deadline: JOB_DEADLINE,
#[cfg(test)]
stub: Mutex::default(),
}),
}
}
}
impl Lookups {
/// whether `did` failed a lookup recently, so a message that needs one is dropped without
/// asking again
pub(crate) fn failing(&self, did: &Did) -> bool {
self.inner
.negative
.read_sync(did, |_, negative| {
negative
.failed_until
.is_some_and(|until| until > Instant::now())
})
.unwrap_or(false)
}
fn background_cooling(&self, did: &Did) -> bool {
self.inner
.negative
.read_sync(did, |_, negative| {
negative
.background_until
.is_some_and(|until| until > Instant::now())
})
.unwrap_or(false)
}
#[cfg(test)]
pub(crate) fn in_flight(&self) -> usize {
self.inner.in_flight.lock().len()
}
/// starts a job for `request.did`, or attaches to the one already running. the answer goes
/// to `reply`, the channel of the shard that owns the did, exactly once per started job
pub(crate) fn submit(&self, request: Request, reply: &AnswerTx) -> Submitted {
let counters = &self.inner.counters;
let cooling =
request.priority == Priority::Background && self.background_cooling(&request.did);
if cooling || self.failing(&request.did) {
counters.refused_failing.fetch_add(1, Ordering::Relaxed);
return Submitted::Failing;
}
let flight = {
let mut in_flight = self.inner.in_flight.lock();
if let Some(flight) = in_flight.get(&request.did) {
if request.priority == Priority::Park {
flight.urgent.store(true, Ordering::Release);
flight.woken.notify_one();
}
counters.joined.fetch_add(1, Ordering::Relaxed);
return Submitted::Joined;
}
if in_flight.len() >= LOOKUP_QUEUE {
counters.refused_full.fetch_add(1, Ordering::Relaxed);
return Submitted::Full;
}
let flight = Arc::new(InFlight {
urgent: AtomicBool::new(request.priority == Priority::Park),
woken: Notify::new(),
});
in_flight.insert(request.did.clone(), flight.clone());
flight
};
counters.started.fetch_add(1, Ordering::Relaxed);
let reply = Reply {
tx: Some(reply.clone()),
did: request.did.clone(),
kind: request.kind,
flight: flight.clone(),
lookups: self.clone(),
};
let lookups = self.clone();
let started = Instant::now();
tokio::spawn(async move {
let job = lookups.run(request, flight, started);
let result = tokio::time::timeout(lookups.inner.deadline, job)
.await
.unwrap_or(Err(LookupError::TimedOut));
reply.send(result);
});
Submitted::Started
}
#[cfg(test)]
pub(crate) fn stats(&self) -> LookupStats {
let counters = &self.inner.counters;
LookupStats {
started: counters.started.load(Ordering::Relaxed),
joined: counters.joined.load(Ordering::Relaxed),
refused_failing: counters.refused_failing.load(Ordering::Relaxed),
refused_full: counters.refused_full.load(Ordering::Relaxed),
failed: counters.failed.load(Ordering::Relaxed),
host_down: counters.host_down.load(Ordering::Relaxed),
parked: counters.parked.load(Ordering::Relaxed),
parked_dropped_cap: counters.parked_dropped_cap.load(Ordering::Relaxed),
parked_dropped_deadline: counters.parked_dropped_deadline.load(Ordering::Relaxed),
}
}
/// room for a did's first held message, `bytes` big, whose lookup goes to `host`
pub(crate) fn reserve_park(&self, host: Option<&SmolStr>, bytes: usize) -> bool {
if let Some(host) = host {
let mut hosts = self.inner.parked_hosts.lock();
let parked = hosts.entry(host.clone()).or_default();
if *parked >= PARKED_PER_HOST {
return false;
}
*parked += 1;
}
if self.grow_park(bytes) {
self.inner.counters.parked.fetch_add(1, Ordering::Relaxed);
return true;
}
if let Some(host) = host {
self.unpark_host(host);
}
false
}
/// room for `bytes` more of held messages
pub(crate) fn grow_park(&self, bytes: usize) -> bool {
self.inner
.parked_bytes
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |parked| {
(parked + bytes <= PARKED_BYTES).then_some(parked + bytes)
})
.is_ok()
}
pub(crate) fn release_park(&self, host: Option<&SmolStr>, bytes: usize) {
self.inner.parked_bytes.fetch_sub(bytes, Ordering::AcqRel);
if let Some(host) = host {
self.unpark_host(host);
}
}
fn unpark_host(&self, host: &SmolStr) {
let mut hosts = self.inner.parked_hosts.lock();
if let Some(parked) = hosts.get_mut(host) {
*parked -= 1;
if *parked == 0 {
hosts.remove(host);
}
}
}
pub(crate) fn count_dropped(&self, why: Dropped) {
let counter = match why {
Dropped::Cap => &self.inner.counters.parked_dropped_cap,
Dropped::Deadline => &self.inner.counters.parked_dropped_deadline,
};
counter.fetch_add(1, Ordering::Relaxed);
}
/// whether a sig failure of `did` may refresh its doc now. saying yes starts the cooldown,
/// whatever the refresh then finds
pub(crate) fn start_sig_refresh(&self, did: &Did) -> bool {
let now = Instant::now();
let (_, mut entry) = self
.inner
.negative
.entry_sync(did.clone())
.or_put_with(Negative::default);
let negative = entry.get_mut();
if negative.sig_refresh_until.is_some_and(|until| until > now) {
return false;
}
negative.sig_refresh_until = Some(now + SIG_REFRESH_COOLDOWN);
true
}
/// an #identity says the did's doc may have changed, so its backoff goes and its next
/// message may ask again, like indigo purging its cache. at most once per
/// [`IDENTITY_CLEAR_COOLDOWN`], and the doubling step stays, so a did whose lookups keep
/// failing still backs off further. true when it cleared one
pub(crate) fn clear_for_identity(&self, did: &Did) -> bool {
let now = Instant::now();
let Some(mut entry) = self.inner.negative.get_sync(did) else {
return false;
};
let negative = entry.get_mut();
let backed_off = negative.failed_until.is_some_and(|until| until > now)
|| negative.background_until.is_some_and(|until| until > now);
let cooling = negative
.identity_clear_until
.is_some_and(|until| until > now);
if !backed_off || cooling {
return false;
}
negative.failed_until = None;
negative.background_until = None;
negative.identity_clear_until = Some(now + IDENTITY_CLEAR_COOLDOWN);
true
}
/// stops lookups for `did` for a while, like a failed one does
pub(crate) fn back_off(&self, did: &Did) {
self.fail(did, Backoff::Doubling);
}
async fn run(
&self,
request: Request,
flight: Arc<InFlight>,
started: Instant,
) -> Result<Resolved, LookupError> {
let _background = if flight.urgent.load(Ordering::Acquire) {
None
} else {
self.background_permit(&flight).await
};
#[cfg(test)]
let stub = self.inner.stub.lock().clone();
#[cfg(test)]
if let Some(stub) = stub {
let _global = self.inner.global.acquire().await.expect("never closed");
return stub(request.did, request.kind).await;
}
let Request {
did,
kind,
resolver,
probe,
..
} = request;
let host = identity_host(&did).map_or(HostKey::Plc, HostKey::Named);
let doc_elapsed;
let doc = {
// the host's permit first, so a job waiting on a host that hangs doesn't sit on a
// global permit every other host needs
let _host = self.host_permit(&host).await?;
let _global = self.inner.global.acquire().await.expect("never closed");
resolver.invalidate(&did).await;
let fetch_started = Instant::now();
let result = resolver.resolve_doc_fresh(&did).await;
doc_elapsed = fetch_started.elapsed();
self.record(&host, result.as_ref().err().map(|error| error.failure()));
result.map_err(|error| {
debug!(%did, err = %error, "identity lookup failed");
LookupError::Resolve(error.failure())
})?
};
resolver.cache_doc(did.clone(), doc.clone()).await;
if kind == LookupKind::Key && doc.key.is_none() {
return Err(LookupError::Resolve(ResolveFailure::NoSigningKey));
}
let probe_started = Instant::now();
let status = match (kind, &probe) {
(LookupKind::NewAccount, Some(http)) => {
// what's left of the deadline, so a slow probe never costs the doc it follows
let left = self
.inner
.deadline
.saturating_sub(started.elapsed())
.saturating_sub(PROBE_MARGIN);
self.probe(http, &did, &doc.pds, left.min(PROBE_TIMEOUT))
.await
}
_ => None,
};
let probe_elapsed =
(kind == LookupKind::NewAccount && probe.is_some()).then(|| probe_started.elapsed());
Ok(Resolved {
doc,
status,
doc_elapsed,
probe_elapsed,
})
}
/// the pds's word on a new account. any failure is no word, and the row starts out
/// backfilling
async fn probe(
&self,
http: &PublicHttpClient,
did: &Did,
pds: &Url,
timeout: Duration,
) -> Option<RepoState<'static>> {
let host = HostKey::Named(SmolStr::new(pds.host_str()?));
let probe = tokio::time::timeout(timeout, async {
let _host = self.host_permit(&host).await.ok()?;
let _global = self.inner.global.acquire().await.expect("never closed");
Some(super::probe_repo_status(http, did, pds).await)
})
.await;
match probe {
// the host's breaker is open
Ok(None) => None,
Ok(Some(StatusProbe::Answered(status))) => {
self.record(&host, None);
status
}
Ok(Some(StatusProbe::Unreachable)) | Err(_) => {
self.record(&host, Some(ResolveFailure::Unreachable));
None
}
}
}
/// a background job waits for one of the few background permits, unless a shard starts
/// waiting on it meanwhile
async fn background_permit(
&self,
flight: &InFlight,
) -> Option<tokio::sync::SemaphorePermit<'_>> {
loop {
if flight.urgent.load(Ordering::Acquire) {
return None;
}
tokio::select! {
permit = self.inner.background.acquire() => {
return Some(permit.expect("never closed"));
}
_ = flight.woken.notified() => {}
}
}
}
/// one of the host's permits, or [`LookupError::HostDown`] while its breaker is open
async fn host_permit(
&self,
host: &HostKey,
) -> Result<tokio::sync::OwnedSemaphorePermit, LookupError> {
let permits = {
let mut hosts = self.inner.hosts.lock();
if hosts.len() >= HOSTS_CAP && !hosts.contains_key(host) {
forget_hosts(&mut hosts);
}
let entry = hosts.entry(host.clone()).or_insert_with(|| Host::new(host));
let breaker = &mut entry.breaker;
let checking = breaker
.probing
.is_some_and(|since| since.elapsed() < self.inner.deadline);
match breaker.open_until {
Some(until) if until > Instant::now() || checking => return Err(self.host_down()),
Some(_) => breaker.probing = Some(Instant::now()),
None => {}
}
entry.permits.clone()
};
let permit = permits.acquire_owned().await.expect("never closed");
// the breaker may have opened while this waited for a permit
let opened = self.inner.hosts.lock().get(host).is_some_and(|entry| {
entry
.breaker
.open_until
.is_some_and(|until| until > Instant::now())
});
match opened {
true => Err(self.host_down()),
false => Ok(permit),
}
}
fn host_down(&self) -> LookupError {
self.inner
.counters
.host_down
.fetch_add(1, Ordering::Relaxed);
LookupError::HostDown
}
/// moves `host`'s breaker by what one request got. anything the host answered closes it,
/// even an answer that the did doesn't exist
fn record(&self, host: &HostKey, failure: Option<ResolveFailure>) {
let mut hosts = self.inner.hosts.lock();
let Some(Host { breaker, .. }) = hosts.get_mut(host) else {
return;
};
breaker.probing = None;
let plc = *host == HostKey::Plc;
match failure {
Some(ResolveFailure::Unreachable) => {
breaker.failures += 1;
let trip = if plc { PLC_BREAKER_TRIP } else { BREAKER_TRIP };
if breaker.failures >= trip {
let open = match (plc, breaker.backoff) {
(true, _) => PLC_BREAKER_OPEN,
(false, None) => BREAKER_OPEN,
(false, Some(last)) => (last * 2).min(BREAKER_OPEN_MAX),
};
breaker.backoff = Some(open);
breaker.open_until = Some(Instant::now() + open);
debug!(?host, ?open, "identity host failing, breaker open");
}
}
// we have no Retry-After here, the resolver drops the headers
Some(ResolveFailure::Ratelimited) => {
breaker.open_until = Some(Instant::now() + BREAKER_OPEN);
}
_ => *breaker = Breaker::default(),
}
}
/// `waited` is whether a shard held messages on the job, which only then backs the did off
/// when it failed. a background refresh that failed only holds off the next background
/// refresh, the row keeps its key and a held message may still ask
fn settle(&self, did: &Did, result: &Result<Resolved, LookupError>, waited: bool) {
self.inner.in_flight.lock().remove(did);
let backoff = match result {
Ok(_) => {
self.inner.negative.remove_if_sync(did, |negative| {
negative.failed_until = None;
negative.backoff = None;
negative.background_until = None;
let now = Instant::now();
negative.sig_refresh_until.is_none_or(|until| until <= now)
&& negative
.identity_clear_until
.is_none_or(|until| until <= now)
});
return;
}
// neither says anything about the did, and the breaker already stops the next ask
Err(LookupError::Resolve(ResolveFailure::Ratelimited) | LookupError::HostDown) => {
return;
}
Err(LookupError::Resolve(ResolveFailure::NotFound)) => Backoff::NotFound,
Err(LookupError::Resolve(
ResolveFailure::MissingPds
| ResolveFailure::NoSigningKey
| ResolveFailure::Tombstoned,
)) => Backoff::Fixed(FAILED_FOR_GOOD),
Err(_) => Backoff::Doubling,
};
self.inner.counters.failed.fetch_add(1, Ordering::Relaxed);
if waited {
self.fail(did, backoff);
return;
}
let (_, mut entry) = self
.inner
.negative
.entry_sync(did.clone())
.or_put_with(Negative::default);
entry.get_mut().background_until = Some(Instant::now() + FAILED_FIRST);
}
fn fail(&self, did: &Did, backoff: Backoff) {
let (_, mut entry) = self
.inner
.negative
.entry_sync(did.clone())
.or_put_with(Negative::default);
let negative = entry.get_mut();
let last = negative.backoff;
let wait = match backoff {
Backoff::Fixed(wait) => wait,
Backoff::Doubling => last.map_or(FAILED_FIRST, |last| (last * 2).min(FAILED_MAX)),
Backoff::NotFound => {
last.map_or(NOT_FOUND_FIRST, |last| (last * 2).min(FAILED_FOR_GOOD))
}
};
negative.backoff = Some(wait);
negative.failed_until = Some(Instant::now() + wait);
}
#[cfg(test)]
pub(crate) fn stub(&self, stub: Stub) {
*self.inner.stub.lock() = Some(stub);
}
/// a pool whose jobs give up after `deadline`
#[cfg(test)]
pub(crate) fn with_deadline(deadline: Duration) -> Self {
let lookups = Self::default();
let mut inner = Arc::into_inner(lookups.inner).expect("not shared yet");
inner.deadline = deadline;
Self {
inner: Arc::new(inner),
}
}
}
/// how long a failed lookup keeps a did from asking again
#[derive(Debug, Clone, Copy)]
enum Backoff {
Fixed(Duration),
/// the next step of the did's doubling backoff
Doubling,
/// like [`Backoff::Doubling`], from [`NOT_FOUND_FIRST`] up to [`FAILED_FOR_GOOD`]
NotFound,
}
/// at the cap, forgets the hosts nobody is asking right now whose breaker is shut, and if that
/// isn't enough, the open breakers too, so a stream of new hosts that each fail once can't grow
/// the map. plc is always kept
fn forget_hosts(hosts: &mut HashMap<HostKey, Host>) {
let now = Instant::now();
hosts.retain(|key, host| {
*key == HostKey::Plc
|| host.busy()
|| host.breaker.open_until.is_some_and(|until| until > now)
});
if hosts.len() >= HOSTS_CAP {
hosts.retain(|key, host| *key == HostKey::Plc || host.busy());
}
}
/// the host serving a did:web doc, `None` for every other method
pub(crate) fn identity_host(did: &Did) -> Option<SmolStr> {
crate::resolver::did_web_url(did)
.ok()?
.host_str()
.map(SmolStr::new)
}
/// answers its job exactly once, with [`LookupError::Lost`] if the job dies first
struct Reply {
tx: Option<AnswerTx>,
did: Did,
kind: LookupKind,
flight: Arc<InFlight>,
lookups: Lookups,
}
impl Reply {
fn send(mut self, result: Result<Resolved, LookupError>) {
self.answer(result);
}
fn answer(&mut self, result: Result<Resolved, LookupError>) {
let Some(tx) = self.tx.take() else {
return;
};
// out of the in-flight set before the shard hears of it, or a lookup the shard starts
// right away would join this finished one and never answer
let waited = self.flight.urgent.load(Ordering::Acquire);
self.lookups.settle(&self.did, &result, waited);
let _ = tx.send(Answer {
did: self.did.clone(),
kind: self.kind,
result,
resolved_at: chrono::Utc::now().timestamp(),
});
}
}
impl Drop for Reply {
fn drop(&mut self) {
self.answer(Err(LookupError::Lost));
}
}
#[cfg(test)]
mod tests {
use std::sync::atomic::AtomicUsize;
use super::*;
fn did(s: &str) -> Did {
Did::new_owned(s).unwrap()
}
fn resolved() -> Resolved {
Resolved {
doc: MiniDoc {
pds: Url::parse("https://pds.example/").unwrap(),
handle: None,
key: None,
},
status: None,
doc_elapsed: Duration::ZERO,
probe_elapsed: None,
}
}
fn resolver(client: reqwest::Client) -> Resolver {
Resolver::with_plc_urls_and_client(
vec![Url::parse("http://plc.example/").unwrap()],
16,
PublicHttpClient::test(client),
)
}
/// building a client loads the system roots, so the stubbed requests share one
fn request(did: Did, priority: Priority) -> Request {
static RESOLVER: std::sync::LazyLock<Resolver> =
std::sync::LazyLock::new(|| resolver(reqwest::Client::new()));
Request {
did,
kind: LookupKind::Refresh,
priority,
resolver: RESOLVER.clone(),
probe: None,
}
}
/// a pool whose jobs run `answer` instead of going to the network, counting the calls
fn stubbed(
answer: impl Fn() -> futures::future::BoxFuture<'static, Result<Resolved, LookupError>>
+ Send
+ Sync
+ 'static,
) -> (Lookups, Arc<AtomicUsize>) {
let lookups = Lookups::with_deadline(Duration::from_millis(200));
let calls = Arc::new(AtomicUsize::new(0));
let counted = calls.clone();
lookups.stub(Arc::new(move |_, _| {
counted.fetch_add(1, Ordering::SeqCst);
answer()
}));
(lookups, calls)
}
#[tokio::test]
async fn a_second_request_for_a_did_joins_the_running_one() {
let gate = Arc::new(tokio::sync::Semaphore::new(0));
let held = gate.clone();
let (lookups, calls) = stubbed(move || {
let held = held.clone();
Box::pin(async move {
let _ = held.acquire().await;
Ok(resolved())
})
});
let (tx, mut rx) = mpsc::unbounded_channel();
let alice = did("did:plc:aaaaaaaaaaaaaaaaaaaaaaaa");
assert_eq!(
lookups.submit(request(alice.clone(), Priority::Background), &tx),
Submitted::Started
);
assert_eq!(
lookups.submit(request(alice.clone(), Priority::Park), &tx),
Submitted::Joined
);
gate.add_permits(1);
let answer = rx.recv().await.unwrap();
assert_eq!(answer.did, alice);
assert!(answer.result.is_ok());
assert_eq!(calls.load(Ordering::SeqCst), 1);
assert_eq!(lookups.in_flight(), 0);
// a finished job is out of the way, so the next request starts its own
assert_eq!(
lookups.submit(request(alice, Priority::Park), &tx),
Submitted::Started
);
}
#[tokio::test]
async fn a_job_that_never_finishes_answers_by_its_deadline() {
let (lookups, _) = stubbed(|| Box::pin(std::future::pending()));
let (tx, mut rx) = mpsc::unbounded_channel();
let alice = did("did:plc:aaaaaaaaaaaaaaaaaaaaaaaa");
let started = Instant::now();
lookups.submit(request(alice.clone(), Priority::Park), &tx);
let answer = rx.recv().await.unwrap();
assert!(matches!(answer.result, Err(LookupError::TimedOut)));
assert!(started.elapsed() < Duration::from_secs(2));
assert!(lookups.failing(&alice), "a timeout backs the did off");
assert_eq!(
lookups.submit(request(alice, Priority::Park), &tx),
Submitted::Failing
);
}
#[tokio::test]
async fn a_job_that_panics_still_answers() {
let (lookups, _) = stubbed(|| Box::pin(async { panic!("stinkpot") }));
let (tx, mut rx) = mpsc::unbounded_channel();
lookups.submit(
request(did("did:plc:aaaaaaaaaaaaaaaaaaaaaaaa"), Priority::Park),
&tx,
);
let answer = tokio::time::timeout(Duration::from_secs(5), rx.recv())
.await
.unwrap()
.unwrap();
assert!(matches!(answer.result, Err(LookupError::Lost)));
assert_eq!(lookups.in_flight(), 0);
}
fn failed_until(lookups: &Lookups, did: &Did) -> Option<Instant> {
lookups
.inner
.negative
.read_sync(did, |_, negative| negative.failed_until)
.flatten()
}
/// a mirror that hasn't seen an account made just now says it isn't there, so the first
/// answer like that only holds the did off briefly, and a did that stays missing backs off
/// further each time
#[tokio::test]
async fn a_did_that_isnt_there_yet_is_asked_again_soon() {
let (lookups, _) =
stubbed(|| Box::pin(async { Err(LookupError::Resolve(ResolveFailure::NotFound)) }));
let (tx, mut rx) = mpsc::unbounded_channel();
let alice = did("did:plc:aaaaaaaaaaaaaaaaaaaaaaaa");
lookups.submit(request(alice.clone(), Priority::Park), &tx);
rx.recv().await.unwrap();
let until = failed_until(&lookups, &alice).expect("backed off");
assert!(until <= Instant::now() + NOT_FOUND_FIRST);
lookups.fail(&alice, Backoff::NotFound);
let until = failed_until(&lookups, &alice).unwrap();
assert!(until > Instant::now() + NOT_FOUND_FIRST);
}
/// nobody waits on a background refresh, and the row keeps the key it has, so its failure
/// mustn't stop the did's next sig failure or pds move from asking
#[tokio::test]
async fn a_failed_background_refresh_doesnt_back_the_did_off() {
let (lookups, _) =
stubbed(|| Box::pin(async { Err(LookupError::Resolve(ResolveFailure::Unreachable)) }));
let (tx, mut rx) = mpsc::unbounded_channel();
let alice = did("did:plc:aaaaaaaaaaaaaaaaaaaaaaaa");
lookups.submit(request(alice.clone(), Priority::Background), &tx);
assert!(rx.recv().await.unwrap().result.is_err());
assert!(!lookups.failing(&alice));
// but the next background refresh waits, so a did whose refreshes keep failing can't
// ask on every message
assert_eq!(
lookups.submit(request(alice.clone(), Priority::Background), &tx),
Submitted::Failing
);
lookups.submit(request(alice.clone(), Priority::Park), &tx);
rx.recv().await.unwrap();
assert!(
lookups.failing(&alice),
"a held message's lookup still backs off"
);
}
/// an #identity lets a backed off did ask again, once per cooldown, and keeps its doubling
#[test]
fn an_identity_clears_a_backoff_once_per_cooldown() {
let lookups = Lookups::default();
let alice = did("did:plc:aaaaaaaaaaaaaaaaaaaaaaaa");
assert!(!lookups.clear_for_identity(&alice), "nothing to clear");
lookups.fail(&alice, Backoff::Doubling);
assert!(lookups.failing(&alice));
assert!(lookups.clear_for_identity(&alice));
assert!(!lookups.failing(&alice));
// its next lookup fails again, and a second #identity inside the cooldown leaves that be
lookups.fail(&alice, Backoff::Doubling);
assert!(!lookups.clear_for_identity(&alice));
assert!(lookups.failing(&alice));
let until = failed_until(&lookups, &alice).unwrap();
assert!(
until > Instant::now() + FAILED_FIRST,
"the doubling was reset"
);
// nor does it touch the sig refresh cooldown
assert!(lookups.start_sig_refresh(&alice));
lookups
.inner
.negative
.get_sync(&alice)
.unwrap()
.get_mut()
.identity_clear_until = None;
assert!(lookups.clear_for_identity(&alice));
assert!(!lookups.start_sig_refresh(&alice));
}
/// an open breaker is about the host, so the did is asked again once it shuts
#[tokio::test]
async fn a_host_that_is_down_doesnt_back_the_did_off() {
let (lookups, _) = stubbed(|| Box::pin(async { Err(LookupError::HostDown) }));
let (tx, mut rx) = mpsc::unbounded_channel();
let alice = did("did:plc:aaaaaaaaaaaaaaaaaaaaaaaa");
lookups.submit(request(alice.clone(), Priority::Park), &tx);
rx.recv().await.unwrap();
assert!(!lookups.failing(&alice));
}
#[tokio::test]
async fn the_queue_refuses_past_its_bound() {
let (lookups, _) = stubbed(|| Box::pin(std::future::pending()));
let (tx, _rx) = mpsc::unbounded_channel();
for i in 0..LOOKUP_QUEUE {
assert_eq!(
lookups.submit(
request(did(&format!("did:web:u{i}.example")), Priority::Park),
&tx
),
Submitted::Started
);
}
assert_eq!(
lookups.submit(
request(did("did:web:one.more.example"), Priority::Park),
&tx
),
Submitted::Full
);
assert_eq!(lookups.stats().refused_full, 1);
}
/// accepts connections and never answers them
async fn silent_host() -> (std::net::SocketAddr, Arc<AtomicUsize>) {
let accepted = Arc::new(AtomicUsize::new(0));
let counter = accepted.clone();
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
let mut held = Vec::new();
while let Ok((stream, _)) = listener.accept().await {
counter.fetch_add(1, Ordering::SeqCst);
held.push(stream);
}
});
(addr, accepted)
}
#[tokio::test]
async fn a_hanging_did_web_host_is_asked_three_times_then_skipped() {
let (addr, accepted) = silent_host().await;
let client = reqwest::Client::builder()
.no_proxy()
.resolve("slow.test", addr)
.timeout(Duration::from_millis(100))
.build()
.unwrap();
let lookups = Lookups::default();
let (tx, mut rx) = mpsc::unbounded_channel();
for i in 0..100 {
let request = Request {
did: did(&format!("did:web:slow.test:u{i}")),
kind: LookupKind::NewAccount,
priority: Priority::Park,
resolver: resolver(client.clone()),
probe: None,
};
assert_eq!(lookups.submit(request, &tx), Submitted::Started);
let answer = rx.recv().await.unwrap();
let expected = match i < BREAKER_TRIP as usize {
true => LookupError::Resolve(ResolveFailure::Unreachable),
false => LookupError::HostDown,
};
assert_eq!(answer.result.unwrap_err(), expected, "lookup {i}");
}
assert_eq!(accepted.load(Ordering::SeqCst), BREAKER_TRIP as usize);
assert_eq!(lookups.stats().host_down, 100 - BREAKER_TRIP as u64);
}
/// how long `host`'s breaker is open from now, after `failures` more failures in a row
fn open_after(lookups: &Lookups, host: &HostKey, failures: u32) -> Option<Duration> {
lookups
.inner
.hosts
.lock()
.entry(host.clone())
.or_insert_with(|| Host::new(host));
for _ in 0..failures {
lookups.record(host, Some(ResolveFailure::Unreachable));
}
let hosts = lookups.inner.hosts.lock();
let until = hosts[host].breaker.open_until?;
Some(until.saturating_duration_since(Instant::now()))
}
#[test]
fn plc_takes_more_to_open_and_never_backs_off_long() {
let lookups = Lookups::default();
let plc = &HostKey::Plc;
assert_eq!(open_after(&lookups, plc, PLC_BREAKER_TRIP - 1), None);
let first = open_after(&lookups, plc, 1).unwrap();
assert!(first <= PLC_BREAKER_OPEN);
// a failed check after the open time doesn't double it
let again = open_after(&lookups, plc, 1).unwrap();
assert!(again <= PLC_BREAKER_OPEN, "reopened for {again:?}");
let host = &HostKey::Named("slow.test".into());
assert!(open_after(&lookups, host, BREAKER_TRIP).unwrap() <= BREAKER_OPEN);
assert!(open_after(&lookups, host, 1).unwrap() > BREAKER_OPEN);
}
/// a plc that answers every did with a doc once `gate` lets it, counting the requests that
/// arrived
async fn plc_stub(
gate: Arc<tokio::sync::Semaphore>,
) -> (std::net::SocketAddr, Arc<AtomicUsize>) {
let arrived = Arc::new(AtomicUsize::new(0));
let counter = arrived.clone();
let app = axum::Router::new().fallback(move |uri: axum::http::Uri| {
let gate = gate.clone();
let counter = counter.clone();
async move {
counter.fetch_add(1, Ordering::SeqCst);
drop(gate.acquire().await);
let doc = serde_json::json!({
"id": uri.path().trim_start_matches('/'),
"service": [{
"id": "#atproto_pds",
"type": "AtprotoPersonalDataServer",
"serviceEndpoint": "https://pds.example",
}],
});
(
[(axum::http::header::CONTENT_TYPE, "application/json")],
doc.to_string(),
)
}
});
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
(addr, arrived)
}
/// hangs up on every connection
async fn rude_host() -> std::net::SocketAddr {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move { while listener.accept().await.is_ok() {} });
addr
}
fn plc_request(did: Did, resolver: &Resolver) -> Request {
Request {
did,
kind: LookupKind::Refresh,
priority: Priority::Park,
resolver: resolver.clone(),
probe: None,
}
}
async fn within<T>(secs: u64, what: &str, future: impl std::future::Future<Output = T>) -> T {
tokio::time::timeout(Duration::from_secs(secs), future)
.await
.unwrap_or_else(|_| panic!("{what}"))
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn plc_takes_its_own_number_of_lookups_at_once() {
let gate = Arc::new(tokio::sync::Semaphore::new(0));
let (plc, arrived) = plc_stub(gate.clone()).await;
let client = reqwest::Client::builder()
.no_proxy()
.resolve("plc.example", plc)
.build()
.unwrap();
let resolver = resolver(client);
let lookups = Lookups::default();
let (tx, mut rx) = mpsc::unbounded_channel();
for i in 0..PLC_PERMITS + 4 {
let name: String = i
.to_string()
.chars()
.map(|c| (c as u8 + 49) as char)
.collect();
let did = did(&format!("did:plc:{name:a>24}"));
lookups.submit(plc_request(did, &resolver), &tx);
}
within(5, "plc got fewer lookups at once than its permits", async {
while arrived.load(Ordering::SeqCst) < PLC_PERMITS {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await;
tokio::time::sleep(Duration::from_millis(100)).await;
assert_eq!(arrived.load(Ordering::SeqCst), PLC_PERMITS);
gate.add_permits(PLC_PERMITS + 4);
for _ in 0..PLC_PERMITS + 4 {
assert!(rx.recv().await.unwrap().result.is_ok());
}
}
/// the jobs waiting on a host that hangs hold none of the global permits, so a did on
/// another host gets its answer at once
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_hanging_host_leaves_the_global_permits_to_everyone_else() {
let (silent, _) = silent_host().await;
let (plc, _) = plc_stub(Arc::new(tokio::sync::Semaphore::new(1024))).await;
let client = reqwest::Client::builder()
.no_proxy()
.resolve("plc.example", plc)
.resolve("slow.test", silent)
.build()
.unwrap();
let resolver = resolver(client);
let lookups = Lookups::with_deadline(Duration::from_secs(20));
let (tx, mut rx) = mpsc::unbounded_channel();
for i in 0..PARKED_PER_HOST {
let did = did(&format!("did:web:slow.test:u{i}"));
lookups.submit(plc_request(did, &resolver), &tx);
}
tokio::time::sleep(Duration::from_millis(100)).await;
let alice = did("did:plc:aaaaaaaaaaaaaaaaaaaaaaaa");
lookups.submit(plc_request(alice.clone(), &resolver), &tx);
let answer = within(
5,
"the plc lookup waited behind the hanging host",
rx.recv(),
)
.await
.unwrap();
assert_eq!(answer.did, alice);
assert!(answer.result.is_ok());
}
/// a did:web host named `plc` has a breaker of its own, so its failures can't stop every
/// did:plc lookup
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_did_web_host_named_plc_cant_open_the_plc_breaker() {
let (plc, _) = plc_stub(Arc::new(tokio::sync::Semaphore::new(1024))).await;
let client = reqwest::Client::builder()
.no_proxy()
.resolve("plc.example", plc)
.resolve("plc", rude_host().await)
.build()
.unwrap();
let resolver = resolver(client);
let lookups = Lookups::default();
let (tx, mut rx) = mpsc::unbounded_channel();
for i in 0..PLC_BREAKER_TRIP {
let did = did(&format!("did:web:plc:u{i}"));
lookups.submit(plc_request(did, &resolver), &tx);
assert!(rx.recv().await.unwrap().result.is_err());
}
let alice = did("did:plc:aaaaaaaaaaaaaaaaaaaaaaaa");
lookups.submit(plc_request(alice, &resolver), &tx);
let answer = rx.recv().await.unwrap();
assert!(answer.result.is_ok(), "{:?}", answer.result.err());
}
/// hosts that failed once and were never asked again are forgotten at the cap
#[tokio::test]
async fn the_host_breakers_stay_under_their_cap() {
let lookups = Lookups::default();
for i in 0..HOSTS_CAP + 1000 {
let host = HostKey::Named(format!("h{i}.test").into());
drop(lookups.host_permit(&host).await.unwrap());
lookups.record(&host, Some(ResolveFailure::Unreachable));
}
assert!(lookups.inner.hosts.lock().len() <= HOSTS_CAP);
}
#[test]
fn did_web_dids_name_their_host() {
assert_eq!(
identity_host(&did("did:web:example.com:user:alice")).as_deref(),
Some("example.com")
);
assert_eq!(
identity_host(&did("did:plc:aaaaaaaaaaaaaaaaaaaaaaaa")),
None
);
}
}