alice.pds.demo.boltless.dev
This repository has no description
did:plc:3fwecdnvtcscjnrx2p4n7a…
1use std::future::Future;
2use std::sync::Arc;
3use std::sync::atomic::{AtomicUsize, Ordering};
4use std::time::Duration;
5
6use futures::future::OptionFuture;
7use futures::{FutureExt, StreamExt};
8use knot_atproto::{Atproto, AtprotoError};
9use knot_index::{
10 Coverage, HostedCoverage, Index, IndexGeneration, KeyLease, KeyRecord, KeyReprieve,
11 KeyReprieved, KeyTtl, MemberWork, Pushers, Resolved, StalePushers, SuspectPushers, SweepFloor,
12};
13use knot_resource::{Burst, HostKey, HostPacer, RateLimit, RefillMicros, SlotPermit, Slots};
14use knot_runtime::{Clock, HttpTransport};
15use knot_types::{AccountDid, OfferedKey, UnixMicros, UnixSeconds};
16use tokio::sync::watch;
17use tokio_util::sync::CancellationToken;
18use url::Url;
19
20const FILL_FANOUT: usize = 8;
21
22const PASS_HEADROOM: u32 = 2;
23
24macro_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
41span!(
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)]
49pub struct AccountBudget(usize);
50
51impl 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)]
62pub struct Cursor(AtomicUsize);
63
64impl 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)]
79pub struct Cursors {
80 members: Cursor,
81 suspected: Cursor,
82}
83
84fn 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)]
97pub 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
110impl 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
130impl 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
138pub 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
168struct 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
179impl<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
201async 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
211async 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
221struct 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
232pub 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
295struct PushWork {
296 generation: IndexGeneration,
297 hosted: HostedCoverage,
298 pushers: Pushers,
299 due: StalePushers,
300 suspected: SuspectPushers,
301 complete: bool,
302}
303
304async 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
316async 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
348async 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
367fn 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
386async 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
410async 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
434fn 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
459fn 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
481async 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
498async 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
504async 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)]
522mod 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}