Something went wrong. Try again.
lightweight `com.atproto.sync.listReposByCollection`
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430//! Bounded, host-fair in-memory queue for discovery-sourced resync work.//!//! Backfill and deep-crawl tasks push newly-discovered repos here instead of//! the on-disk resync queue. The dispatcher pops items round-robin across PDS//! hosts, achieving fair throughput distribution without timestamp stagger//! arithmetic.//!//! The queue is bounded in two dimensions://! - **Global capacity**: total items across all hosts; prevents unbounded//! memory growth.//! - **Per-host limit**: items for a single PDS; prevents one large PDS from//! monopolising the queue and starving other hosts.//!//! When either limit is reached, producers block on [`DiscoveryQueue::push`]//! until the dispatcher drains items, applying natural backpressure to crawl//! tasks.//!//! On process restart, queued items are lost. This is acceptable because//! backfill/deep-crawl cursors are persisted — the gap self-heals on the next//! crawl pass.
use std::collections::{HashMap, HashSet, VecDeque};use std::pin::pin;use std::sync::atomic::{AtomicUsize, Ordering};use std::sync::{Arc, Mutex};
use jacquard_common::types::crypto::PublicKey;use jacquard_common::types::string::Did;use jacquard_common::url::Url;
/// A newly-discovered repo awaiting its initial resync.////// Carries the resolved identity (PDS + pubkey) from backfill so the/// dispatcher can warm the identity cache without a redundant network call.pub struct DiscoveryItem { pub did: Did<'static>, pub pds: Arc<Url>, pub pubkey: PublicKey<'static>,}
pub struct DiscoveryQueue { inner: Mutex<Inner>, /// Wakes producers when space becomes available (after a pop). space_available: tokio::sync::Notify, /// Wakes the dispatcher when items arrive. item_available: tokio::sync::Notify, /// Maximum total items across all host queues. capacity: usize, /// Maximum items for a single PDS host. per_host_limit: usize, /// Total items, readable without the lock (kept in sync under the lock). len: AtomicUsize,}
struct Inner { /// Ordered list of PDS hosts for round-robin traversal. hosts: Vec<Arc<Url>>, /// Per-host FIFO queues of DIDs awaiting resync. queues: HashMap<Arc<Url>, VecDeque<(Did<'static>, PublicKey<'static>)>>, /// Round-robin cursor: index into `hosts` for the next pop. cursor: usize, /// Total items across all queues. len: usize,}
impl Inner { fn can_push(&self, pds: &Arc<Url>, capacity: usize, per_host_limit: usize) -> bool { if self.len >= capacity { return false; } let host_len = self.queues.get(pds).map_or(0, |q| q.len()); host_len < per_host_limit }
fn insert(&mut self, item: DiscoveryItem) { if !self.queues.contains_key(&item.pds) { self.hosts.push(item.pds.clone()); self.queues.insert(item.pds.clone(), VecDeque::new()); } self.queues .get_mut(&item.pds) .unwrap() .push_back((item.did, item.pubkey)); self.len += 1; }}
impl DiscoveryQueue { pub fn new(capacity: usize, per_host_limit: usize) -> Self { Self { inner: Mutex::new(Inner { hosts: Vec::new(), queues: HashMap::new(), cursor: 0, len: 0, }), space_available: tokio::sync::Notify::new(), item_available: tokio::sync::Notify::new(), capacity, per_host_limit, len: AtomicUsize::new(0), } }
/// Push an item, blocking if the global capacity or per-host limit is /// reached. /// /// Callers should `tokio::select!` this against their cancellation token /// for clean shutdown. pub async fn push(&self, item: DiscoveryItem) { loop { // Register for wake-up BEFORE checking, so a pop between check // and await still wakes us. let mut notified = pin!(self.space_available.notified()); notified.as_mut().enable();
{ let mut inner = self.inner.lock().unwrap(); if inner.can_push(&item.pds, self.capacity, self.per_host_limit) { inner.insert(item); self.len.store(inner.len, Ordering::Relaxed); drop(inner); self.item_available.notify_one(); return; } }
notified.await; } }
/// Pop up to `n` items round-robin across hosts. /// /// Hosts in `skip_hosts` are left in the queue (e.g. cooling after 429). /// Returns immediately with whatever is available (may be empty). pub fn pop_batch(&self, n: usize, skip_hosts: &HashSet<Arc<Url>>) -> Vec<DiscoveryItem> { let mut inner = self.inner.lock().unwrap(); let mut result = Vec::with_capacity(n);
if inner.hosts.is_empty() || n == 0 { return result; }
let mut consecutive_skips = 0; while result.len() < n && consecutive_skips <= inner.hosts.len() { if inner.hosts.is_empty() { break; } let idx = inner.cursor % inner.hosts.len(); let host = inner.hosts[idx].clone();
if skip_hosts.contains(&host) { inner.cursor = idx + 1; consecutive_skips += 1; continue; }
let queue = inner .queues .get_mut(&host) .expect("host in ring must have a queue"); let (did, pubkey) = queue.pop_front().expect("host in ring must have items"); let host_empty = queue.is_empty(); inner.len -= 1; consecutive_skips = 0;
if host_empty { inner.hosts.remove(idx); inner.queues.remove(&host); // Don't advance cursor — the next host slid into this index. if !inner.hosts.is_empty() { inner.cursor = idx % inner.hosts.len(); } } else { inner.cursor = idx + 1; }
result.push(DiscoveryItem { did, pds: host, pubkey, }); }
let popped = result.len(); self.len.store(inner.len, Ordering::Relaxed); drop(inner);
if popped > 0 { self.space_available.notify_waiters(); } result }
/// Number of items across all host queues. pub fn len(&self) -> usize { self.len.load(Ordering::Relaxed) }
pub fn is_empty(&self) -> bool { self.len() == 0 }
/// Returns a future that completes when an item is pushed. /// /// Use in `tokio::select!` alongside the idle-poll timer so the dispatcher /// wakes immediately when discovery work arrives instead of waiting for /// the next poll cycle. pub fn notified(&self) -> tokio::sync::futures::Notified<'_> { self.item_available.notified() }}
#[cfg(test)]mod tests { use super::*;
fn url(s: &str) -> Arc<Url> { Arc::new(s.parse().unwrap()) }
fn did(s: &str) -> Did<'static> { Did::new_owned(s).unwrap() }
fn item(did_str: &str, pds_str: &str) -> DiscoveryItem { use jacquard_common::types::crypto::KeyCodec; DiscoveryItem { did: did(did_str), pds: url(pds_str), pubkey: PublicKey { codec: KeyCodec::Ed25519, bytes: std::borrow::Cow::Owned(vec![0u8; 32]), }, } }
#[tokio::test] async fn push_and_pop_single() { let q = DiscoveryQueue::new(10, 10); q.push(item("did:web:alice.test", "https://pds1.test")) .await; assert_eq!(q.len(), 1);
let batch = q.pop_batch(10, &HashSet::new()); assert_eq!(batch.len(), 1); assert_eq!(batch[0].did.as_str(), "did:web:alice.test"); assert_eq!(q.len(), 0); }
#[tokio::test] async fn round_robin_across_hosts() { let q = DiscoveryQueue::new(100, 100); // Two items on pds1, two on pds2. q.push(item("did:web:a1.test", "https://pds1.test")).await; q.push(item("did:web:a2.test", "https://pds1.test")).await; q.push(item("did:web:b1.test", "https://pds2.test")).await; q.push(item("did:web:b2.test", "https://pds2.test")).await;
let batch = q.pop_batch(4, &HashSet::new()); assert_eq!(batch.len(), 4); // Should alternate: pds1, pds2, pds1, pds2. assert_eq!(batch[0].did.as_str(), "did:web:a1.test"); assert_eq!(batch[1].did.as_str(), "did:web:b1.test"); assert_eq!(batch[2].did.as_str(), "did:web:a2.test"); assert_eq!(batch[3].did.as_str(), "did:web:b2.test"); }
#[tokio::test] async fn skip_hosts_leaves_items_in_queue() { let q = DiscoveryQueue::new(100, 100); let pds1 = url("https://pds1.test"); q.push(item("did:web:a1.test", "https://pds1.test")).await; q.push(item("did:web:b1.test", "https://pds2.test")).await;
let skip = HashSet::from([pds1]); let batch = q.pop_batch(10, &skip); assert_eq!(batch.len(), 1); assert_eq!(batch[0].did.as_str(), "did:web:b1.test"); // pds1 item is still queued. assert_eq!(q.len(), 1); }
#[tokio::test] async fn pop_batch_respects_limit() { let q = DiscoveryQueue::new(100, 100); for i in 0..10 { q.push(item(&format!("did:web:user{i}.test"), "https://pds1.test")) .await; }
let batch = q.pop_batch(3, &HashSet::new()); assert_eq!(batch.len(), 3); assert_eq!(q.len(), 7); }
#[tokio::test] async fn empty_pop_returns_empty() { let q = DiscoveryQueue::new(10, 10); let batch = q.pop_batch(10, &HashSet::new()); assert!(batch.is_empty()); }
/// Helper: yield to the runtime repeatedly so a spawned task has a /// chance to make progress. If a push would complete, it will within /// a few yields; if it's genuinely blocked, it stays blocked. async fn yield_to_runtime() { for _ in 0..5 { tokio::task::yield_now().await; } }
#[tokio::test] async fn global_backpressure_blocks_producer() { let q = Arc::new(DiscoveryQueue::new(2, 10)); q.push(item("did:web:a.test", "https://pds1.test")).await; q.push(item("did:web:b.test", "https://pds2.test")).await; assert_eq!(q.len(), 2);
// Third push should block (global capacity 2). let q2 = q.clone(); let handle = tokio::spawn(async move { q2.push(item("did:web:c.test", "https://pds3.test")).await; });
yield_to_runtime().await; assert!(!handle.is_finished());
// Pop one to free capacity — push should complete. q.pop_batch(1, &HashSet::new()); yield_to_runtime().await; assert!(handle.is_finished()); handle.await.expect("push task shouldn't panic"); assert_eq!(q.len(), 2); }
#[tokio::test] async fn per_host_limit_blocks_producer() { let q = Arc::new(DiscoveryQueue::new(100, 2)); // Fill host pds1 to its per-host limit of 2. q.push(item("did:web:a.test", "https://pds1.test")).await; q.push(item("did:web:b.test", "https://pds1.test")).await;
// Third push to same host should block despite global capacity. let q2 = q.clone(); let handle = tokio::spawn(async move { q2.push(item("did:web:c.test", "https://pds1.test")).await; });
yield_to_runtime().await; assert!(!handle.is_finished());
// A different host should still be pushable. q.push(item("did:web:d.test", "https://pds2.test")).await; assert_eq!(q.len(), 3);
// Pop one from pds1 to unblock the waiting producer. let batch = q.pop_batch(1, &HashSet::new()); assert_eq!(batch[0].pds.as_str(), "https://pds1.test/"); yield_to_runtime().await; assert!(handle.is_finished()); handle.await.expect("push task shouldn't panic"); assert_eq!(q.len(), 3); }
#[tokio::test] async fn per_host_limit_prevents_monopolisation() { let q = DiscoveryQueue::new(100, 3); // Push 3 items for pds1 (at limit), then 3 for pds2. for i in 0..3 { q.push(item(&format!("did:web:a{i}.test"), "https://pds1.test")) .await; } for i in 0..3 { q.push(item(&format!("did:web:b{i}.test"), "https://pds2.test")) .await; } assert_eq!(q.len(), 6);
// Pop 6: should interleave thanks to round-robin. let batch = q.pop_batch(6, &HashSet::new()); assert_eq!(batch.len(), 6); assert_eq!(batch[0].pds.as_str(), "https://pds1.test/"); assert_eq!(batch[1].pds.as_str(), "https://pds2.test/"); assert_eq!(batch[2].pds.as_str(), "https://pds1.test/"); assert_eq!(batch[3].pds.as_str(), "https://pds2.test/"); }
#[tokio::test] async fn cursor_persists_across_batches() { let q = DiscoveryQueue::new(100, 100); // 3 hosts, 1 item each. q.push(item("did:web:a.test", "https://pds1.test")).await; q.push(item("did:web:b.test", "https://pds2.test")).await; q.push(item("did:web:c.test", "https://pds3.test")).await; // Add second item to each. q.push(item("did:web:a2.test", "https://pds1.test")).await; q.push(item("did:web:b2.test", "https://pds2.test")).await; q.push(item("did:web:c2.test", "https://pds3.test")).await;
// Pop 2: should get pds1, pds2. Cursor now at pds3. let batch1 = q.pop_batch(2, &HashSet::new()); assert_eq!(batch1[0].did.as_str(), "did:web:a.test"); assert_eq!(batch1[1].did.as_str(), "did:web:b.test");
// Pop 2 more: cursor resumes at pds3, then wraps to pds1. let batch2 = q.pop_batch(2, &HashSet::new()); assert_eq!(batch2[0].did.as_str(), "did:web:c.test"); assert_eq!(batch2[1].did.as_str(), "did:web:a2.test"); }
#[tokio::test] async fn host_removed_when_drained() { let q = DiscoveryQueue::new(100, 100); q.push(item("did:web:a.test", "https://pds1.test")).await; q.push(item("did:web:b.test", "https://pds2.test")).await;
// Drain pds1 by skipping pds2. let pds2 = url("https://pds2.test"); let batch = q.pop_batch(10, &HashSet::from([pds2])); assert_eq!(batch.len(), 1);
// Now un-skip pds2 and pop. let batch = q.pop_batch(10, &HashSet::new()); assert_eq!(batch.len(), 1); assert_eq!(batch[0].did.as_str(), "did:web:b.test"); assert!(q.is_empty()); }}