diff --git a/Cargo.lock b/Cargo.lock --- a/Cargo.lock +++ b/Cargo.lock @@ -4555,6 +4555,29 @@ "tracing", ] [[package]] +name = "knot-keyfill" +version = "2.0.0" +dependencies = [ + "bytes", + "futures", + "http", + "knot-atproto", + "knot-cob", + "knot-cobs", + "knot-git", + "knot-index", + "knot-resource", + "knot-runtime", + "knot-types", + "serde_json", + "tempfile", + "tokio", + "tokio-util", + "tracing", + "url", +] + +[[package]] name = "knot-langs" version = "2.0.0" dependencies = [ diff --git a/Cargo.toml b/Cargo.toml --- a/Cargo.toml +++ b/Cargo.toml @@ -57,6 +57,7 @@ knot-pack = { path = "knot2/crates/knot-pack" } knot-cob = { path = "knot2/crates/knot-cob" } knot-cobs = { path = "knot2/crates/knot-cobs" } knot-index = { path = "knot2/crates/knot-index" } +knot-keyfill = { path = "knot2/crates/knot-keyfill" } knot-cache = { path = "knot2/crates/knot-cache" } knot-acl = { path = "knot2/crates/knot-acl" } knot-atproto = { path = "knot2/crates/knot-atproto" } diff --git a/knot2/crates/knot-keyfill/Cargo.toml b/knot2/crates/knot-keyfill/Cargo.toml new file mode 100644 --- /dev/null +++ b/knot2/crates/knot-keyfill/Cargo.toml @@ -0,0 +1,27 @@ +[package] +name = "knot-keyfill" +version = "2.0.0" +edition.workspace = true +rust-version.workspace = true +license.workspace = true + +[dependencies] +knot-atproto = { workspace = true } +knot-index = { workspace = true } +knot-resource = { workspace = true } +knot-runtime = { workspace = true } +knot-types = { workspace = true } +futures = { workspace = true } +tokio = { workspace = true } +tokio-util = { workspace = true } +tracing = { workspace = true } +url = { workspace = true } + +[dev-dependencies] +knot-cob = { workspace = true } +knot-cobs = { workspace = true } +knot-git = { workspace = true } +bytes = { workspace = true } +http = { workspace = true } +serde_json = { workspace = true } +tempfile = { workspace = true } diff --git a/knot2/crates/knot-keyfill/src/lib.rs b/knot2/crates/knot-keyfill/src/lib.rs new file mode 100644 --- /dev/null +++ b/knot2/crates/knot-keyfill/src/lib.rs @@ -0,0 +1,619 @@ +use std::future::Future; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +use futures::future::OptionFuture; +use futures::{FutureExt, StreamExt}; +use knot_atproto::{Atproto, AtprotoError}; +use knot_index::{ + Coverage, HostedCoverage, Index, IndexGeneration, KeyLease, KeyRecord, KeyReprieve, + KeyReprieved, KeyTtl, MemberWork, Pushers, Resolved, StalePushers, SuspectPushers, SweepFloor, +}; +use knot_resource::{Burst, HostKey, HostPacer, RateLimit, RefillMicros, SlotPermit, Slots}; +use knot_runtime::{Clock, HttpTransport}; +use knot_types::{AccountDid, OfferedKey, UnixMicros, UnixSeconds}; +use tokio::sync::watch; +use tokio_util::sync::CancellationToken; +use url::Url; + +const FILL_FANOUT: usize = 8; + +const PASS_HEADROOM: u32 = 2; + +macro_rules! span { + ($($name:ident from $unit:ident),+ $(,)?) => {$( + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + pub struct $name(Duration); + + impl $name { + pub const fn $unit(value: u64) -> Self { + Self(Duration::$unit(value)) + } + + pub const fn get(self) -> Duration { + self.0 + } + } + )+}; +} + +span!( + BusyRetry from from_millis, + SettleFloor from from_millis, + StalledBackoff from from_secs, + SettledPause from from_secs, +); + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AccountBudget(usize); + +impl AccountBudget { + pub const fn new(accounts: usize) -> Self { + Self(if accounts == 0 { 1 } else { accounts }) + } + + pub const fn get(self) -> usize { + self.0 + } +} + +#[derive(Debug, Default)] +pub struct Cursor(AtomicUsize); + +impl Cursor { + fn advance(&self, by: usize, len: usize) -> usize { + match len { + 0 => 0, + len => self + .0 + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |seen| { + Some((seen % len).wrapping_add(by) % len) + }) + .map_or(0, |seen| seen % len), + } + } +} + +#[derive(Debug, Default)] +pub struct Cursors { + members: Cursor, + suspected: Cursor, +} + +fn portion(accounts: &[AccountDid], budget: AccountBudget, cursor: &Cursor) -> Vec { + let taken = budget.get().min(accounts.len()); + let start = cursor.advance(taken, accounts.len()); + accounts + .iter() + .cycle() + .skip(start) + .take(taken) + .cloned() + .collect() +} + +#[derive(Debug, Clone, Copy)] +pub struct Pace { + pub busy: BusyRetry, + pub floor: SettleFloor, + pub ttl: KeyTtl, + pub reprieve: KeyReprieve, + pub sweep: SweepFloor, + pub stalled: StalledBackoff, + pub settled: SettledPause, + pub members: AccountBudget, + pub suspected: AccountBudget, + pub host: RateLimit, +} + +impl Default for Pace { + fn default() -> Self { + Self { + busy: BusyRetry::from_millis(50), + floor: SettleFloor::from_millis(1_000), + ttl: KeyTtl::DEFAULT, + reprieve: KeyReprieve::DEFAULT, + sweep: SweepFloor::DEFAULT, + stalled: StalledBackoff::from_secs(30), + settled: SettledPause::from_secs(60), + members: AccountBudget::new(64), + suspected: AccountBudget::new(256), + host: RateLimit { + burst: Burst::new(10), + refill: RefillMicros::new(200_000), + }, + } + } +} + +impl Pace { + fn ttl_covering(self, accounts: usize) -> KeyTtl { + let micros = (accounts as u64).saturating_mul(self.host.interval().get()); + let secs = (micros / 1_000_000).saturating_mul(u64::from(PASS_HEADROOM)); + self.ttl.longest(KeyTtl::from_secs(secs)) + } +} + +pub fn spawn( + index: Arc, + atproto: Arc>, + slots: Slots, + pace: Pace, + shutdown: CancellationToken, +) -> tokio::task::JoinHandle<()> { + let stop = shutdown.clone(); + let driver = Driver { + generations: index.generations(), + pacer: HostPacer::new(pace.host), + cursors: Cursors::default(), + index, + atproto, + slots, + pace, + shutdown, + }; + tokio::spawn(async move { + futures::stream::unfold(driver, |mut driver| async move { + driver.pass().await; + Some(((), driver)) + }) + .take_until(stop.cancelled_owned()) + .for_each(|()| std::future::ready(())) + .await; + tracing::info!("key fill stopped"); + }) +} + +struct Driver { + index: Arc, + atproto: Arc>, + slots: Slots, + pacer: HostPacer, + cursors: Cursors, + pace: Pace, + shutdown: CancellationToken, + generations: watch::Receiver, +} + +impl Driver { + async fn pass(&mut self) { + self.generations.mark_unchanged(); + let filling = fill_once( + &self.index, + &self.atproto, + &self.slots, + &self.pacer, + self.pace, + &self.cursors, + ); + let pause = tokio::select! { + pause = guarded(filling, self.pace) => pause, + () = self.shutdown.cancelled() => self.pace.stalled.get(), + }; + tokio::select! { + () = self.shutdown.cancelled() => {} + () = settle(&mut self.generations, pause, self.pace.floor) => {} + } + } +} + +async fn guarded(pass: impl Future, pace: Pace) -> Duration { + std::panic::AssertUnwindSafe(pass) + .catch_unwind() + .await + .unwrap_or_else(|_| { + tracing::error!("key fill pass panicked, backing off before the next pass"); + pace.stalled.get() + }) +} + +async fn settle( + generations: &mut watch::Receiver, + pause: Duration, + floor: SettleFloor, +) { + let held = floor.get().min(pause); + tokio::time::sleep(held).await; + let _ = tokio::time::timeout(pause.saturating_sub(held), generations.changed()).await; +} + +struct Pass<'a, H, C> { + index: &'a Arc, + atproto: &'a Arc>, + slots: &'a Slots, + pacer: &'a HostPacer, + pace: Pace, + now: UnixSeconds, + lease: KeyLease, + reprieve: KeyReprieve, +} + +pub async fn fill_once( + index: &Arc, + atproto: &Arc>, + slots: &Slots, + pacer: &HostPacer, + pace: Pace, + cursors: &Cursors, +) -> Duration { + fold_hosted(index).await; + let now = atproto.now().seconds(); + let Resolved::Ready(work) = index.keys().work(now, pace.sweep) else { + index.keys().mark_warming(); + return pace.stalled.get(); + }; + if let HostedCoverage::Partial { unread } = work.hosted { + tracing::debug!( + repos = unread, + "partial grant set, a repo was registered while this pass was working out who may push" + ); + } + let members = match (work.hosted, work.members) { + (HostedCoverage::Whole, Resolved::Ready(members)) => { + index.keys().retain(&members.kept); + Some(members) + } + (HostedCoverage::Partial { .. }, Resolved::Ready(members)) => Some(members), + (_, Resolved::Warming) => None, + }; + let ttl = pace.ttl_covering(work.tracked); + if ttl != pace.ttl { + tracing::debug!( + tracked = work.tracked, + ttl_secs = ttl.get().as_secs(), + "one paced pass over the grant set outruns the key ttl, so the fill stretches it" + ); + } + let pass = Pass { + index, + atproto, + slots, + pacer, + pace, + now, + lease: ttl.lease_from(now), + reprieve: pace.reprieve.budgeted_for(ttl), + }; + let pause = fill_pushers( + &pass, + PushWork { + generation: work.generation, + hosted: work.hosted, + pushers: work.pushers, + due: work.due, + suspected: work.suspected, + complete: work.complete, + }, + &cursors.suspected, + ) + .await; + OptionFuture::from(members.map(|work| fill_members(&pass, work, &cursors.members))).await; + pause +} + +struct PushWork { + generation: IndexGeneration, + hosted: HostedCoverage, + pushers: Pushers, + due: StalePushers, + suspected: SuspectPushers, + complete: bool, +} + +async fn fold_hosted(index: &Arc) { + let index = Arc::clone(index); + match tokio::task::spawn_blocking(move || index.warm_collaborators()).await { + Ok(0) | Err(_) => {} + Ok(unreadable) => tracing::warn!( + repos = unreadable, + "the knot hosts registered repos it can't open, so it won't grant anybody through \ + them until an operator restores or deregisters each repo" + ), + } +} + +async fn fill_pushers( + pass: &Pass<'_, H, C>, + work: PushWork, + cursor: &Cursor, +) -> Duration { + let PushWork { + generation, + hosted, + pushers, + due, + suspected, + complete, + } = work; + if !complete { + pass.index.keys().mark_warming(); + } + let wanted = due.len(); + let recorded = match wanted { + 0 => 0, + _ => { + let recorded = record_each(pass, due.into_vec()).await; + tracing::debug!(wanted, recorded, "pusher key fill pass"); + recorded + } + }; + recheck_suspected(pass, &suspected, cursor).await; + match recorded < wanted { + true => pass.pace.stalled.get(), + false => claim_ready(pass, generation, hosted, &pushers), + } +} + +async fn recheck_suspected( + pass: &Pass<'_, H, C>, + suspected: &SuspectPushers, + cursor: &Cursor, +) { + if suspected.is_empty() { + return; + } + let batch = portion(suspected.as_slice(), pass.pace.suspected, cursor); + let wanted = batch.len(); + let recorded = record_each(pass, batch).await; + tracing::debug!( + wanted, + recorded, + deferred = suspected.len().saturating_sub(wanted), + "pusher key recheck pass, a client offered a key the accounts on file don't publish" + ); +} + +fn claim_ready( + pass: &Pass<'_, H, C>, + generation: IndexGeneration, + hosted: HostedCoverage, + pushers: &Pushers, +) -> Duration { + let keys = pass.index.keys(); + let settled = + hosted == HostedCoverage::Whole && keys.all_live(pushers, pass.atproto.now().seconds()); + if !settled { + return pass.pace.stalled.get(); + } + keys.mark_ready(generation); + match keys.coverage() { + Coverage::Ready => pass.pace.settled.get(), + Coverage::Warming => pass.pace.stalled.get(), + } +} + +async fn fill_members( + pass: &Pass<'_, H, C>, + work: MemberWork, + cursor: &Cursor, +) { + if !work.unread.is_empty() { + let wanted = work.unread.len(); + let recorded = record_each(pass, work.unread.into_vec()).await; + tracing::debug!(wanted, recorded, "member key first-read pass"); + } + if work.due.is_empty() { + return; + } + let batch = portion(work.due.as_slice(), pass.pace.members, cursor); + let wanted = batch.len(); + let recorded = record_each(pass, batch).await; + tracing::debug!( + wanted, + recorded, + deferred = work.due.len().saturating_sub(wanted), + "member key renewal pass" + ); +} + +async fn record_each( + pass: &Pass<'_, H, C>, + stale: Vec, +) -> usize { + futures::stream::iter(stale) + .map(|did| async move { + match published_keys(pass, &did).await { + Ok(keys) => usize::from(record(pass, &did, keys)), + Err(error) if error.is_gone() => { + tracing::debug!( + did = did.as_str(), + %error, + "key fill records an empty key set for an account whose DID document is gone" + ); + usize::from(record(pass, &did, Vec::new())) + } + Err(error) => reprieve(pass, &did, error), + } + }) + .buffer_unordered(FILL_FANOUT) + .fold(0, |total, recorded| async move { total + recorded }) + .await +} + +fn reprieve(pass: &Pass<'_, H, C>, did: &AccountDid, error: AtprotoError) -> usize { + let outcome = pass + .index + .keys() + .reprieve(did, pass.now, pass.reprieve, pass.lease); + match outcome { + KeyReprieved::Exhausted => tracing::warn!( + did = did.as_str(), + %error, + "the knot spent the whole reprieve failing to read an account, so it records an \ + empty key set for the account until a later read succeeds" + ), + KeyReprieved::Extended | KeyReprieved::Pending => tracing::debug!( + did = did.as_str(), + %error, + ?outcome, + "key fill couldn't read an account's records" + ), + } + match outcome { + KeyReprieved::Extended | KeyReprieved::Exhausted => 1, + KeyReprieved::Pending => 0, + } +} + +fn record(pass: &Pass<'_, H, C>, did: &AccountDid, keys: Vec) -> bool { + match pass.index.keys().record(did, keys, pass.lease) { + KeyRecord::Stored => true, + KeyRecord::Unheld => { + tracing::warn!( + did = did.as_str(), + "the key set is full, so the knot will check this account's pushes against its \ + PDS every time instead of against the set" + ); + true + } + KeyRecord::Saturated => { + tracing::warn!( + did = did.as_str(), + "the key budget can't record even that it read this account, so the knot keeps \ + reporting the set incomplete and defers every offered key to the push check" + ); + false + } + } +} + +async fn published_keys( + pass: &Pass<'_, H, C>, + did: &AccountDid, +) -> Result, AtprotoError> { + let atproto = pass.atproto; + if let Ok(document) = atproto.document_url(did) { + wait_for_turn(pass.pacer, &document, atproto.now()).await; + } + let identity = { + let _permit = idle_permit(pass.slots, pass.pace).await; + atproto.resolve_identity(did).await? + }; + wait_for_turn(pass.pacer, identity.pds.url(), atproto.now()).await; + let _permit = idle_permit(pass.slots, pass.pace).await; + atproto.pubkeys_at(&identity, did).await +} + +async fn wait_for_turn(pacer: &HostPacer, url: &Url, now: UnixMicros) { + if let Some(host) = url.host_str().map(HostKey::new) { + tokio::time::sleep(pacer.reserve(&host, now)).await; + } +} + +async fn idle_permit(slots: &Slots, pace: Pace) -> SlotPermit { + let attempts = futures::stream::repeat(()).filter_map(|()| async { + match slots.resolve.try_acquire() { + Some(permit) => Some(permit), + None => { + tokio::time::sleep(pace.busy.get()).await; + None + } + } + }); + futures::pin_mut!(attempts); + attempts + .next() + .await + .expect("an endless stream of attempts yields a permit") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_grant_set_the_pacer_cant_reread_within_the_ttl_stretches_it() { + let pace = Pace::default(); + assert_eq!( + pace.ttl_covering(1_000), + KeyTtl::DEFAULT, + "a set one paced pass covers well inside the ttl keeps the ttl it was configured with" + ); + assert_eq!( + pace.ttl_covering(100_000), + KeyTtl::from_secs(40_000), + "at one directory turn per 200ms a hundred thousand accounts take 20_000s to reread, \ + so the entries the pass wrote first must outlive the pass that writes the last" + ); + } + + #[test] + fn the_reprieve_budget_never_undercuts_the_ttl_the_fill_is_working_to() { + let stretched = KeyTtl::from_secs(40_000); + assert_eq!( + KeyReprieve::DEFAULT.budgeted_for(stretched), + KeyReprieve::from_secs(300, 40_000), + "an account would be released while the pass that would reread it is still running, \ + if the budget stayed under the ttl" + ); + assert_eq!( + KeyReprieve::DEFAULT.budgeted_for(KeyTtl::DEFAULT), + KeyReprieve::DEFAULT, + "a ttl the budget already covers leaves the budget alone" + ); + } + + #[test] + fn the_member_cursor_turns_over_without_running_past_its_type() { + let cursor = Cursor(AtomicUsize::new(usize::MAX)); + assert_eq!( + cursor.advance(3, 4), + usize::MAX % 4, + "a cursor at the end of its range wraps instead of overflowing" + ); + assert_eq!(cursor.advance(3, 4), (usize::MAX % 4 + 3) % 4); + assert_eq!( + cursor.advance(3, 4), + (usize::MAX % 4 + 6) % 4, + "every pass after the wrap steps by its budget, so one member can't keep the front \ + of the queue" + ); + } + + #[test] + fn a_set_larger_than_its_budget_comes_round_in_turns_that_cover_everybody() { + let accounts: Vec = ["nel", "olaren", "teq", "bailey", "cuttle"] + .iter() + .map(|name| AccountDid::new(format!("did:plc:{name}")).unwrap()) + .collect(); + let cursor = Cursor::default(); + let budget = AccountBudget::new(2); + let turns: Vec> = (0..3) + .map(|_| { + portion(&accounts, budget, &cursor) + .iter() + .map(|did| did.as_str().to_string()) + .collect() + }) + .collect(); + assert_eq!( + turns, + vec![ + vec!["did:plc:nel", "did:plc:olaren"], + vec!["did:plc:teq", "did:plc:bailey"], + vec!["did:plc:cuttle", "did:plc:nel"], + ], + "a stranger offering an unrecognized key mustn't cost the knot a read of every \ + account it grants. The accounts it defers must come round on later passes" + ); + } + + #[test] + fn a_set_the_budget_covers_is_read_whole_without_repeats() { + let accounts: Vec = ["nel", "olaren"] + .iter() + .map(|name| AccountDid::new(format!("did:plc:{name}")).unwrap()) + .collect(); + let cursor = Cursor::default(); + assert_eq!( + portion(&accounts, AccountBudget::new(256), &cursor), + accounts, + "a set that fits inside one budget must behave as though there were no budget" + ); + assert!( + portion(&[], AccountBudget::new(256), &cursor).is_empty(), + "an empty set must yield an empty portion, or the cycle turns forever" + ); + } +} diff --git a/knot2/crates/knot-keyfill/tests/fill.rs b/knot2/crates/knot-keyfill/tests/fill.rs new file mode 100644 --- /dev/null +++ b/knot2/crates/knot-keyfill/tests/fill.rs @@ -0,0 +1,442 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; + +use futures::StreamExt; +use knot_atproto::Atproto; +use knot_cob::{CobHome, CobStore}; +use knot_cobs::{Grant, MembersChange, Registration, RegistryChange}; +use knot_git::{Layout, Repo}; +use knot_index::{Coverage, Index, KeyReprieve, KeyTtl, Resolved, SweepFloor}; +use knot_keyfill::{ + AccountBudget, BusyRetry, Cursors, Pace, SettleFloor, SettledPause, StalledBackoff, fill_once, +}; +use knot_resource::{Burst, HostKey, HostPacer, RateLimit, RefillMicros, Slots}; +use knot_runtime::{ + FakeHttp, HttpRequest, HttpResponse, K256Signer, ManualClock, NetworkError, SeededEntropy, + Signer, UnixMicros, +}; +use knot_types::{ + AccountDid, KnotId, OfferedKey, OwnerDid, RepoDid, RepoName, RepoRkey, UnixSeconds, +}; +use tempfile::TempDir; +use tokio_util::sync::CancellationToken; +use url::Url; + +const KNOT_DID: &str = "did:web:nel.pet"; +const NEL: &str = "did:plc:nel"; +const OLAREN: &str = "did:plc:olaren"; +const TEQ: &str = "did:plc:teq"; +const BAILEY: &str = "did:plc:bailey"; +const NEL_PDS: &str = "https://pds.nel.pet"; +const OLAREN_PDS: &str = "https://pds.olaren.dev"; + +type Responder = Box Result + Send + Sync>; + +fn responding(status: http::StatusCode, body: bytes::Bytes) -> HttpResponse { + HttpResponse { + status, + headers: http::HeaderMap::new(), + body, + } +} + +struct Hosted { + owner: OwnerDid, + repo: RepoDid, + name: RepoName, +} + +fn account(did: &str) -> AccountDid { + AccountDid::new(did).unwrap() +} + +fn anemone() -> Hosted { + Hosted { + owner: OwnerDid::new(NEL).unwrap(), + repo: RepoDid::new("did:plc:squid").unwrap(), + name: RepoName::new("anemone").unwrap(), + } +} + +fn barnacle() -> Hosted { + Hosted { + owner: OwnerDid::new(OLAREN).unwrap(), + repo: RepoDid::new("did:plc:limpet").unwrap(), + name: RepoName::new("barnacle").unwrap(), + } +} + +fn hosted_index(scratch: &TempDir, repos: &[Hosted], members: &[AccountDid]) -> Arc { + let meta_path = scratch.path().join("meta"); + Repo::create(&meta_path).unwrap(); + let layout = Layout::new(scratch.path().join("repos")); + let meta = Repo::open(&meta_path).unwrap(); + let store = CobStore::new(&meta); + let home = CobHome::from(&KnotId::new(KNOT_DID).unwrap()); + let signer = K256Signer::generate(&SeededEntropy::new(3)); + let registration = |hosted: &Hosted| { + layout.create(&hosted.repo).unwrap(); + RegistryChange::Register(Registration { + owner: hosted.owner.clone(), + rkey: RepoRkey::new(hosted.name.as_str()).unwrap(), + name: hosted.name.clone(), + repo: hosted.repo.clone(), + created_at: UnixSeconds::new(1), + }) + }; + let (first, rest) = repos.split_first().expect("a knot under test hosts a repo"); + let registry = store + .create(&home, ®istration(first), &signer, UnixSeconds::new(1)) + .unwrap() + .object; + rest.iter().for_each(|hosted| { + store + .update( + &home, + registry, + ®istration(hosted), + &signer, + UnixSeconds::new(2), + ) + .unwrap(); + }); + + let granted = |subject: &AccountDid| { + MembersChange::Add(Grant { + subject: subject.clone(), + added_by: account(NEL), + created_at: UnixSeconds::new(1), + }) + }; + if let Some((first, rest)) = members.split_first() { + let roll = store + .create(&home, &granted(first), &signer, UnixSeconds::new(1)) + .unwrap() + .object; + rest.iter().for_each(|subject| { + store + .update(&home, roll, &granted(subject), &signer, UnixSeconds::new(2)) + .unwrap(); + }); + } + + let index = Arc::new(Index::new(meta_path, layout)); + index.rebuild().unwrap(); + index.warm_collaborators(); + index +} + +fn atproto_with(responder: Responder) -> Arc, ManualClock>> { + Arc::new(Atproto::new( + FakeHttp::new(responder), + ManualClock::new(UnixMicros::new(1_000_000_000)), + KnotId::new(KNOT_DID).unwrap(), + knot_atproto::PlcDirectory::new(Url::parse("https://plc.directory/").unwrap()).unwrap(), + )) +} + +fn atproto_answering( + status: http::StatusCode, + calls: Arc, +) -> Arc, ManualClock>> { + atproto_with(Box::new(move |_request: &HttpRequest| { + calls.fetch_add(1, Ordering::SeqCst); + Ok(responding(status, bytes::Bytes::new())) + })) +} + +fn did_document(did: &str, handle: &str, pds: &str) -> bytes::Bytes { + let signing = K256Signer::generate(&SeededEntropy::new(5)); + let multibase = knot_types::crypto::multikey(0xe7, signing.public_key().as_bytes()); + let body = serde_json::json!({ + "id": did, + "alsoKnownAs": [format!("at://{handle}")], + "verificationMethod": [{ + "id": format!("{did}#atproto"), + "type": "Multikey", + "controller": did, + "publicKeyMultibase": multibase, + }], + "service": [{ + "id": "#atproto_pds", + "type": "AtprotoPersonalDataServer", + "serviceEndpoint": pds, + }] + }); + bytes::Bytes::from(serde_json::to_vec(&body).unwrap()) +} + +fn atproto_serving_two_accounts() -> Arc, ManualClock>> { + atproto_with(Box::new(move |request: &HttpRequest| { + let url = request.url.as_str(); + let body = if url.contains("listRecords") { + bytes::Bytes::from_static(br#"{"records":[]}"#) + } else if url.ends_with(NEL) { + did_document(NEL, "nel.pet", NEL_PDS) + } else if url.ends_with(OLAREN) { + did_document(OLAREN, "olaren.dev", OLAREN_PDS) + } else { + return Ok(responding(http::StatusCode::NOT_FOUND, bytes::Bytes::new())); + }; + Ok(responding(http::StatusCode::OK, body)) + })) +} + +fn atproto_recording( + seen: Arc>>, +) -> Arc, ManualClock>> { + atproto_with(Box::new(move |request: &HttpRequest| { + let url = request.url.as_str(); + if url.contains("listRecords") { + return Ok(responding( + http::StatusCode::OK, + bytes::Bytes::from_static(br#"{"records":[]}"#), + )); + } + if url.ends_with(NEL) { + return Ok(responding( + http::StatusCode::OK, + did_document(NEL, "nel.pet", NEL_PDS), + )); + } + if let Some(did) = url.rsplit('/').next() { + seen.lock().unwrap().push(did.to_string()); + } + Ok(responding( + http::StatusCode::SERVICE_UNAVAILABLE, + bytes::Bytes::new(), + )) + })) +} + +fn now() -> UnixSeconds { + UnixSeconds::new(1_000) +} + +fn brisk() -> Pace { + Pace { + busy: BusyRetry::from_millis(0), + floor: SettleFloor::from_millis(0), + ttl: KeyTtl::from_secs(3_600), + reprieve: KeyReprieve::from_secs(300, 21_600), + sweep: SweepFloor::DEFAULT, + stalled: StalledBackoff::from_secs(1), + settled: SettledPause::from_secs(1), + members: AccountBudget::new(64), + suspected: AccountBudget::new(256), + host: RateLimit { + burst: Burst::new(1), + refill: RefillMicros::new(1_000), + }, + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_gone_document_completes_the_set_and_an_outage_leaves_it_warming() { + let outcome = |status: http::StatusCode| async move { + let scratch = tempfile::tempdir().unwrap(); + let index = hosted_index(&scratch, &[anemone()], &[]); + let calls = Arc::new(AtomicUsize::new(0)); + let atproto = atproto_answering(status, Arc::clone(&calls)); + let pacer = HostPacer::new(brisk().host); + assert_eq!(index.keys().coverage(), Coverage::Warming); + fill_once( + &index, + &atproto, + &Slots::testing(4), + &pacer, + brisk(), + &Cursors::default(), + ) + .await; + assert!( + calls.load(Ordering::SeqCst) > 0, + "the fill made an outbound request" + ); + index.keys().coverage() + }; + + assert_eq!( + outcome(http::StatusCode::NOT_FOUND).await, + Coverage::Ready, + "a permanently unresolvable owner is recorded with an empty key set, so the set is complete" + ); + assert_eq!( + outcome(http::StatusCode::SERVICE_UNAVAILABLE).await, + Coverage::Warming, + "a transient failure doesn't teach the knot anything, so it mustn't claim a complete set" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn keys_the_knot_already_read_survive_an_outage_and_a_refused_listing() { + let scratch = tempfile::tempdir().unwrap(); + let index = hosted_index(&scratch, &[anemone()], &[]); + let key = OfferedKey::from_bytes(vec![9]); + let spent = KeyTtl::from_secs(1).lease_from(UnixSeconds::new(100)); + index.keys().record(&account(NEL), vec![key.clone()], spent); + let pacer = HostPacer::new(brisk().host); + + let unreachable = atproto_answering(http::StatusCode::SERVICE_UNAVAILABLE, Arc::default()); + fill_once( + &index, + &unreachable, + &Slots::testing(4), + &pacer, + brisk(), + &Cursors::default(), + ) + .await; + assert_eq!( + index.keys().coverage(), + Coverage::Ready, + "an owner the knot has read before keeps its last keys through an outage, so one \ + unreachable PDS mustn't reopen the knot to every offered key" + ); + assert_eq!( + index.owner_of_key(&key, now()), + Resolved::Ready(Some(account(NEL))), + "the reprieve keeps the keys the knot last read" + ); + + index.keys().record(&account(NEL), vec![key.clone()], spent); + let listing = Arc::new(AtomicUsize::new(0)); + let refusing = { + let listing = Arc::clone(&listing); + atproto_with(Box::new(move |request: &HttpRequest| { + let url = request.url.as_str(); + if url.contains("listRecords") { + listing.fetch_add(1, Ordering::SeqCst); + return Ok(responding( + http::StatusCode::BAD_REQUEST, + bytes::Bytes::new(), + )); + } + Ok(responding( + http::StatusCode::OK, + did_document(NEL, "nel.pet", NEL_PDS), + )) + })) + }; + fill_once( + &index, + &refusing, + &Slots::testing(4), + &pacer, + brisk(), + &Cursors::default(), + ) + .await; + assert!( + listing.load(Ordering::SeqCst) > 0, + "the fill read from the PDS" + ); + assert_eq!( + index.owner_of_key(&key, now()), + Resolved::Ready(Some(account(NEL))), + "an account whose DID document resolves hasn't gone anywhere, so the knot mustn't read \ + a 400 from a record listing as proof the account stopped publishing keys" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn the_fill_takes_a_turn_at_each_host_it_reads_from_and_stops_on_shutdown() { + let scratch = tempfile::tempdir().unwrap(); + let index = hosted_index(&scratch, &[anemone(), barnacle()], &[]); + let atproto = atproto_serving_two_accounts(); + let pacer = HostPacer::new(brisk().host); + + fill_once( + &index, + &atproto, + &Slots::testing(4), + &pacer, + brisk(), + &Cursors::default(), + ) + .await; + + assert_eq!( + index.keys().coverage(), + Coverage::Ready, + "both owners resolved, so every account that may push has a record" + ); + ["plc.directory", "pds.nel.pet", "pds.olaren.dev"] + .iter() + .for_each(|host| { + assert!( + !pacer.reserve_now(&HostKey::new(host), UnixMicros::new(0)), + "the fill must take a turn at {host} before it reads from it, \ + or a knot whose members share one PDS spends its whole rate at that host" + ); + }); + assert!( + pacer.reserve_now(&HostKey::new("pds.teq.dev"), UnixMicros::new(0)), + "a host the fill never read from is due immediately. The bookings above are the fill's \ + own work" + ); + + let shutdown = CancellationToken::new(); + let task = knot_keyfill::spawn( + Arc::clone(&index), + Arc::clone(&atproto), + Slots::testing(4), + brisk(), + shutdown.clone(), + ); + shutdown.cancel(); + tokio::time::timeout(std::time::Duration::from_secs(5), task) + .await + .expect("a shutting-down knot mustn't wait out the pause between fill passes") + .expect("the fill task stops without panicking"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn members_are_read_whole_on_first_contact_then_renewed_one_budget_turn_per_pass() { + let scratch = tempfile::tempdir().unwrap(); + let members = [account(OLAREN), account(TEQ), account(BAILEY)]; + let index = hosted_index(&scratch, &[anemone()], &members); + let seen = Arc::new(Mutex::new(Vec::new())); + let atproto = atproto_recording(Arc::clone(&seen)); + let pacer = HostPacer::new(brisk().host); + let pace = Pace { + members: AccountBudget::new(1), + ..brisk() + }; + let cursor = Cursors::default(); + + fill_once(&index, &atproto, &Slots::testing(4), &pacer, pace, &cursor).await; + assert_eq!( + index.keys().coverage(), + Coverage::Ready, + "the pushers are what coverage waits on, so an unreadable member mustn't make the knot \ + doubt the keys it checks pushes against" + ); + let mut attempted = seen.lock().unwrap().clone(); + attempted.sort(); + assert_eq!( + attempted, + vec![BAILEY.to_string(), OLAREN.to_string(), TEQ.to_string()], + "the renewal budget paces rereads, so a member the knot has never read mustn't wait \ + its turn behind it and be refused at the handshake for the passes in between" + ); + + let spent = KeyTtl::from_secs(1).lease_from(UnixSeconds::new(0)); + members + .iter() + .for_each(|member| _ = index.keys().record(member, Vec::new(), spent)); + seen.lock().unwrap().clear(); + + futures::stream::iter(0..3) + .for_each(|_| async { + fill_once(&index, &atproto, &Slots::testing(4), &pacer, pace, &cursor).await; + }) + .await; + let order = seen.lock().unwrap().clone(); + assert_eq!( + order.iter().map(String::as_str).collect::>(), + vec![BAILEY, OLAREN, TEQ], + "one member per pass in turn, or members whose PDS stays down keep the front of \ + the queue and the knot never reads the rest" + ); +}