knot2/keyfill: fill index w/ published keys off auth path · alice.pds.demo.boltless.dev/core@ef66668

+23

Cargo.lock

··· 4555 4555 ] 4556 4556 4557 4557 [[package]] 4558 + name = "knot-keyfill" 4559 + version = "2.0.0" 4560 + dependencies = [ 4561 + "bytes", 4562 + "futures", 4563 + "http", 4564 + "knot-atproto", 4565 + "knot-cob", 4566 + "knot-cobs", 4567 + "knot-git", 4568 + "knot-index", 4569 + "knot-resource", 4570 + "knot-runtime", 4571 + "knot-types", 4572 + "serde_json", 4573 + "tempfile", 4574 + "tokio", 4575 + "tokio-util", 4576 + "tracing", 4577 + "url", 4578 + ] 4579 + 4580 + [[package]] 4558 4581 name = "knot-langs" 4559 4582 version = "2.0.0" 4560 4583 dependencies = [

+1

Cargo.toml

View file

··· 57 57 knot-cob = { path = "knot2/crates/knot-cob" } 58 58 knot-cobs = { path = "knot2/crates/knot-cobs" } 59 59 knot-index = { path = "knot2/crates/knot-index" } 60 + knot-keyfill = { path = "knot2/crates/knot-keyfill" } 60 61 knot-cache = { path = "knot2/crates/knot-cache" } 61 62 knot-acl = { path = "knot2/crates/knot-acl" } 62 63 knot-atproto = { path = "knot2/crates/knot-atproto" }

+27

knot2/crates/knot-keyfill/Cargo.toml

View file

··· 1 + [package] 2 + name = "knot-keyfill" 3 + version = "2.0.0" 4 + edition.workspace = true 5 + rust-version.workspace = true 6 + license.workspace = true 7 + 8 + [dependencies] 9 + knot-atproto = { workspace = true } 10 + knot-index = { workspace = true } 11 + knot-resource = { workspace = true } 12 + knot-runtime = { workspace = true } 13 + knot-types = { workspace = true } 14 + futures = { workspace = true } 15 + tokio = { workspace = true } 16 + tokio-util = { workspace = true } 17 + tracing = { workspace = true } 18 + url = { workspace = true } 19 + 20 + [dev-dependencies] 21 + knot-cob = { workspace = true } 22 + knot-cobs = { workspace = true } 23 + knot-git = { workspace = true } 24 + bytes = { workspace = true } 25 + http = { workspace = true } 26 + serde_json = { workspace = true } 27 + tempfile = { workspace = true }

+619

knot2/crates/knot-keyfill/src/lib.rs

View file

··· 1 + use std::future::Future; 2 + use std::sync::Arc; 3 + use std::sync::atomic::{AtomicUsize, Ordering}; 4 + use std::time::Duration; 5 + 6 + use futures::future::OptionFuture; 7 + use futures::{FutureExt, StreamExt}; 8 + use knot_atproto::{Atproto, AtprotoError}; 9 + use knot_index::{ 10 + Coverage, HostedCoverage, Index, IndexGeneration, KeyLease, KeyRecord, KeyReprieve, 11 + KeyReprieved, KeyTtl, MemberWork, Pushers, Resolved, StalePushers, SuspectPushers, SweepFloor, 12 + }; 13 + use knot_resource::{Burst, HostKey, HostPacer, RateLimit, RefillMicros, SlotPermit, Slots}; 14 + use knot_runtime::{Clock, HttpTransport}; 15 + use knot_types::{AccountDid, OfferedKey, UnixMicros, UnixSeconds}; 16 + use tokio::sync::watch; 17 + use tokio_util::sync::CancellationToken; 18 + use url::Url; 19 + 20 + const FILL_FANOUT: usize = 8; 21 + 22 + const PASS_HEADROOM: u32 = 2; 23 + 24 + macro_rules! span { 25 + ($($name:ident from $unit:ident),+ $(,)?) => {$( 26 + #[derive(Debug, Clone, Copy, PartialEq, Eq)] 27 + pub struct $name(Duration); 28 + 29 + impl $name { 30 + pub const fn $unit(value: u64) -> Self { 31 + Self(Duration::$unit(value)) 32 + } 33 + 34 + pub const fn get(self) -> Duration { 35 + self.0 36 + } 37 + } 38 + )+}; 39 + } 40 + 41 + span!( 42 + BusyRetry from from_millis, 43 + SettleFloor from from_millis, 44 + StalledBackoff from from_secs, 45 + SettledPause from from_secs, 46 + ); 47 + 48 + #[derive(Debug, Clone, Copy, PartialEq, Eq)] 49 + pub struct AccountBudget(usize); 50 + 51 + impl AccountBudget { 52 + pub const fn new(accounts: usize) -> Self { 53 + Self(if accounts == 0 { 1 } else { accounts }) 54 + } 55 + 56 + pub const fn get(self) -> usize { 57 + self.0 58 + } 59 + } 60 + 61 + #[derive(Debug, Default)] 62 + pub struct Cursor(AtomicUsize); 63 + 64 + impl Cursor { 65 + fn advance(&self, by: usize, len: usize) -> usize { 66 + match len { 67 + 0 => 0, 68 + len => self 69 + .0 70 + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |seen| { 71 + Some((seen % len).wrapping_add(by) % len) 72 + }) 73 + .map_or(0, |seen| seen % len), 74 + } 75 + } 76 + } 77 + 78 + #[derive(Debug, Default)] 79 + pub struct Cursors { 80 + members: Cursor, 81 + suspected: Cursor, 82 + } 83 + 84 + fn portion(accounts: &[AccountDid], budget: AccountBudget, cursor: &Cursor) -> Vec<AccountDid> { 85 + let taken = budget.get().min(accounts.len()); 86 + let start = cursor.advance(taken, accounts.len()); 87 + accounts 88 + .iter() 89 + .cycle() 90 + .skip(start) 91 + .take(taken) 92 + .cloned() 93 + .collect() 94 + } 95 + 96 + #[derive(Debug, Clone, Copy)] 97 + pub struct Pace { 98 + pub busy: BusyRetry, 99 + pub floor: SettleFloor, 100 + pub ttl: KeyTtl, 101 + pub reprieve: KeyReprieve, 102 + pub sweep: SweepFloor, 103 + pub stalled: StalledBackoff, 104 + pub settled: SettledPause, 105 + pub members: AccountBudget, 106 + pub suspected: AccountBudget, 107 + pub host: RateLimit, 108 + } 109 + 110 + impl Default for Pace { 111 + fn default() -> Self { 112 + Self { 113 + busy: BusyRetry::from_millis(50), 114 + floor: SettleFloor::from_millis(1_000), 115 + ttl: KeyTtl::DEFAULT, 116 + reprieve: KeyReprieve::DEFAULT, 117 + sweep: SweepFloor::DEFAULT, 118 + stalled: StalledBackoff::from_secs(30), 119 + settled: SettledPause::from_secs(60), 120 + members: AccountBudget::new(64), 121 + suspected: AccountBudget::new(256), 122 + host: RateLimit { 123 + burst: Burst::new(10), 124 + refill: RefillMicros::new(200_000), 125 + }, 126 + } 127 + } 128 + } 129 + 130 + impl Pace { 131 + fn ttl_covering(self, accounts: usize) -> KeyTtl { 132 + let micros = (accounts as u64).saturating_mul(self.host.interval().get()); 133 + let secs = (micros / 1_000_000).saturating_mul(u64::from(PASS_HEADROOM)); 134 + self.ttl.longest(KeyTtl::from_secs(secs)) 135 + } 136 + } 137 + 138 + pub fn spawn<H: HttpTransport, C: Clock>( 139 + index: Arc<Index>, 140 + atproto: Arc<Atproto<H, C>>, 141 + slots: Slots, 142 + pace: Pace, 143 + shutdown: CancellationToken, 144 + ) -> tokio::task::JoinHandle<()> { 145 + let stop = shutdown.clone(); 146 + let driver = Driver { 147 + generations: index.generations(), 148 + pacer: HostPacer::new(pace.host), 149 + cursors: Cursors::default(), 150 + index, 151 + atproto, 152 + slots, 153 + pace, 154 + shutdown, 155 + }; 156 + tokio::spawn(async move { 157 + futures::stream::unfold(driver, |mut driver| async move { 158 + driver.pass().await; 159 + Some(((), driver)) 160 + }) 161 + .take_until(stop.cancelled_owned()) 162 + .for_each(|()| std::future::ready(())) 163 + .await; 164 + tracing::info!("key fill stopped"); 165 + }) 166 + } 167 + 168 + struct Driver<H, C> { 169 + index: Arc<Index>, 170 + atproto: Arc<Atproto<H, C>>, 171 + slots: Slots, 172 + pacer: HostPacer, 173 + cursors: Cursors, 174 + pace: Pace, 175 + shutdown: CancellationToken, 176 + generations: watch::Receiver<IndexGeneration>, 177 + } 178 + 179 + impl<H: HttpTransport, C: Clock> Driver<H, C> { 180 + async fn pass(&mut self) { 181 + self.generations.mark_unchanged(); 182 + let filling = fill_once( 183 + &self.index, 184 + &self.atproto, 185 + &self.slots, 186 + &self.pacer, 187 + self.pace, 188 + &self.cursors, 189 + ); 190 + let pause = tokio::select! { 191 + pause = guarded(filling, self.pace) => pause, 192 + () = self.shutdown.cancelled() => self.pace.stalled.get(), 193 + }; 194 + tokio::select! { 195 + () = self.shutdown.cancelled() => {} 196 + () = settle(&mut self.generations, pause, self.pace.floor) => {} 197 + } 198 + } 199 + } 200 + 201 + async fn guarded(pass: impl Future<Output = Duration>, pace: Pace) -> Duration { 202 + std::panic::AssertUnwindSafe(pass) 203 + .catch_unwind() 204 + .await 205 + .unwrap_or_else(|_| { 206 + tracing::error!("key fill pass panicked, backing off before the next pass"); 207 + pace.stalled.get() 208 + }) 209 + } 210 + 211 + async fn settle( 212 + generations: &mut watch::Receiver<IndexGeneration>, 213 + pause: Duration, 214 + floor: SettleFloor, 215 + ) { 216 + let held = floor.get().min(pause); 217 + tokio::time::sleep(held).await; 218 + let _ = tokio::time::timeout(pause.saturating_sub(held), generations.changed()).await; 219 + } 220 + 221 + struct Pass<'a, H, C> { 222 + index: &'a Arc<Index>, 223 + atproto: &'a Arc<Atproto<H, C>>, 224 + slots: &'a Slots, 225 + pacer: &'a HostPacer, 226 + pace: Pace, 227 + now: UnixSeconds, 228 + lease: KeyLease, 229 + reprieve: KeyReprieve, 230 + } 231 + 232 + pub async fn fill_once<H: HttpTransport, C: Clock>( 233 + index: &Arc<Index>, 234 + atproto: &Arc<Atproto<H, C>>, 235 + slots: &Slots, 236 + pacer: &HostPacer, 237 + pace: Pace, 238 + cursors: &Cursors, 239 + ) -> Duration { 240 + fold_hosted(index).await; 241 + let now = atproto.now().seconds(); 242 + let Resolved::Ready(work) = index.keys().work(now, pace.sweep) else { 243 + index.keys().mark_warming(); 244 + return pace.stalled.get(); 245 + }; 246 + if let HostedCoverage::Partial { unread } = work.hosted { 247 + tracing::debug!( 248 + repos = unread, 249 + "partial grant set, a repo was registered while this pass was working out who may push" 250 + ); 251 + } 252 + let members = match (work.hosted, work.members) { 253 + (HostedCoverage::Whole, Resolved::Ready(members)) => { 254 + index.keys().retain(&members.kept); 255 + Some(members) 256 + } 257 + (HostedCoverage::Partial { .. }, Resolved::Ready(members)) => Some(members), 258 + (_, Resolved::Warming) => None, 259 + }; 260 + let ttl = pace.ttl_covering(work.tracked); 261 + if ttl != pace.ttl { 262 + tracing::debug!( 263 + tracked = work.tracked, 264 + ttl_secs = ttl.get().as_secs(), 265 + "one paced pass over the grant set outruns the key ttl, so the fill stretches it" 266 + ); 267 + } 268 + let pass = Pass { 269 + index, 270 + atproto, 271 + slots, 272 + pacer, 273 + pace, 274 + now, 275 + lease: ttl.lease_from(now), 276 + reprieve: pace.reprieve.budgeted_for(ttl), 277 + }; 278 + let pause = fill_pushers( 279 + &pass, 280 + PushWork { 281 + generation: work.generation, 282 + hosted: work.hosted, 283 + pushers: work.pushers, 284 + due: work.due, 285 + suspected: work.suspected, 286 + complete: work.complete, 287 + }, 288 + &cursors.suspected, 289 + ) 290 + .await; 291 + OptionFuture::from(members.map(|work| fill_members(&pass, work, &cursors.members))).await; 292 + pause 293 + } 294 + 295 + struct PushWork { 296 + generation: IndexGeneration, 297 + hosted: HostedCoverage, 298 + pushers: Pushers, 299 + due: StalePushers, 300 + suspected: SuspectPushers, 301 + complete: bool, 302 + } 303 + 304 + async fn fold_hosted(index: &Arc<Index>) { 305 + let index = Arc::clone(index); 306 + match tokio::task::spawn_blocking(move || index.warm_collaborators()).await { 307 + Ok(0) | Err(_) => {} 308 + Ok(unreadable) => tracing::warn!( 309 + repos = unreadable, 310 + "the knot hosts registered repos it can't open, so it won't grant anybody through \ 311 + them until an operator restores or deregisters each repo" 312 + ), 313 + } 314 + } 315 + 316 + async fn fill_pushers<H: HttpTransport, C: Clock>( 317 + pass: &Pass<'_, H, C>, 318 + work: PushWork, 319 + cursor: &Cursor, 320 + ) -> Duration { 321 + let PushWork { 322 + generation, 323 + hosted, 324 + pushers, 325 + due, 326 + suspected, 327 + complete, 328 + } = work; 329 + if !complete { 330 + pass.index.keys().mark_warming(); 331 + } 332 + let wanted = due.len(); 333 + let recorded = match wanted { 334 + 0 => 0, 335 + _ => { 336 + let recorded = record_each(pass, due.into_vec()).await; 337 + tracing::debug!(wanted, recorded, "pusher key fill pass"); 338 + recorded 339 + } 340 + }; 341 + recheck_suspected(pass, &suspected, cursor).await; 342 + match recorded < wanted { 343 + true => pass.pace.stalled.get(), 344 + false => claim_ready(pass, generation, hosted, &pushers), 345 + } 346 + } 347 + 348 + async fn recheck_suspected<H: HttpTransport, C: Clock>( 349 + pass: &Pass<'_, H, C>, 350 + suspected: &SuspectPushers, 351 + cursor: &Cursor, 352 + ) { 353 + if suspected.is_empty() { 354 + return; 355 + } 356 + let batch = portion(suspected.as_slice(), pass.pace.suspected, cursor); 357 + let wanted = batch.len(); 358 + let recorded = record_each(pass, batch).await; 359 + tracing::debug!( 360 + wanted, 361 + recorded, 362 + deferred = suspected.len().saturating_sub(wanted), 363 + "pusher key recheck pass, a client offered a key the accounts on file don't publish" 364 + ); 365 + } 366 + 367 + fn claim_ready<H: HttpTransport, C: Clock>( 368 + pass: &Pass<'_, H, C>, 369 + generation: IndexGeneration, 370 + hosted: HostedCoverage, 371 + pushers: &Pushers, 372 + ) -> Duration { 373 + let keys = pass.index.keys(); 374 + let settled = 375 + hosted == HostedCoverage::Whole && keys.all_live(pushers, pass.atproto.now().seconds()); 376 + if !settled { 377 + return pass.pace.stalled.get(); 378 + } 379 + keys.mark_ready(generation); 380 + match keys.coverage() { 381 + Coverage::Ready => pass.pace.settled.get(), 382 + Coverage::Warming => pass.pace.stalled.get(), 383 + } 384 + } 385 + 386 + async fn fill_members<H: HttpTransport, C: Clock>( 387 + pass: &Pass<'_, H, C>, 388 + work: MemberWork, 389 + cursor: &Cursor, 390 + ) { 391 + if !work.unread.is_empty() { 392 + let wanted = work.unread.len(); 393 + let recorded = record_each(pass, work.unread.into_vec()).await; 394 + tracing::debug!(wanted, recorded, "member key first-read pass"); 395 + } 396 + if work.due.is_empty() { 397 + return; 398 + } 399 + let batch = portion(work.due.as_slice(), pass.pace.members, cursor); 400 + let wanted = batch.len(); 401 + let recorded = record_each(pass, batch).await; 402 + tracing::debug!( 403 + wanted, 404 + recorded, 405 + deferred = work.due.len().saturating_sub(wanted), 406 + "member key renewal pass" 407 + ); 408 + } 409 + 410 + async fn record_each<H: HttpTransport, C: Clock>( 411 + pass: &Pass<'_, H, C>, 412 + stale: Vec<AccountDid>, 413 + ) -> usize { 414 + futures::stream::iter(stale) 415 + .map(|did| async move { 416 + match published_keys(pass, &did).await { 417 + Ok(keys) => usize::from(record(pass, &did, keys)), 418 + Err(error) if error.is_gone() => { 419 + tracing::debug!( 420 + did = did.as_str(), 421 + %error, 422 + "key fill records an empty key set for an account whose DID document is gone" 423 + ); 424 + usize::from(record(pass, &did, Vec::new())) 425 + } 426 + Err(error) => reprieve(pass, &did, error), 427 + } 428 + }) 429 + .buffer_unordered(FILL_FANOUT) 430 + .fold(0, |total, recorded| async move { total + recorded }) 431 + .await 432 + } 433 + 434 + fn reprieve<H, C>(pass: &Pass<'_, H, C>, did: &AccountDid, error: AtprotoError) -> usize { 435 + let outcome = pass 436 + .index 437 + .keys() 438 + .reprieve(did, pass.now, pass.reprieve, pass.lease); 439 + match outcome { 440 + KeyReprieved::Exhausted => tracing::warn!( 441 + did = did.as_str(), 442 + %error, 443 + "the knot spent the whole reprieve failing to read an account, so it records an \ 444 + empty key set for the account until a later read succeeds" 445 + ), 446 + KeyReprieved::Extended | KeyReprieved::Pending => tracing::debug!( 447 + did = did.as_str(), 448 + %error, 449 + ?outcome, 450 + "key fill couldn't read an account's records" 451 + ), 452 + } 453 + match outcome { 454 + KeyReprieved::Extended | KeyReprieved::Exhausted => 1, 455 + KeyReprieved::Pending => 0, 456 + } 457 + } 458 + 459 + fn record<H, C>(pass: &Pass<'_, H, C>, did: &AccountDid, keys: Vec<OfferedKey>) -> bool { 460 + match pass.index.keys().record(did, keys, pass.lease) { 461 + KeyRecord::Stored => true, 462 + KeyRecord::Unheld => { 463 + tracing::warn!( 464 + did = did.as_str(), 465 + "the key set is full, so the knot will check this account's pushes against its \ 466 + PDS every time instead of against the set" 467 + ); 468 + true 469 + } 470 + KeyRecord::Saturated => { 471 + tracing::warn!( 472 + did = did.as_str(), 473 + "the key budget can't record even that it read this account, so the knot keeps \ 474 + reporting the set incomplete and defers every offered key to the push check" 475 + ); 476 + false 477 + } 478 + } 479 + } 480 + 481 + async fn published_keys<H: HttpTransport, C: Clock>( 482 + pass: &Pass<'_, H, C>, 483 + did: &AccountDid, 484 + ) -> Result<Vec<OfferedKey>, AtprotoError> { 485 + let atproto = pass.atproto; 486 + if let Ok(document) = atproto.document_url(did) { 487 + wait_for_turn(pass.pacer, &document, atproto.now()).await; 488 + } 489 + let identity = { 490 + let _permit = idle_permit(pass.slots, pass.pace).await; 491 + atproto.resolve_identity(did).await? 492 + }; 493 + wait_for_turn(pass.pacer, identity.pds.url(), atproto.now()).await; 494 + let _permit = idle_permit(pass.slots, pass.pace).await; 495 + atproto.pubkeys_at(&identity, did).await 496 + } 497 + 498 + async fn wait_for_turn(pacer: &HostPacer, url: &Url, now: UnixMicros) { 499 + if let Some(host) = url.host_str().map(HostKey::new) { 500 + tokio::time::sleep(pacer.reserve(&host, now)).await; 501 + } 502 + } 503 + 504 + async fn idle_permit(slots: &Slots, pace: Pace) -> SlotPermit { 505 + let attempts = futures::stream::repeat(()).filter_map(|()| async { 506 + match slots.resolve.try_acquire() { 507 + Some(permit) => Some(permit), 508 + None => { 509 + tokio::time::sleep(pace.busy.get()).await; 510 + None 511 + } 512 + } 513 + }); 514 + futures::pin_mut!(attempts); 515 + attempts 516 + .next() 517 + .await 518 + .expect("an endless stream of attempts yields a permit") 519 + } 520 + 521 + #[cfg(test)] 522 + mod tests { 523 + use super::*; 524 + 525 + #[test] 526 + fn a_grant_set_the_pacer_cant_reread_within_the_ttl_stretches_it() { 527 + let pace = Pace::default(); 528 + assert_eq!( 529 + pace.ttl_covering(1_000), 530 + KeyTtl::DEFAULT, 531 + "a set one paced pass covers well inside the ttl keeps the ttl it was configured with" 532 + ); 533 + assert_eq!( 534 + pace.ttl_covering(100_000), 535 + KeyTtl::from_secs(40_000), 536 + "at one directory turn per 200ms a hundred thousand accounts take 20_000s to reread, \ 537 + so the entries the pass wrote first must outlive the pass that writes the last" 538 + ); 539 + } 540 + 541 + #[test] 542 + fn the_reprieve_budget_never_undercuts_the_ttl_the_fill_is_working_to() { 543 + let stretched = KeyTtl::from_secs(40_000); 544 + assert_eq!( 545 + KeyReprieve::DEFAULT.budgeted_for(stretched), 546 + KeyReprieve::from_secs(300, 40_000), 547 + "an account would be released while the pass that would reread it is still running, \ 548 + if the budget stayed under the ttl" 549 + ); 550 + assert_eq!( 551 + KeyReprieve::DEFAULT.budgeted_for(KeyTtl::DEFAULT), 552 + KeyReprieve::DEFAULT, 553 + "a ttl the budget already covers leaves the budget alone" 554 + ); 555 + } 556 + 557 + #[test] 558 + fn the_member_cursor_turns_over_without_running_past_its_type() { 559 + let cursor = Cursor(AtomicUsize::new(usize::MAX)); 560 + assert_eq!( 561 + cursor.advance(3, 4), 562 + usize::MAX % 4, 563 + "a cursor at the end of its range wraps instead of overflowing" 564 + ); 565 + assert_eq!(cursor.advance(3, 4), (usize::MAX % 4 + 3) % 4); 566 + assert_eq!( 567 + cursor.advance(3, 4), 568 + (usize::MAX % 4 + 6) % 4, 569 + "every pass after the wrap steps by its budget, so one member can't keep the front \ 570 + of the queue" 571 + ); 572 + } 573 + 574 + #[test] 575 + fn a_set_larger_than_its_budget_comes_round_in_turns_that_cover_everybody() { 576 + let accounts: Vec<AccountDid> = ["nel", "olaren", "teq", "bailey", "cuttle"] 577 + .iter() 578 + .map(|name| AccountDid::new(format!("did:plc:{name}")).unwrap()) 579 + .collect(); 580 + let cursor = Cursor::default(); 581 + let budget = AccountBudget::new(2); 582 + let turns: Vec<Vec<String>> = (0..3) 583 + .map(|_| { 584 + portion(&accounts, budget, &cursor) 585 + .iter() 586 + .map(|did| did.as_str().to_string()) 587 + .collect() 588 + }) 589 + .collect(); 590 + assert_eq!( 591 + turns, 592 + vec![ 593 + vec!["did:plc:nel", "did:plc:olaren"], 594 + vec!["did:plc:teq", "did:plc:bailey"], 595 + vec!["did:plc:cuttle", "did:plc:nel"], 596 + ], 597 + "a stranger offering an unrecognized key mustn't cost the knot a read of every \ 598 + account it grants. The accounts it defers must come round on later passes" 599 + ); 600 + } 601 + 602 + #[test] 603 + fn a_set_the_budget_covers_is_read_whole_without_repeats() { 604 + let accounts: Vec<AccountDid> = ["nel", "olaren"] 605 + .iter() 606 + .map(|name| AccountDid::new(format!("did:plc:{name}")).unwrap()) 607 + .collect(); 608 + let cursor = Cursor::default(); 609 + assert_eq!( 610 + portion(&accounts, AccountBudget::new(256), &cursor), 611 + accounts, 612 + "a set that fits inside one budget must behave as though there were no budget" 613 + ); 614 + assert!( 615 + portion(&[], AccountBudget::new(256), &cursor).is_empty(), 616 + "an empty set must yield an empty portion, or the cycle turns forever" 617 + ); 618 + } 619 + }

+442

knot2/crates/knot-keyfill/tests/fill.rs

View file

··· 1 + use std::sync::atomic::{AtomicUsize, Ordering}; 2 + use std::sync::{Arc, Mutex}; 3 + 4 + use futures::StreamExt; 5 + use knot_atproto::Atproto; 6 + use knot_cob::{CobHome, CobStore}; 7 + use knot_cobs::{Grant, MembersChange, Registration, RegistryChange}; 8 + use knot_git::{Layout, Repo}; 9 + use knot_index::{Coverage, Index, KeyReprieve, KeyTtl, Resolved, SweepFloor}; 10 + use knot_keyfill::{ 11 + AccountBudget, BusyRetry, Cursors, Pace, SettleFloor, SettledPause, StalledBackoff, fill_once, 12 + }; 13 + use knot_resource::{Burst, HostKey, HostPacer, RateLimit, RefillMicros, Slots}; 14 + use knot_runtime::{ 15 + FakeHttp, HttpRequest, HttpResponse, K256Signer, ManualClock, NetworkError, SeededEntropy, 16 + Signer, UnixMicros, 17 + }; 18 + use knot_types::{ 19 + AccountDid, KnotId, OfferedKey, OwnerDid, RepoDid, RepoName, RepoRkey, UnixSeconds, 20 + }; 21 + use tempfile::TempDir; 22 + use tokio_util::sync::CancellationToken; 23 + use url::Url; 24 + 25 + const KNOT_DID: &str = "did:web:nel.pet"; 26 + const NEL: &str = "did:plc:nel"; 27 + const OLAREN: &str = "did:plc:olaren"; 28 + const TEQ: &str = "did:plc:teq"; 29 + const BAILEY: &str = "did:plc:bailey"; 30 + const NEL_PDS: &str = "https://pds.nel.pet"; 31 + const OLAREN_PDS: &str = "https://pds.olaren.dev"; 32 + 33 + type Responder = Box<dyn Fn(&HttpRequest) -> Result<HttpResponse, NetworkError> + Send + Sync>; 34 + 35 + fn responding(status: http::StatusCode, body: bytes::Bytes) -> HttpResponse { 36 + HttpResponse { 37 + status, 38 + headers: http::HeaderMap::new(), 39 + body, 40 + } 41 + } 42 + 43 + struct Hosted { 44 + owner: OwnerDid, 45 + repo: RepoDid, 46 + name: RepoName, 47 + } 48 + 49 + fn account(did: &str) -> AccountDid { 50 + AccountDid::new(did).unwrap() 51 + } 52 + 53 + fn anemone() -> Hosted { 54 + Hosted { 55 + owner: OwnerDid::new(NEL).unwrap(), 56 + repo: RepoDid::new("did:plc:squid").unwrap(), 57 + name: RepoName::new("anemone").unwrap(), 58 + } 59 + } 60 + 61 + fn barnacle() -> Hosted { 62 + Hosted { 63 + owner: OwnerDid::new(OLAREN).unwrap(), 64 + repo: RepoDid::new("did:plc:limpet").unwrap(), 65 + name: RepoName::new("barnacle").unwrap(), 66 + } 67 + } 68 + 69 + fn hosted_index(scratch: &TempDir, repos: &[Hosted], members: &[AccountDid]) -> Arc<Index> { 70 + let meta_path = scratch.path().join("meta"); 71 + Repo::create(&meta_path).unwrap(); 72 + let layout = Layout::new(scratch.path().join("repos")); 73 + let meta = Repo::open(&meta_path).unwrap(); 74 + let store = CobStore::new(&meta); 75 + let home = CobHome::from(&KnotId::new(KNOT_DID).unwrap()); 76 + let signer = K256Signer::generate(&SeededEntropy::new(3)); 77 + let registration = |hosted: &Hosted| { 78 + layout.create(&hosted.repo).unwrap(); 79 + RegistryChange::Register(Registration { 80 + owner: hosted.owner.clone(), 81 + rkey: RepoRkey::new(hosted.name.as_str()).unwrap(), 82 + name: hosted.name.clone(), 83 + repo: hosted.repo.clone(), 84 + created_at: UnixSeconds::new(1), 85 + }) 86 + }; 87 + let (first, rest) = repos.split_first().expect("a knot under test hosts a repo"); 88 + let registry = store 89 + .create(&home, &registration(first), &signer, UnixSeconds::new(1)) 90 + .unwrap() 91 + .object; 92 + rest.iter().for_each(|hosted| { 93 + store 94 + .update( 95 + &home, 96 + registry, 97 + &registration(hosted), 98 + &signer, 99 + UnixSeconds::new(2), 100 + ) 101 + .unwrap(); 102 + }); 103 + 104 + let granted = |subject: &AccountDid| { 105 + MembersChange::Add(Grant { 106 + subject: subject.clone(), 107 + added_by: account(NEL), 108 + created_at: UnixSeconds::new(1), 109 + }) 110 + }; 111 + if let Some((first, rest)) = members.split_first() { 112 + let roll = store 113 + .create(&home, &granted(first), &signer, UnixSeconds::new(1)) 114 + .unwrap() 115 + .object; 116 + rest.iter().for_each(|subject| { 117 + store 118 + .update(&home, roll, &granted(subject), &signer, UnixSeconds::new(2)) 119 + .unwrap(); 120 + }); 121 + } 122 + 123 + let index = Arc::new(Index::new(meta_path, layout)); 124 + index.rebuild().unwrap(); 125 + index.warm_collaborators(); 126 + index 127 + } 128 + 129 + fn atproto_with(responder: Responder) -> Arc<Atproto<FakeHttp<Responder>, ManualClock>> { 130 + Arc::new(Atproto::new( 131 + FakeHttp::new(responder), 132 + ManualClock::new(UnixMicros::new(1_000_000_000)), 133 + KnotId::new(KNOT_DID).unwrap(), 134 + knot_atproto::PlcDirectory::new(Url::parse("https://plc.directory/").unwrap()).unwrap(), 135 + )) 136 + } 137 + 138 + fn atproto_answering( 139 + status: http::StatusCode, 140 + calls: Arc<AtomicUsize>, 141 + ) -> Arc<Atproto<FakeHttp<Responder>, ManualClock>> { 142 + atproto_with(Box::new(move |_request: &HttpRequest| { 143 + calls.fetch_add(1, Ordering::SeqCst); 144 + Ok(responding(status, bytes::Bytes::new())) 145 + })) 146 + } 147 + 148 + fn did_document(did: &str, handle: &str, pds: &str) -> bytes::Bytes { 149 + let signing = K256Signer::generate(&SeededEntropy::new(5)); 150 + let multibase = knot_types::crypto::multikey(0xe7, signing.public_key().as_bytes()); 151 + let body = serde_json::json!({ 152 + "id": did, 153 + "alsoKnownAs": [format!("at://{handle}")], 154 + "verificationMethod": [{ 155 + "id": format!("{did}#atproto"), 156 + "type": "Multikey", 157 + "controller": did, 158 + "publicKeyMultibase": multibase, 159 + }], 160 + "service": [{ 161 + "id": "#atproto_pds", 162 + "type": "AtprotoPersonalDataServer", 163 + "serviceEndpoint": pds, 164 + }] 165 + }); 166 + bytes::Bytes::from(serde_json::to_vec(&body).unwrap()) 167 + } 168 + 169 + fn atproto_serving_two_accounts() -> Arc<Atproto<FakeHttp<Responder>, ManualClock>> { 170 + atproto_with(Box::new(move |request: &HttpRequest| { 171 + let url = request.url.as_str(); 172 + let body = if url.contains("listRecords") { 173 + bytes::Bytes::from_static(br#"{"records":[]}"#) 174 + } else if url.ends_with(NEL) { 175 + did_document(NEL, "nel.pet", NEL_PDS) 176 + } else if url.ends_with(OLAREN) { 177 + did_document(OLAREN, "olaren.dev", OLAREN_PDS) 178 + } else { 179 + return Ok(responding(http::StatusCode::NOT_FOUND, bytes::Bytes::new())); 180 + }; 181 + Ok(responding(http::StatusCode::OK, body)) 182 + })) 183 + } 184 + 185 + fn atproto_recording( 186 + seen: Arc<Mutex<Vec<String>>>, 187 + ) -> Arc<Atproto<FakeHttp<Responder>, ManualClock>> { 188 + atproto_with(Box::new(move |request: &HttpRequest| { 189 + let url = request.url.as_str(); 190 + if url.contains("listRecords") { 191 + return Ok(responding( 192 + http::StatusCode::OK, 193 + bytes::Bytes::from_static(br#"{"records":[]}"#), 194 + )); 195 + } 196 + if url.ends_with(NEL) { 197 + return Ok(responding( 198 + http::StatusCode::OK, 199 + did_document(NEL, "nel.pet", NEL_PDS), 200 + )); 201 + } 202 + if let Some(did) = url.rsplit('/').next() { 203 + seen.lock().unwrap().push(did.to_string()); 204 + } 205 + Ok(responding( 206 + http::StatusCode::SERVICE_UNAVAILABLE, 207 + bytes::Bytes::new(), 208 + )) 209 + })) 210 + } 211 + 212 + fn now() -> UnixSeconds { 213 + UnixSeconds::new(1_000) 214 + } 215 + 216 + fn brisk() -> Pace { 217 + Pace { 218 + busy: BusyRetry::from_millis(0), 219 + floor: SettleFloor::from_millis(0), 220 + ttl: KeyTtl::from_secs(3_600), 221 + reprieve: KeyReprieve::from_secs(300, 21_600), 222 + sweep: SweepFloor::DEFAULT, 223 + stalled: StalledBackoff::from_secs(1), 224 + settled: SettledPause::from_secs(1), 225 + members: AccountBudget::new(64), 226 + suspected: AccountBudget::new(256), 227 + host: RateLimit { 228 + burst: Burst::new(1), 229 + refill: RefillMicros::new(1_000), 230 + }, 231 + } 232 + } 233 + 234 + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] 235 + async fn a_gone_document_completes_the_set_and_an_outage_leaves_it_warming() { 236 + let outcome = |status: http::StatusCode| async move { 237 + let scratch = tempfile::tempdir().unwrap(); 238 + let index = hosted_index(&scratch, &[anemone()], &[]); 239 + let calls = Arc::new(AtomicUsize::new(0)); 240 + let atproto = atproto_answering(status, Arc::clone(&calls)); 241 + let pacer = HostPacer::new(brisk().host); 242 + assert_eq!(index.keys().coverage(), Coverage::Warming); 243 + fill_once( 244 + &index, 245 + &atproto, 246 + &Slots::testing(4), 247 + &pacer, 248 + brisk(), 249 + &Cursors::default(), 250 + ) 251 + .await; 252 + assert!( 253 + calls.load(Ordering::SeqCst) > 0, 254 + "the fill made an outbound request" 255 + ); 256 + index.keys().coverage() 257 + }; 258 + 259 + assert_eq!( 260 + outcome(http::StatusCode::NOT_FOUND).await, 261 + Coverage::Ready, 262 + "a permanently unresolvable owner is recorded with an empty key set, so the set is complete" 263 + ); 264 + assert_eq!( 265 + outcome(http::StatusCode::SERVICE_UNAVAILABLE).await, 266 + Coverage::Warming, 267 + "a transient failure doesn't teach the knot anything, so it mustn't claim a complete set" 268 + ); 269 + } 270 + 271 + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] 272 + async fn keys_the_knot_already_read_survive_an_outage_and_a_refused_listing() { 273 + let scratch = tempfile::tempdir().unwrap(); 274 + let index = hosted_index(&scratch, &[anemone()], &[]); 275 + let key = OfferedKey::from_bytes(vec![9]); 276 + let spent = KeyTtl::from_secs(1).lease_from(UnixSeconds::new(100)); 277 + index.keys().record(&account(NEL), vec![key.clone()], spent); 278 + let pacer = HostPacer::new(brisk().host); 279 + 280 + let unreachable = atproto_answering(http::StatusCode::SERVICE_UNAVAILABLE, Arc::default()); 281 + fill_once( 282 + &index, 283 + &unreachable, 284 + &Slots::testing(4), 285 + &pacer, 286 + brisk(), 287 + &Cursors::default(), 288 + ) 289 + .await; 290 + assert_eq!( 291 + index.keys().coverage(), 292 + Coverage::Ready, 293 + "an owner the knot has read before keeps its last keys through an outage, so one \ 294 + unreachable PDS mustn't reopen the knot to every offered key" 295 + ); 296 + assert_eq!( 297 + index.owner_of_key(&key, now()), 298 + Resolved::Ready(Some(account(NEL))), 299 + "the reprieve keeps the keys the knot last read" 300 + ); 301 + 302 + index.keys().record(&account(NEL), vec![key.clone()], spent); 303 + let listing = Arc::new(AtomicUsize::new(0)); 304 + let refusing = { 305 + let listing = Arc::clone(&listing); 306 + atproto_with(Box::new(move |request: &HttpRequest| { 307 + let url = request.url.as_str(); 308 + if url.contains("listRecords") { 309 + listing.fetch_add(1, Ordering::SeqCst); 310 + return Ok(responding( 311 + http::StatusCode::BAD_REQUEST, 312 + bytes::Bytes::new(), 313 + )); 314 + } 315 + Ok(responding( 316 + http::StatusCode::OK, 317 + did_document(NEL, "nel.pet", NEL_PDS), 318 + )) 319 + })) 320 + }; 321 + fill_once( 322 + &index, 323 + &refusing, 324 + &Slots::testing(4), 325 + &pacer, 326 + brisk(), 327 + &Cursors::default(), 328 + ) 329 + .await; 330 + assert!( 331 + listing.load(Ordering::SeqCst) > 0, 332 + "the fill read from the PDS" 333 + ); 334 + assert_eq!( 335 + index.owner_of_key(&key, now()), 336 + Resolved::Ready(Some(account(NEL))), 337 + "an account whose DID document resolves hasn't gone anywhere, so the knot mustn't read \ 338 + a 400 from a record listing as proof the account stopped publishing keys" 339 + ); 340 + } 341 + 342 + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] 343 + async fn the_fill_takes_a_turn_at_each_host_it_reads_from_and_stops_on_shutdown() { 344 + let scratch = tempfile::tempdir().unwrap(); 345 + let index = hosted_index(&scratch, &[anemone(), barnacle()], &[]); 346 + let atproto = atproto_serving_two_accounts(); 347 + let pacer = HostPacer::new(brisk().host); 348 + 349 + fill_once( 350 + &index, 351 + &atproto, 352 + &Slots::testing(4), 353 + &pacer, 354 + brisk(), 355 + &Cursors::default(), 356 + ) 357 + .await; 358 + 359 + assert_eq!( 360 + index.keys().coverage(), 361 + Coverage::Ready, 362 + "both owners resolved, so every account that may push has a record" 363 + ); 364 + ["plc.directory", "pds.nel.pet", "pds.olaren.dev"] 365 + .iter() 366 + .for_each(|host| { 367 + assert!( 368 + !pacer.reserve_now(&HostKey::new(host), UnixMicros::new(0)), 369 + "the fill must take a turn at {host} before it reads from it, \ 370 + or a knot whose members share one PDS spends its whole rate at that host" 371 + ); 372 + }); 373 + assert!( 374 + pacer.reserve_now(&HostKey::new("pds.teq.dev"), UnixMicros::new(0)), 375 + "a host the fill never read from is due immediately. The bookings above are the fill's \ 376 + own work" 377 + ); 378 + 379 + let shutdown = CancellationToken::new(); 380 + let task = knot_keyfill::spawn( 381 + Arc::clone(&index), 382 + Arc::clone(&atproto), 383 + Slots::testing(4), 384 + brisk(), 385 + shutdown.clone(), 386 + ); 387 + shutdown.cancel(); 388 + tokio::time::timeout(std::time::Duration::from_secs(5), task) 389 + .await 390 + .expect("a shutting-down knot mustn't wait out the pause between fill passes") 391 + .expect("the fill task stops without panicking"); 392 + } 393 + 394 + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] 395 + async fn members_are_read_whole_on_first_contact_then_renewed_one_budget_turn_per_pass() { 396 + let scratch = tempfile::tempdir().unwrap(); 397 + let members = [account(OLAREN), account(TEQ), account(BAILEY)]; 398 + let index = hosted_index(&scratch, &[anemone()], &members); 399 + let seen = Arc::new(Mutex::new(Vec::new())); 400 + let atproto = atproto_recording(Arc::clone(&seen)); 401 + let pacer = HostPacer::new(brisk().host); 402 + let pace = Pace { 403 + members: AccountBudget::new(1), 404 + ..brisk() 405 + }; 406 + let cursor = Cursors::default(); 407 + 408 + fill_once(&index, &atproto, &Slots::testing(4), &pacer, pace, &cursor).await; 409 + assert_eq!( 410 + index.keys().coverage(), 411 + Coverage::Ready, 412 + "the pushers are what coverage waits on, so an unreadable member mustn't make the knot \ 413 + doubt the keys it checks pushes against" 414 + ); 415 + let mut attempted = seen.lock().unwrap().clone(); 416 + attempted.sort(); 417 + assert_eq!( 418 + attempted, 419 + vec![BAILEY.to_string(), OLAREN.to_string(), TEQ.to_string()], 420 + "the renewal budget paces rereads, so a member the knot has never read mustn't wait \ 421 + its turn behind it and be refused at the handshake for the passes in between" 422 + ); 423 + 424 + let spent = KeyTtl::from_secs(1).lease_from(UnixSeconds::new(0)); 425 + members 426 + .iter() 427 + .for_each(|member| _ = index.keys().record(member, Vec::new(), spent)); 428 + seen.lock().unwrap().clear(); 429 + 430 + futures::stream::iter(0..3) 431 + .for_each(|_| async { 432 + fill_once(&index, &atproto, &Slots::testing(4), &pacer, pace, &cursor).await; 433 + }) 434 + .await; 435 + let order = seen.lock().unwrap().clone(); 436 + assert_eq!( 437 + order.iter().map(String::as_str).collect::<Vec<_>>(), 438 + vec![BAILEY, OLAREN, TEQ], 439 + "one member per pass in turn, or members whose PDS stays down keep the front of \ 440 + the queue and the knot never reads the rest" 441 + ); 442 + }

Configure Feed

Configure Feed