use crate::p2p::{score::Score, wire};
use ahash::{AHashMap, AHashSet};
use commonware_actor::{Feedback, Unreliable};
use commonware_cryptography::PublicKey;
use commonware_p2p::{Recipients, Sender, utils::codec::WrappedSender};
use commonware_runtime::{
Clock, Metrics,
telemetry::metrics::{
EncodeStruct, GaugeExt, GaugeFamily, Histogram, MetricsExt as _,
histogram::Buckets,
status::{self, Status},
},
};
use commonware_utils::{
NonZeroDuration, PrioritySet, Span, SystemTimeExt,
sampling::{WeightedIndex, WithoutReplacement},
};
use rand::RngExt as _;
use rand_core::Rng;
use std::{
collections::hash_map::Entry,
mem,
time::{Duration, SystemTime},
};
use tracing::debug;
#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeStruct)]
struct Peer<P: PublicKey> {
peer: P,
}
pub type ID = u64;
struct ActiveRequest<P, Key> {
key: Key,
peer: P,
start: SystemTime,
}
struct Participants<P> {
peers: Vec<(P, Score)>,
weights: WeightedIndex,
}
impl<P: PublicKey> Participants<P> {
fn index(&self, peer: &P) -> Option<usize> {
self.peers.binary_search_by(|(p, _)| p.cmp(peer)).ok()
}
fn eligible_subset(&self, peers: &AHashSet<P>) -> Vec<usize> {
let lookup_steps = self.peers.len().checked_ilog2().unwrap_or(0) as usize + 1;
if peers.len().saturating_mul(lookup_steps) < self.peers.len() {
let mut indices: Vec<_> = peers
.iter()
.filter_map(|peer| self.index(peer))
.filter(|&index| self.weights.weight(index) != 0)
.collect();
indices.sort_unstable();
indices
} else {
self.peers
.iter()
.enumerate()
.filter(|(index, (peer, _))| {
self.weights.weight(*index) != 0 && peers.contains(peer)
})
.map(|(index, _)| index)
.collect()
}
}
}
enum Candidates {
All,
Targets {
peers: Vec<usize>,
weights: WeightedIndex,
},
}
impl Candidates {
fn select(self, shared: &mut WeightedIndex) -> Selection<'_> {
let (peers, draws) = match self {
Self::All => (None, shared.without_replacement()),
Self::Targets { peers, weights } => (Some(peers), weights.into_without_replacement()),
};
Selection { peers, draws }
}
}
struct Selection<'a> {
peers: Option<Vec<usize>>,
draws: WithoutReplacement<'a>,
}
impl Selection<'_> {
fn next<R: Rng>(&mut self, rng: &mut R, retries: u32) -> Option<usize> {
if self.draws.is_empty() {
return None;
}
let index = if rng.random_range(0..u64::from(retries) + 4) < 4 {
self.draws.sample(rng)?
} else {
self.draws.sample_uniform(rng)?
};
Some(self.peers.as_ref().map_or(index, |peers| peers[index]))
}
}
pub struct Config<P: PublicKey> {
pub me: Option<P>,
pub fetch_timeout: NonZeroDuration,
pub retry_delay: Duration,
pub priority_requests: bool,
}
pub struct Fetcher<E, P, Key>
where
E: Clock + Rng + Metrics,
P: PublicKey,
Key: Span,
{
context: E,
me: Option<P>,
blocked: AHashSet<P>,
participants: Participants<P>,
request_id: Option<ID>,
active: PrioritySet<ID, SystemTime>,
requests: AHashMap<ID, ActiveRequest<P, Key>>,
fetch_timeout: NonZeroDuration,
pending: PrioritySet<Key, (bool, SystemTime)>,
waiter: Option<SystemTime>,
retry_delay: Duration,
priority_requests: bool,
targets: AHashMap<Key, AHashSet<P>>,
retries: AHashMap<Key, u32>,
performance: GaugeFamily<Peer<P>>,
requests_created: status::Counter,
requests_sent: status::Counter,
resolves: Histogram,
}
impl<E, P, Key> Fetcher<E, P, Key>
where
E: Clock + Rng + Metrics,
P: PublicKey,
Key: Span,
{
pub fn new(context: E, config: Config<P>) -> Self {
let performance = context.family(
"peer_performance",
"Per-peer throughput score expressed as a sampling weight",
);
let requests_created =
context.family("requests_created", "Status of request creation attempts");
let requests_sent = context.family(
"requests_sent",
"Status of individual network requests sent to peers",
);
let resolves = context.histogram(
"resolves",
"Number and duration of useful responses",
Buckets::NETWORK,
);
Self {
context,
me: config.me,
blocked: AHashSet::new(),
participants: Participants {
peers: Vec::new(),
weights: WeightedIndex::new([]),
},
request_id: None,
active: PrioritySet::new(),
requests: AHashMap::new(),
fetch_timeout: config.fetch_timeout,
pending: PrioritySet::new(),
waiter: None,
retry_delay: config.retry_delay,
priority_requests: config.priority_requests,
targets: AHashMap::new(),
retries: AHashMap::new(),
performance,
requests_created,
requests_sent,
resolves,
}
}
fn update_performance(&mut self, participant: &P, elapsed: Duration, bytes: usize) {
let Some(index) = self.participants.index(participant) else {
return;
};
let score = &mut self.participants.peers[index].1;
score.record(elapsed, bytes);
let next = score.weight();
self.participants
.weights
.update(index, |weight| if weight == 0 { 0 } else { next });
let _ = self.performance.get_or_create_by(participant).try_set(next);
}
fn candidates(&self, key: &Key) -> Option<(Candidates, u32)> {
let candidates = if let Some(targets) = self.targets.get(key) {
let peers = self.participants.eligible_subset(targets);
if peers.is_empty() {
return None;
}
let weights = WeightedIndex::new(
peers
.iter()
.map(|&index| self.participants.weights.weight(index)),
);
Candidates::Targets { peers, weights }
} else {
if self.participants.weights.total() == 0 {
return None;
}
Candidates::All
};
let retries = self.retries.get(key).copied().unwrap_or(0);
Some((candidates, retries))
}
#[cfg(test)]
fn get_eligible_peers(&mut self, key: &Key) -> Vec<P> {
let Some((candidates, retries)) = self.candidates(key) else {
return Vec::new();
};
let mut peers = Vec::new();
let mut selection = candidates.select(&mut self.participants.weights);
while let Some(index) = selection.next(&mut self.context, retries) {
peers.push(self.participants.peers[index].0.clone());
}
peers
}
pub fn fetch<NetS: Sender<PublicKey = P>>(
&mut self,
sender: &mut WrappedSender<NetS, wire::Message<Key>>,
) {
self.waiter = None;
let now = self.context.current();
let mut next_deadline = SystemTime::limit();
let mut earliest_rate_limit: Option<SystemTime> = None;
let mut found_eligible_peers = false;
let pending = mem::replace(&mut self.pending, PrioritySet::new());
let mut sent = None;
for (key, (_, deadline)) in pending.iter() {
if *deadline > now {
next_deadline = next_deadline.min(*deadline);
continue;
}
let Some((candidates, retries)) = self.candidates(key) else {
self.requests_created.inc(Status::Dropped);
continue;
};
self.requests_created.inc(Status::Success);
found_eligible_peers = true;
let mut selection = candidates.select(&mut self.participants.weights);
while let Some(index) = selection.next(&mut self.context, retries) {
let peer = self.participants.peers[index].0.clone();
let checked = match sender.check(Recipients::One(peer.clone())) {
Ok(checked) => checked,
Err(not_until) => {
earliest_rate_limit =
Some(earliest_rate_limit.map_or(not_until, |t| t.min(not_until)));
continue;
}
};
let id = self.request_id.unwrap_or_else(|| self.context.next_u64());
self.request_id = Some(id.wrapping_add(1));
let message = wire::Message {
id,
payload: wire::Payload::Request(key.clone()),
};
match checked.send(message, self.priority_requests) {
Unreliable::Outcome(Feedback::Ok | Feedback::Backoff) => {
self.requests_sent.inc(Status::Success);
let now = self.context.current();
sent = Some((key.clone(), id, peer, now));
break;
}
feedback @ (Unreliable::Rejected | Unreliable::Outcome(Feedback::Closed)) => {
self.requests_sent.inc(Status::Dropped);
debug!(?peer, ?feedback, "send failed");
}
}
}
if sent.is_some() {
break;
}
}
self.pending = pending;
if let Some((key, id, peer, start)) = sent {
assert!(self.pending.remove(&key));
let deadline = start
.checked_add(self.fetch_timeout.get())
.expect("time overflowed");
self.active.put(id, deadline);
self.requests.insert(id, ActiveRequest { key, peer, start });
return;
}
if let Some(rate_limit_time) = earliest_rate_limit {
next_deadline = next_deadline.min(rate_limit_time);
} else if found_eligible_peers {
next_deadline = next_deadline.min(self.context.current() + self.retry_delay);
}
self.waiter = Some(next_deadline);
}
pub fn retain(&mut self, predicate: impl Fn(&Key) -> bool) {
for (id, _) in self.requests.extract_if(|_, req| !predicate(&req.key)) {
self.active.remove(&id);
}
self.pending.retain(&predicate);
self.targets.retain(|k, _| predicate(k));
self.retries.retain(|k, _| predicate(k));
self.waiter = None;
}
pub fn add_ready(&mut self, key: Key) {
assert!(!self.pending.contains(&key));
self.waiter = None;
self.pending.put(key, (false, self.context.current()));
}
pub fn add_retry(&mut self, key: Key) {
assert!(!self.pending.contains(&key));
let retries = self.retries.entry(key.clone()).or_default();
*retries = retries.saturating_add(1);
self.waiter = None;
let deadline = self.context.current() + self.retry_delay;
self.pending.put(key, (true, deadline));
}
pub fn get_pending_deadline(&self) -> Option<SystemTime> {
if self.pending.is_empty() {
return None;
}
let pending_deadline = self.pending.peek().map(|(_, (_, deadline))| *deadline);
pending_deadline.max(self.waiter)
}
pub fn get_active_deadline(&self) -> Option<SystemTime> {
self.active.peek().map(|(_, deadline)| *deadline)
}
pub fn pop_active(&mut self) -> Option<Key> {
let (id, _) = self.active.pop()?;
let req = self.requests.remove(&id)?;
self.update_performance(
&req.peer,
self.fetch_timeout.get().saturating_add(self.retry_delay),
0,
);
Some(req.key)
}
pub fn pop_response(&mut self, id: ID, peer: &P) -> Option<(Key, Duration)> {
let req = match self.requests.entry(id) {
Entry::Occupied(entry) if &entry.get().peer == peer => entry.remove(),
_ => return None,
};
self.active.remove(&id);
let elapsed = self
.context
.current()
.duration_since(req.start)
.unwrap_or_default();
Some((req.key, elapsed))
}
pub fn record_response(&mut self, peer: &P, elapsed: Duration, bytes: usize) {
self.update_performance(peer, elapsed, bytes);
self.resolves.observe(elapsed.as_secs_f64());
}
pub fn pop_missing(&mut self, id: ID, peer: &P) -> Option<Key> {
let (key, elapsed) = self.pop_response(id, peer)?;
self.update_performance(peer, elapsed.saturating_add(self.retry_delay), 0);
Some(key)
}
pub fn reconcile(&mut self, peers: &[P]) {
let mut peers = peers.to_vec();
peers.sort_unstable();
peers.dedup();
let peers: Vec<_> = peers
.into_iter()
.map(|peer| {
let score = self
.participants
.index(&peer)
.map_or_else(Score::default, |index| self.participants.peers[index].1);
(peer, score)
})
.collect();
let weights = WeightedIndex::new(peers.iter().map(|(peer, score)| {
if self.me.as_ref() == Some(peer) || self.blocked.contains(peer) {
0
} else {
score.weight()
}
}));
self.participants = Participants { peers, weights };
self.waiter = None;
}
pub fn set_blocked(&mut self, blocked: impl IntoIterator<Item = P>) {
let blocked: AHashSet<P> = blocked.into_iter().collect();
for peer in self.blocked.symmetric_difference(&blocked) {
if self.me.as_ref() == Some(peer) {
continue;
}
if let Some(index) = self.participants.index(peer) {
let weight = if blocked.contains(peer) {
0
} else {
self.participants.peers[index].1.weight()
};
self.participants.weights.set(index, weight);
}
}
self.blocked = blocked;
self.waiter = None;
}
pub fn add_targets(&mut self, key: Key, peers: impl IntoIterator<Item = P>) {
self.targets.entry(key).or_default().extend(peers);
self.waiter = None;
}
pub fn clear_targets(&mut self, key: &Key) {
self.targets.remove(key);
self.waiter = None;
}
pub fn clear(&mut self, key: &Key) {
self.targets.remove(key);
self.retries.remove(key);
}
pub fn has_targets(&self, key: &Key) -> bool {
self.targets.contains_key(key)
}
#[cfg(test)]
pub fn len(&self) -> usize {
self.pending.len() + self.requests.len()
}
pub fn len_pending(&self) -> usize {
self.pending.len()
}
pub fn len_active(&self) -> usize {
self.requests.len()
}
#[cfg(test)]
pub fn contains(&self, key: &Key) -> bool {
self.requests.values().any(|req| &req.key == key) || self.pending.contains(key)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::p2p::mocks::Key as MockKey;
use commonware_actor::Unreliable;
use commonware_cryptography::{
Signer,
ed25519::{PrivateKey, PublicKey},
};
use commonware_p2p::{LimitedSender, Recipients, UnlimitedSender};
use commonware_runtime::{
IoBufs, KeyedRateLimiter, Quota, Runner as _, Supervisor as _,
deterministic::{self, Context, Runner},
};
use commonware_utils::{NZDuration, NZU32, sync::RwLock};
use std::{slice, sync::Arc, time::Duration};
#[derive(Debug)]
struct CheckedSender<'a, S: UnlimitedSender> {
sender: &'a mut S,
recipients: Recipients<S::PublicKey>,
}
impl<'a, S: UnlimitedSender> commonware_p2p::CheckedSender for CheckedSender<'a, S> {
type PublicKey = S::PublicKey;
fn recipients(&self) -> Vec<Self::PublicKey> {
match &self.recipients {
Recipients::All => Vec::new(),
Recipients::Some(peers) => peers.clone(),
Recipients::One(peer) => vec![peer.clone()],
}
}
fn send(self, message: impl Into<IoBufs> + Send, priority: bool) -> Unreliable<Feedback> {
self.sender.send(self.recipients, message, priority)
}
}
#[derive(Default, Clone, Debug)]
struct FailMockSenderInner;
impl UnlimitedSender for FailMockSenderInner {
type PublicKey = PublicKey;
fn send(
&mut self,
_recipients: Recipients<Self::PublicKey>,
_message: impl Into<IoBufs> + Send,
_priority: bool,
) -> Unreliable<Feedback> {
Unreliable::Rejected
}
}
#[derive(Default, Clone, Debug)]
struct FailMockSender(FailMockSenderInner);
impl LimitedSender for FailMockSender {
type PublicKey = PublicKey;
type Checked<'a> = CheckedSender<'a, FailMockSenderInner>;
fn check(
&mut self,
recipients: Recipients<Self::PublicKey>,
) -> Result<Self::Checked<'_>, SystemTime> {
Ok(CheckedSender {
sender: &mut self.0,
recipients,
})
}
}
#[derive(Default, Clone, Debug)]
struct SuccessMockSenderInner;
impl UnlimitedSender for SuccessMockSenderInner {
type PublicKey = PublicKey;
fn send(
&mut self,
recipients: Recipients<Self::PublicKey>,
_message: impl Into<IoBufs> + Send,
_priority: bool,
) -> Unreliable<Feedback> {
match recipients {
Recipients::One(_) => Unreliable::new(Feedback::Ok),
_ => unimplemented!(),
}
}
}
#[derive(Default, Clone, Debug)]
struct SuccessMockSender(SuccessMockSenderInner);
impl LimitedSender for SuccessMockSender {
type PublicKey = PublicKey;
type Checked<'a> = CheckedSender<'a, SuccessMockSenderInner>;
fn check(
&mut self,
recipients: Recipients<Self::PublicKey>,
) -> Result<Self::Checked<'_>, SystemTime> {
Ok(CheckedSender {
sender: &mut self.0,
recipients,
})
}
}
struct LimitedMockSender<E: Clock, S> {
inner: S,
rate_limiter: Arc<RwLock<KeyedRateLimiter<PublicKey, E>>>,
}
impl<E: Clock, S: Clone> Clone for LimitedMockSender<E, S> {
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
rate_limiter: self.rate_limiter.clone(),
}
}
}
impl<E: Clock, S> LimitedMockSender<E, S> {
fn new(inner: S, quota: Quota, clock: E) -> Self {
Self {
inner,
rate_limiter: Arc::new(RwLock::new(KeyedRateLimiter::hashmap_with_clock(
quota, clock,
))),
}
}
}
impl<E: Clock, S: UnlimitedSender<PublicKey = PublicKey>> LimitedSender
for LimitedMockSender<E, S>
{
type PublicKey = PublicKey;
type Checked<'a> = CheckedSender<'a, S>;
fn check(
&mut self,
recipients: Recipients<Self::PublicKey>,
) -> Result<Self::Checked<'_>, SystemTime> {
let peer = match &recipients {
Recipients::One(p) => p,
_ => unimplemented!(),
};
{
let rate_limiter = self.rate_limiter.write();
if let Err(not_until) = rate_limiter.check_key(peer) {
return Err(not_until.earliest_possible());
}
}
Ok(CheckedSender {
sender: &mut self.inner,
recipients,
})
}
}
fn create_test_fetcher(context: Context) -> Fetcher<Context, PublicKey, MockKey> {
let public_key = PrivateKey::from_seed(0).public_key();
let config = Config {
me: Some(public_key),
fetch_timeout: NZDuration!(Duration::from_secs(5)),
retry_delay: Duration::from_millis(100),
priority_requests: false,
};
Fetcher::new(context, config)
}
fn create_unservable_fetcher(
context: &Context,
) -> (Fetcher<Context, PublicKey, MockKey>, PublicKey) {
let public_key = PrivateKey::from_seed(0).public_key();
let peer = PrivateKey::from_seed(1).public_key();
let missing_peer = PrivateKey::from_seed(2).public_key();
let config = Config {
me: Some(public_key.clone()),
fetch_timeout: NZDuration!(Duration::from_secs(5)),
retry_delay: Duration::from_millis(100),
priority_requests: false,
};
let mut fetcher = Fetcher::new(context.child("fetcher"), config);
fetcher.reconcile(&[public_key, peer.clone()]);
fetcher.add_targets(MockKey(1), [missing_peer]);
fetcher.add_ready(MockKey(1));
(fetcher, peer)
}
fn add_test_active(fetcher: &mut Fetcher<Context, PublicKey, MockKey>, id: ID, key: MockKey) {
let peer = PrivateKey::from_seed(1).public_key();
let now = fetcher.context.current();
let deadline = now + Duration::from_secs(5);
fetcher.active.put(id, deadline);
fetcher.requests.insert(
id,
ActiveRequest {
key,
peer,
start: now,
},
);
}
const DRAWS: usize = 1_200;
fn peer_weight(fetcher: &Fetcher<Context, PublicKey, MockKey>, peer: &PublicKey) -> u64 {
let index = fetcher.participants.index(peer).unwrap();
fetcher.participants.peers[index].1.weight()
}
fn first_picks(
fetcher: &mut Fetcher<Context, PublicKey, MockKey>,
key: &MockKey,
) -> AHashMap<PublicKey, usize> {
let mut picks = AHashMap::new();
for _ in 0..DRAWS {
let first = fetcher.get_eligible_peers(key)[0].clone();
*picks.entry(first).or_default() += 1;
}
picks
}
fn send_next<S: Sender<PublicKey = PublicKey>>(
fetcher: &mut Fetcher<Context, PublicKey, MockKey>,
sender: &mut WrappedSender<S, wire::Message<MockKey>>,
) -> (ID, PublicKey) {
fetcher.fetch(sender);
assert_eq!(fetcher.len_pending(), 0);
assert_eq!(fetcher.len_active(), 1);
let (&id, req) = fetcher.requests.iter().next().unwrap();
(id, req.peer.clone())
}
#[test]
fn test_retain_function() {
let runner = Runner::default();
runner.start(|context| async {
let mut fetcher = create_test_fetcher(context);
fetcher.add_retry(MockKey(1));
fetcher.add_retry(MockKey(2));
fetcher.add_retry(MockKey(3));
add_test_active(&mut fetcher, 100, MockKey(10));
add_test_active(&mut fetcher, 101, MockKey(20));
add_test_active(&mut fetcher, 102, MockKey(30));
assert_eq!(fetcher.len(), 6);
assert_eq!(fetcher.len_pending(), 3);
assert_eq!(fetcher.len_active(), 3);
fetcher.retain(|key| key.0 <= 10);
assert_eq!(fetcher.len(), 4); assert_eq!(fetcher.len_pending(), 3); assert_eq!(fetcher.len_active(), 1);
assert!(fetcher.pending.contains(&MockKey(1)));
assert!(fetcher.pending.contains(&MockKey(2)));
assert!(fetcher.pending.contains(&MockKey(3)));
assert!(fetcher.requests.values().any(|req| req.key == MockKey(10)));
assert!(!fetcher.requests.values().any(|req| req.key == MockKey(20)));
assert!(!fetcher.requests.values().any(|req| req.key == MockKey(30)));
assert_eq!(fetcher.active.len(), 1);
assert!(fetcher.active.contains(&100));
});
}
#[test]
fn test_len_functions() {
let runner = Runner::default();
runner.start(|context| async {
let mut fetcher = create_test_fetcher(context);
assert_eq!(fetcher.len(), 0);
assert_eq!(fetcher.len_pending(), 0);
assert_eq!(fetcher.len_active(), 0);
fetcher.add_retry(MockKey(1));
fetcher.add_retry(MockKey(2));
assert_eq!(fetcher.len(), 2);
assert_eq!(fetcher.len_pending(), 2);
assert_eq!(fetcher.len_active(), 0);
add_test_active(&mut fetcher, 100, MockKey(10));
add_test_active(&mut fetcher, 101, MockKey(20));
assert_eq!(fetcher.len(), 4);
assert_eq!(fetcher.len_pending(), 2);
assert_eq!(fetcher.len_active(), 2);
assert!(fetcher.pending.remove(&MockKey(1)));
assert_eq!(fetcher.len(), 3);
assert_eq!(fetcher.len_pending(), 1);
assert_eq!(fetcher.len_active(), 2);
fetcher.retain(|key| *key != MockKey(10));
assert_eq!(fetcher.len(), 2);
assert_eq!(fetcher.len_pending(), 1);
assert_eq!(fetcher.len_active(), 1);
});
}
#[test]
fn test_retain_with_empty_collections() {
let runner = Runner::default();
runner.start(|context| async {
let mut fetcher = create_test_fetcher(context);
fetcher.retain(|_| true);
assert_eq!(fetcher.len(), 0);
fetcher.retain(|_| false);
assert_eq!(fetcher.len(), 0);
});
}
#[test]
fn test_retain_all_elements_match_predicate() {
let runner = Runner::default();
runner.start(|context| async {
let mut fetcher = create_test_fetcher(context);
fetcher.add_retry(MockKey(1));
fetcher.add_retry(MockKey(2));
add_test_active(&mut fetcher, 100, MockKey(10));
add_test_active(&mut fetcher, 101, MockKey(20));
let initial_len = fetcher.len();
fetcher.retain(|_| true);
assert_eq!(fetcher.len(), initial_len);
assert_eq!(fetcher.len_pending(), 2);
assert_eq!(fetcher.len_active(), 2);
});
}
#[test]
fn test_retain_no_elements_match_predicate() {
let runner = Runner::default();
runner.start(|context| async {
let mut fetcher = create_test_fetcher(context);
fetcher.add_retry(MockKey(1));
fetcher.add_retry(MockKey(2));
add_test_active(&mut fetcher, 100, MockKey(10));
add_test_active(&mut fetcher, 101, MockKey(20));
fetcher.retain(|_| false);
assert_eq!(fetcher.len(), 0);
assert_eq!(fetcher.len_pending(), 0);
assert_eq!(fetcher.len_active(), 0);
});
}
#[test]
fn test_retain_drops_selected_keys() {
let runner = Runner::default();
runner.start(|context| async {
let mut fetcher = create_test_fetcher(context);
fetcher.add_retry(MockKey(1));
fetcher.add_retry(MockKey(2));
add_test_active(&mut fetcher, 100, MockKey(10));
add_test_active(&mut fetcher, 101, MockKey(20));
fetcher.retain(|key| *key != MockKey(1));
assert_eq!(fetcher.len_pending(), 1);
assert!(!fetcher.contains(&MockKey(1)));
fetcher.retain(|key| *key != MockKey(10));
assert_eq!(fetcher.len_active(), 1);
assert!(!fetcher.contains(&MockKey(10)));
let len = fetcher.len();
fetcher.retain(|key| *key != MockKey(99));
assert_eq!(fetcher.len(), len);
fetcher.retain(|key| *key != MockKey(2));
assert_eq!(fetcher.len_pending(), 0);
assert!(fetcher.get_pending_deadline().is_none());
});
}
#[test]
fn test_contains_function() {
let runner = Runner::default();
runner.start(|context| async {
let mut fetcher = create_test_fetcher(context);
assert!(!fetcher.contains(&MockKey(1)));
fetcher.add_retry(MockKey(1));
assert!(fetcher.contains(&MockKey(1)));
add_test_active(&mut fetcher, 100, MockKey(10));
assert!(fetcher.contains(&MockKey(10)));
assert!(!fetcher.contains(&MockKey(99)));
fetcher.pending.remove(&MockKey(1));
assert!(!fetcher.contains(&MockKey(1)));
fetcher.retain(|key| *key != MockKey(10));
assert!(!fetcher.contains(&MockKey(10)));
});
}
#[test]
fn test_add_retry_function() {
let runner = Runner::default();
runner.start(|context| async {
let mut fetcher = create_test_fetcher(context);
fetcher.add_retry(MockKey(1));
assert_eq!(fetcher.len_pending(), 1);
assert!(fetcher.contains(&MockKey(1)));
fetcher.add_retry(MockKey(2));
assert_eq!(fetcher.len_pending(), 2);
assert!(fetcher.contains(&MockKey(2)));
assert!(fetcher.get_pending_deadline().is_some());
});
}
#[test]
#[should_panic(expected = "assertion failed")]
fn test_add_retry_duplicate_panics() {
let runner = Runner::default();
runner.start(|context| async {
let mut fetcher = create_test_fetcher(context);
fetcher.add_retry(MockKey(1));
fetcher.add_retry(MockKey(1));
});
}
#[test]
fn test_get_pending_deadline() {
let runner = Runner::default();
runner.start(|context| async {
let mut fetcher = create_test_fetcher(context);
assert!(fetcher.get_pending_deadline().is_none());
fetcher.add_retry(MockKey(1));
assert!(fetcher.get_pending_deadline().is_some());
fetcher.add_retry(MockKey(2));
assert!(fetcher.get_pending_deadline().is_some());
fetcher.pending.clear();
assert!(fetcher.get_pending_deadline().is_none());
});
}
#[test]
fn test_get_active_deadline() {
let runner = Runner::default();
runner.start(|context| async {
let fetcher = create_test_fetcher(context);
assert!(fetcher.get_active_deadline().is_none());
});
}
#[test]
fn test_pop_active() {
let runner = Runner::default();
runner.start(|context| async {
let fetcher = create_test_fetcher(context);
assert!(fetcher.get_active_deadline().is_none());
});
}
#[test]
fn test_pop_response_defers_peer_rating() {
let runner = Runner::default();
runner.start(|context| async move {
let mut fetcher = create_test_fetcher(context);
let local_peer = PrivateKey::from_seed(0).public_key();
let peer = PrivateKey::from_seed(1).public_key();
fetcher.reconcile(&[local_peer, peer.clone()]);
add_test_active(&mut fetcher, 100, MockKey(10));
assert!(fetcher.pop_response(999, &peer).is_none());
assert_eq!(fetcher.len_active(), 1);
fetcher.context.sleep(Duration::from_millis(20)).await;
let (key, elapsed) = fetcher.pop_response(100, &peer).expect("matching response");
assert_eq!(key, MockKey(10));
assert_eq!(elapsed, Duration::from_millis(20));
assert_eq!(fetcher.len_active(), 0);
let initial = Score::default().weight();
assert_eq!(peer_weight(&fetcher, &peer), initial);
fetcher.record_response(&peer, elapsed, 1024);
assert!(peer_weight(&fetcher, &peer) > initial);
});
}
#[test]
fn test_pop_response_requires_matching_peer() {
let runner = Runner::default();
runner.start(|context| async {
let mut fetcher = create_test_fetcher(context);
let peer = PrivateKey::from_seed(1).public_key();
let other_peer = PrivateKey::from_seed(2).public_key();
add_test_active(&mut fetcher, 100, MockKey(10));
add_test_active(&mut fetcher, 101, MockKey(20));
assert!(fetcher.pop_response(100, &other_peer).is_none());
assert_eq!(fetcher.len_active(), 2);
assert!(fetcher.active.contains(&100));
assert_eq!(
fetcher.requests.get(&100).map(|req| &req.key),
Some(&MockKey(10))
);
let (key, _) = fetcher
.pop_response(100, &peer)
.expect("matching peer's response must pop the request");
assert_eq!(key, MockKey(10));
assert_eq!(fetcher.len_active(), 1);
assert!(!fetcher.active.contains(&100));
assert!(fetcher.active.contains(&101));
assert_eq!(
fetcher.requests.get(&101).map(|req| &req.key),
Some(&MockKey(20))
);
assert!(fetcher.pop_response(100, &peer).is_none());
assert_eq!(fetcher.len_active(), 1);
assert!(fetcher.active.contains(&101));
});
}
#[test]
fn test_record_response_scores_by_size() {
let runner = Runner::default();
runner.start(|context| async move {
let mut fetcher = create_test_fetcher(context);
let local_peer = PrivateKey::from_seed(0).public_key();
let small = PrivateKey::from_seed(1).public_key();
let large = PrivateKey::from_seed(2).public_key();
fetcher.reconcile(&[local_peer, small.clone(), large.clone()]);
let elapsed = Duration::from_millis(20);
fetcher.record_response(&small, elapsed, 1);
fetcher.record_response(&large, elapsed, 1_000);
assert!(peer_weight(&fetcher, &large) > peer_weight(&fetcher, &small));
let picks = first_picks(&mut fetcher, &MockKey(1));
assert!(picks[&large] > DRAWS * 9 / 10, "{picks:?}");
});
}
#[test]
fn test_empty_fast_reply_cannot_outrank_real_payload() {
let runner = Runner::default();
runner.start(|context| async move {
let mut fetcher = create_test_fetcher(context);
let local_peer = PrivateKey::from_seed(0).public_key();
let empty = PrivateKey::from_seed(1).public_key();
let real = PrivateKey::from_seed(2).public_key();
fetcher.reconcile(&[local_peer, empty.clone(), real.clone()]);
for _ in 0..3 {
fetcher.record_response(&empty, Duration::from_millis(1), 0);
fetcher.record_response(&real, Duration::from_millis(50), 1024);
}
assert_eq!(peer_weight(&fetcher, &empty), 1);
assert!(peer_weight(&fetcher, &real) > peer_weight(&fetcher, &empty));
let picks = first_picks(&mut fetcher, &MockKey(1));
assert!(picks[&real] > DRAWS * 9 / 10, "{picks:?}");
});
}
#[test]
fn test_failures_charge_elapsed_time_and_retry_delay() {
for (retry_delay, expected) in [
(Duration::ZERO, [1_024_001, 10_139, 205]),
(Duration::from_millis(100), [10_139, 5_095, 201]),
(Duration::from_secs(1), [1_023, 931, 171]),
] {
let weights = [Some(Duration::ZERO), Some(Duration::from_millis(100)), None].map(
|response_time| {
Runner::timed(Duration::from_secs(10)).start(|context| async move {
let mut fetcher = create_test_fetcher(context.child("fetcher"));
fetcher.retry_delay = retry_delay;
let peer = PrivateKey::from_seed(1).public_key();
fetcher.reconcile(slice::from_ref(&peer));
fetcher.record_response(&peer, Duration::from_millis(1), 1024);
let initial = peer_weight(&fetcher, &peer);
let key = MockKey(1);
fetcher.add_ready(key.clone());
let mut sender = WrappedSender::new(SuccessMockSender::default());
let (id, _) = send_next(&mut fetcher, &mut sender);
let elapsed = response_time.unwrap_or(fetcher.fetch_timeout.get());
if !elapsed.is_zero() {
context.sleep(elapsed).await;
}
let failed = if response_time.is_some() {
fetcher.pop_missing(id, &peer)
} else {
fetcher.pop_active()
};
assert_eq!(failed, Some(key.clone()));
let weight = peer_weight(&fetcher, &peer);
assert!(weight <= initial);
fetcher.add_retry(key);
assert_eq!(
fetcher.get_pending_deadline(),
Some(context.current() + retry_delay)
);
weight
})
},
);
assert_eq!(weights, expected, "retry delay: {retry_delay:?}");
let [immediate, delayed, silent] = weights;
assert!(immediate > delayed);
assert!(delayed > silent);
}
}
#[test]
fn test_failures_demote_fast_peers_until_successful_service() {
for timeout in [false, true] {
Runner::default().start(|context| async move {
let mut fetcher = create_test_fetcher(context.child("fetcher"));
let fast = PrivateKey::from_seed(1).public_key();
let steady = PrivateKey::from_seed(2).public_key();
fetcher.reconcile(&[fast.clone(), steady.clone()]);
let mut expected = Score::default();
expected.record(Duration::from_millis(1), 1024);
fetcher.record_response(&fast, Duration::from_millis(1), 1024);
fetcher.record_response(&steady, Duration::from_millis(50), 1024);
let initial = peer_weight(&fetcher, &fast);
assert!(initial > peer_weight(&fetcher, &steady));
let key = MockKey(1);
fetcher.add_targets(key.clone(), [fast.clone()]);
fetcher.add_ready(key.clone());
let mut sender = WrappedSender::new(SuccessMockSender::default());
let (id, peer) = send_next(&mut fetcher, &mut sender);
assert_eq!(peer, fast);
let elapsed = if timeout {
fetcher.fetch_timeout.get()
} else {
Duration::from_millis(100)
};
context.sleep(elapsed).await;
let failed = if timeout {
fetcher.pop_active()
} else {
fetcher.pop_missing(id, &peer)
};
assert_eq!(failed, Some(key.clone()));
expected.record(elapsed + fetcher.retry_delay, 0);
let demoted = expected.weight();
assert_eq!(peer_weight(&fetcher, &fast), demoted);
assert!(demoted < peer_weight(&fetcher, &steady));
fetcher.add_retry(key.clone());
context.sleep(Duration::from_secs(3_600)).await;
assert_eq!(peer_weight(&fetcher, &fast), demoted);
let (id, peer) = send_next(&mut fetcher, &mut sender);
assert_eq!(peer_weight(&fetcher, &fast), demoted);
context.sleep(Duration::from_millis(1)).await;
let (_, elapsed) = fetcher.pop_response(id, &peer).unwrap();
assert_eq!(peer_weight(&fetcher, &fast), demoted);
fetcher.record_response(&peer, elapsed, 1024);
expected.record(elapsed, 1024);
assert_eq!(peer_weight(&fetcher, &fast), expected.weight());
assert!(expected.weight() > demoted);
assert!(expected.weight() < initial);
});
}
}
#[test]
fn test_reconcile_and_block() {
let runner = Runner::default();
runner.start(|context| async {
let mut fetcher = create_test_fetcher(context);
let peer1 = PrivateKey::from_seed(1).public_key();
let peer2 = PrivateKey::from_seed(2).public_key();
fetcher.reconcile(&[peer1.clone(), peer2.clone()]);
fetcher.set_blocked([peer1.clone()]);
assert_eq!(fetcher.get_eligible_peers(&MockKey(1)), vec![peer2]);
assert!(fetcher.participants.index(&peer1).is_some());
});
}
#[test]
fn test_edge_cases_empty_state() {
let runner = Runner::default();
runner.start(|context| async {
let fetcher = create_test_fetcher(context);
assert_eq!(fetcher.len(), 0);
assert_eq!(fetcher.len_pending(), 0);
assert_eq!(fetcher.len_active(), 0);
assert!(!fetcher.contains(&MockKey(1)));
assert!(fetcher.get_pending_deadline().is_none());
assert!(fetcher.get_active_deadline().is_none());
});
}
#[test]
fn test_retain_edge_cases() {
let runner = Runner::default();
runner.start(|context| async {
let mut fetcher = create_test_fetcher(context);
fetcher.retain(|key| *key != MockKey(1));
assert_eq!(fetcher.len(), 0);
fetcher.add_retry(MockKey(1));
fetcher.retain(|key| *key != MockKey(1));
assert_eq!(fetcher.len(), 0);
fetcher.retain(|key| *key != MockKey(1));
assert_eq!(fetcher.len(), 0);
});
}
#[test]
fn test_retain_preserves_active_state() {
let runner = Runner::default();
runner.start(|context| async {
let mut fetcher = create_test_fetcher(context);
add_test_active(&mut fetcher, 100, MockKey(1));
add_test_active(&mut fetcher, 101, MockKey(2));
fetcher.retain(|key| key.0 == 1);
assert_eq!(fetcher.len_active(), 1);
assert!(fetcher.requests.values().any(|req| req.key == MockKey(1)));
assert!(!fetcher.requests.values().any(|req| req.key == MockKey(2)));
assert!(fetcher.requests.contains_key(&100));
});
}
#[test]
fn test_mixed_operations() {
let runner = Runner::default();
runner.start(|context| async {
let mut fetcher = create_test_fetcher(context);
fetcher.add_retry(MockKey(1));
fetcher.add_retry(MockKey(2));
add_test_active(&mut fetcher, 100, MockKey(10));
add_test_active(&mut fetcher, 101, MockKey(20));
assert_eq!(fetcher.len(), 4);
fetcher.retain(|key| *key != MockKey(1) && *key != MockKey(10));
assert_eq!(fetcher.len(), 2);
fetcher.retain(|key| key.0 <= 20);
assert_eq!(fetcher.len(), 2);
assert!(fetcher.contains(&MockKey(2)));
assert!(fetcher.contains(&MockKey(20)));
fetcher.retain(|_| false);
assert_eq!(fetcher.len(), 0);
});
}
#[test]
fn test_ready_vs_retry() {
let runner = Runner::default();
runner.start(|context| async move {
let mut fetcher = create_test_fetcher(context.child("fetcher"));
fetcher.add_retry(MockKey(1));
fetcher.add_ready(MockKey(2));
assert_eq!(fetcher.len(), 2);
assert_eq!(fetcher.len_pending(), 2);
assert_eq!(fetcher.len_active(), 0);
let deadline = fetcher.get_pending_deadline().unwrap();
assert_eq!(deadline, context.current());
let (key, _) = fetcher.pending.pop().unwrap();
assert_eq!(key, MockKey(2));
let deadline = fetcher.get_pending_deadline().unwrap();
assert_eq!(deadline, context.current() + Duration::from_millis(100));
let (key, _) = fetcher.pending.pop().unwrap();
assert_eq!(key, MockKey(1));
});
}
#[test]
fn test_ready_requests_precede_all_retries() {
let runner = Runner::default();
runner.start(|context| async move {
let mut fetcher = create_test_fetcher(context.child("fetcher"));
fetcher.add_retry(MockKey(1));
context.sleep(Duration::from_millis(50)).await;
fetcher.add_retry(MockKey(2));
context.sleep(Duration::from_millis(200)).await;
fetcher.add_ready(MockKey(4));
fetcher.add_ready(MockKey(3));
let keys: Vec<_> =
std::iter::from_fn(|| fetcher.pending.pop().map(|(key, _)| key)).collect();
assert_eq!(keys, vec![MockKey(3), MockKey(4), MockKey(1), MockKey(2)]);
});
}
#[test]
fn test_waiter_after_empty() {
let runner = Runner::default();
runner.start(|context| async move {
let public_key = PrivateKey::from_seed(0).public_key();
let other_public_key = PrivateKey::from_seed(1).public_key();
let config = Config {
me: Some(public_key.clone()),
fetch_timeout: NZDuration!(Duration::from_secs(5)),
retry_delay: Duration::from_millis(100),
priority_requests: false,
};
let mut fetcher: Fetcher<_, _, MockKey> =
Fetcher::new(context.child("fetcher"), config);
fetcher.reconcile(&[public_key, other_public_key]);
let mut sender = WrappedSender::new(FailMockSender::default());
fetcher.add_ready(MockKey(1));
fetcher.fetch(&mut sender);
fetcher.fetch(&mut sender);
assert_eq!(fetcher.len_pending(), 1);
let pending_deadline = fetcher.get_pending_deadline().unwrap();
assert!(pending_deadline > context.current());
fetcher.retain(|key| *key != MockKey(1));
assert!(fetcher.get_pending_deadline().is_none());
context.sleep(Duration::from_secs(10)).await;
fetcher.add_retry(MockKey(2));
let next_deadline = fetcher.get_pending_deadline().unwrap();
assert_eq!(
next_deadline,
context.current() + Duration::from_millis(100)
);
});
}
#[test]
fn test_waiter_cleared_on_target_modification() {
let runner = Runner::default();
runner.start(|context| async move {
let public_key = PrivateKey::from_seed(0).public_key();
let peer1 = PrivateKey::from_seed(1).public_key();
let blocked_peer = PrivateKey::from_seed(99).public_key();
let config = Config {
me: Some(public_key.clone()),
fetch_timeout: NZDuration!(Duration::from_secs(5)),
retry_delay: Duration::from_millis(100),
priority_requests: false,
};
let mut fetcher: Fetcher<_, _, MockKey> =
Fetcher::new(context.child("fetcher"), config);
fetcher.reconcile(&[public_key, peer1.clone()]);
let mut sender = WrappedSender::new(FailMockSender::default());
fetcher.set_blocked([blocked_peer.clone()]);
fetcher.add_ready(MockKey(1));
fetcher.add_targets(MockKey(1), [blocked_peer.clone()]);
fetcher.fetch(&mut sender);
assert!(fetcher.waiter.is_some());
let waiter_time = fetcher.waiter.unwrap();
assert!(waiter_time > context.current() + Duration::from_secs(1000));
fetcher.add_targets(MockKey(1), [peer1]);
assert!(fetcher.waiter.is_none());
let deadline = fetcher.get_pending_deadline().unwrap();
assert!(deadline <= context.current() + Duration::from_millis(100));
fetcher.clear_targets(&MockKey(1));
fetcher.add_targets(MockKey(1), [blocked_peer]);
fetcher.fetch(&mut sender);
assert!(fetcher.waiter.is_some());
fetcher.clear_targets(&MockKey(1));
assert!(fetcher.waiter.is_none());
});
}
#[test]
fn test_add_ready_clears_waiter_for_new_fetch() {
let runner = Runner::default();
runner.start(|context| async move {
let public_key = PrivateKey::from_seed(0).public_key();
let peer = PrivateKey::from_seed(1).public_key();
let missing_peer = PrivateKey::from_seed(2).public_key();
let config = Config {
me: Some(public_key.clone()),
fetch_timeout: NZDuration!(Duration::from_secs(5)),
retry_delay: Duration::from_millis(100),
priority_requests: false,
};
let mut fetcher: Fetcher<_, _, MockKey> =
Fetcher::new(context.child("fetcher"), config);
fetcher.reconcile(&[public_key, peer]);
let mut sender = WrappedSender::new(SuccessMockSender::default());
fetcher.add_targets(MockKey(1), [missing_peer]);
fetcher.add_ready(MockKey(1));
fetcher.fetch(&mut sender);
assert!(fetcher.waiter.is_some());
fetcher.add_ready(MockKey(2));
assert_eq!(fetcher.get_pending_deadline(), Some(context.current()));
fetcher.fetch(&mut sender);
assert!(fetcher.pending.contains(&MockKey(1)));
assert!(!fetcher.pending.contains(&MockKey(2)));
assert_eq!(fetcher.len_active(), 1);
});
}
#[test]
fn test_add_retry_clears_waiter_for_retry_deadline() {
let runner = Runner::default();
runner.start(|context| async move {
let (mut fetcher, peer) = create_unservable_fetcher(&context);
let mut sender = WrappedSender::new(SuccessMockSender::default());
fetcher.add_ready(MockKey(2));
fetcher.fetch(&mut sender);
let id = *fetcher.requests.keys().next().unwrap();
fetcher.fetch(&mut sender);
assert!(fetcher.waiter.is_some());
assert_eq!(fetcher.pop_missing(id, &peer), Some(MockKey(2)));
let retry_deadline = context.current() + fetcher.retry_delay;
fetcher.add_retry(MockKey(2));
assert_eq!(fetcher.get_pending_deadline(), Some(context.current()));
fetcher.fetch(&mut sender);
assert_eq!(fetcher.len_active(), 0);
assert!(fetcher.pending.contains(&MockKey(2)));
assert_eq!(fetcher.get_pending_deadline(), Some(retry_deadline));
context.sleep_until(retry_deadline).await;
fetcher.fetch(&mut sender);
assert!(fetcher.pending.contains(&MockKey(1)));
assert!(!fetcher.pending.contains(&MockKey(2)));
assert_eq!(fetcher.len_active(), 1);
});
}
#[test]
fn test_active_timeout_reenqueue_clears_waiter_for_retry_deadline() {
fn run(fetch_first: bool) {
let runner = Runner::default();
runner.start(move |context| async move {
let (mut fetcher, _) = create_unservable_fetcher(&context);
add_test_active(&mut fetcher, 0, MockKey(2));
if fetch_first {
let mut sender = WrappedSender::new(SuccessMockSender::default());
fetcher.fetch(&mut sender);
assert!(fetcher.waiter.is_some());
}
context
.sleep_until(fetcher.get_active_deadline().unwrap())
.await;
let key = fetcher.pop_active().expect("active key should exist");
let retry_deadline = context.current() + fetcher.retry_delay;
fetcher.add_retry(key);
let mut sender = WrappedSender::new(SuccessMockSender::default());
fetcher.fetch(&mut sender);
assert_eq!(fetcher.len_active(), 0);
assert!(fetcher.pending.contains(&MockKey(2)));
assert_eq!(fetcher.get_pending_deadline(), Some(retry_deadline));
context.sleep_until(retry_deadline).await;
fetcher.fetch(&mut sender);
assert!(fetcher.pending.contains(&MockKey(1)));
assert!(!fetcher.pending.contains(&MockKey(2)));
assert_eq!(fetcher.len_active(), 1);
});
}
run(true);
run(false);
}
#[test]
fn test_rate_limited_key_preserves_later_retry_deadline() {
Runner::default().start(|context| async move {
let mut fetcher = create_test_fetcher(context.child("fetcher"));
let peers: Vec<_> = (1..=3)
.map(|seed| PrivateKey::from_seed(seed).public_key())
.collect();
fetcher.reconcile(&peers);
let mut sender = WrappedSender::new(LimitedMockSender::new(
SuccessMockSenderInner,
Quota::per_second(NZU32!(1)),
context.child("rate_limiter"),
));
fetcher.add_targets(MockKey(1), [peers[0].clone()]);
fetcher.add_ready(MockKey(1));
let (id, peer) = send_next(&mut fetcher, &mut sender);
assert_eq!(fetcher.pop_missing(id, &peer), Some(MockKey(1)));
let first_deadline = context.current() + fetcher.retry_delay;
fetcher.add_retry(MockKey(1));
context.sleep(fetcher.retry_delay / 2).await;
fetcher.add_targets(MockKey(2), [peers[1].clone()]);
fetcher.add_ready(MockKey(2));
fetcher.fetch(&mut sender);
assert_eq!(fetcher.len_active(), 1);
let id = *fetcher.requests.keys().next().unwrap();
assert_eq!(fetcher.pop_missing(id, &peers[1]), Some(MockKey(2)));
let second_deadline = context.current() + fetcher.retry_delay;
fetcher.add_retry(MockKey(2));
fetcher.add_targets(MockKey(2), [peers[2].clone()]);
context.sleep_until(first_deadline).await;
fetcher.fetch(&mut sender);
assert_eq!(fetcher.len_active(), 0);
assert!(fetcher.pending.contains(&MockKey(2)));
assert_eq!(fetcher.get_pending_deadline(), Some(second_deadline));
context.sleep_until(second_deadline).await;
fetcher.fetch(&mut sender);
assert!(fetcher.pending.contains(&MockKey(1)));
assert!(!fetcher.pending.contains(&MockKey(2)));
assert_eq!(fetcher.len_active(), 1);
assert_eq!(fetcher.requests.values().next().unwrap().peer, peers[2]);
});
}
#[test]
fn test_add_targets_clears_waiter_for_new_target() {
let runner = Runner::default();
runner.start(|context| async move {
let (mut fetcher, peer) = create_unservable_fetcher(&context);
let mut sender = WrappedSender::new(SuccessMockSender::default());
fetcher.fetch(&mut sender);
assert!(fetcher.waiter.is_some());
fetcher.add_targets(MockKey(1), [peer]);
assert_eq!(fetcher.get_pending_deadline(), Some(context.current()));
fetcher.fetch(&mut sender);
assert!(!fetcher.pending.contains(&MockKey(1)));
assert_eq!(fetcher.len_active(), 1);
});
}
#[test]
fn test_clear_targets_clears_waiter_for_fallback() {
let runner = Runner::default();
runner.start(|context| async move {
let (mut fetcher, _) = create_unservable_fetcher(&context);
let mut sender = WrappedSender::new(SuccessMockSender::default());
fetcher.fetch(&mut sender);
assert!(fetcher.waiter.is_some());
fetcher.clear_targets(&MockKey(1));
assert_eq!(fetcher.get_pending_deadline(), Some(context.current()));
fetcher.fetch(&mut sender);
assert!(!fetcher.pending.contains(&MockKey(1)));
assert_eq!(fetcher.len_active(), 1);
});
}
#[test]
fn test_zero_retry_delay_respects_rate_limit_after_send_failure() {
Runner::timed(Duration::from_secs(3)).start(|context| async move {
let mut fetcher = create_test_fetcher(context.child("fetcher"));
fetcher.retry_delay = Duration::ZERO;
let peer = PrivateKey::from_seed(1).public_key();
fetcher.reconcile(slice::from_ref(&peer));
let mut sender = WrappedSender::new(LimitedMockSender::new(
FailMockSenderInner,
Quota::per_second(NZU32!(1)),
context.child("rate_limiter"),
));
assert_eq!(fetcher.request_id, None);
fetcher.add_ready(MockKey(1));
fetcher.fetch(&mut sender);
assert_eq!(fetcher.len_pending(), 1);
assert_eq!(fetcher.len_active(), 0);
let next_id = fetcher.request_id.unwrap();
assert_eq!(fetcher.get_pending_deadline(), Some(context.current()));
fetcher.fetch(&mut sender);
let refill = context.current() + Duration::from_secs(1);
assert_eq!(fetcher.request_id, Some(next_id));
assert_eq!(fetcher.get_pending_deadline(), Some(refill));
context.sleep_until(refill).await;
fetcher.fetch(&mut sender);
assert_eq!(fetcher.request_id, Some(next_id.wrapping_add(1)));
fetcher.fetch(&mut sender);
assert_eq!(fetcher.request_id, Some(next_id.wrapping_add(1)));
assert_eq!(
fetcher.get_pending_deadline(),
Some(refill + Duration::from_secs(1))
);
});
}
#[test]
fn test_waiter_uses_retry_delay_on_send_failure() {
let cfg = deterministic::Config::default().with_timeout(Some(Duration::from_secs(5)));
let runner = Runner::new(cfg);
runner.start(|context| async move {
let public_key = PrivateKey::from_seed(0).public_key();
let peer1 = PrivateKey::from_seed(1).public_key();
let peer2 = PrivateKey::from_seed(2).public_key();
let retry_delay = Duration::from_millis(100);
let config = Config {
me: Some(public_key.clone()),
fetch_timeout: NZDuration!(Duration::from_secs(5)),
retry_delay,
priority_requests: false,
};
let mut fetcher: Fetcher<_, _, MockKey> =
Fetcher::new(context.child("fetcher"), config);
fetcher.reconcile(&[public_key, peer1, peer2]);
let mut sender = WrappedSender::new(FailMockSender::default());
fetcher.add_ready(MockKey(1));
fetcher.fetch(&mut sender);
assert_eq!(fetcher.len_pending(), 1);
let pending_deadline = fetcher.get_pending_deadline().unwrap();
let max_expected = context.current() + retry_delay + Duration::from_millis(10);
assert!(
pending_deadline <= max_expected,
"pending deadline {:?} exceeds the retry delay",
pending_deadline.duration_since(context.current())
);
let wait_duration = pending_deadline
.duration_since(context.current())
.unwrap_or(Duration::ZERO);
context.sleep(wait_duration).await;
fetcher.fetch(&mut sender);
});
}
#[test]
fn test_eligible_peers_follow_latest_set_and_targets() {
let runner = Runner::default();
runner.start(|context| async move {
let mut fetcher = create_test_fetcher(context);
let me = PrivateKey::from_seed(0).public_key();
let primary = PrivateKey::from_seed(1).public_key();
let secondary = PrivateKey::from_seed(2).public_key();
let other = PrivateKey::from_seed(3).public_key();
let untracked = PrivateKey::from_seed(4).public_key();
fetcher.reconcile(&[me, primary.clone(), secondary.clone(), other.clone()]);
let mut eligible = fetcher.get_eligible_peers(&MockKey(1));
eligible.sort();
let mut expected = vec![primary.clone(), secondary.clone(), other];
expected.sort();
assert_eq!(eligible, expected);
fetcher.add_targets(MockKey(1), [primary.clone(), secondary.clone(), untracked]);
let eligible = fetcher.get_eligible_peers(&MockKey(1));
assert_eq!(eligible.len(), 2);
assert!(eligible.contains(&primary));
assert!(eligible.contains(&secondary));
fetcher.add_targets(MockKey(2), [secondary.clone()]);
assert_eq!(fetcher.get_eligible_peers(&MockKey(2)), vec![secondary]);
});
}
#[test]
fn test_add_targets() {
let runner = Runner::default();
runner.start(|context| async {
let mut fetcher = create_test_fetcher(context);
let peer1 = PrivateKey::from_seed(1).public_key();
let peer2 = PrivateKey::from_seed(2).public_key();
let peer3 = PrivateKey::from_seed(3).public_key();
assert!(fetcher.targets.is_empty());
fetcher.add_targets(MockKey(1), [peer1.clone()]);
assert_eq!(fetcher.targets.len(), 1);
assert!(fetcher.targets.get(&MockKey(1)).unwrap().contains(&peer1));
fetcher.add_targets(MockKey(1), [peer2.clone()]);
assert_eq!(fetcher.targets.len(), 1);
let targets = fetcher.targets.get(&MockKey(1)).unwrap();
assert_eq!(targets.len(), 2);
assert!(targets.contains(&peer1));
assert!(targets.contains(&peer2));
fetcher.add_targets(MockKey(2), [peer1.clone()]);
assert_eq!(fetcher.targets.len(), 2);
assert!(fetcher.targets.get(&MockKey(2)).unwrap().contains(&peer1));
fetcher.add_targets(MockKey(1), [peer1.clone()]);
assert_eq!(fetcher.targets.get(&MockKey(1)).unwrap().len(), 2);
fetcher.add_targets(MockKey(1), [peer3.clone()]);
assert_eq!(fetcher.targets.get(&MockKey(1)).unwrap().len(), 3);
assert!(fetcher.targets.get(&MockKey(1)).unwrap().contains(&peer3));
fetcher.clear_targets(&MockKey(1));
assert!(!fetcher.targets.contains_key(&MockKey(1)));
fetcher.add_targets(MockKey(3), [peer1.clone()]);
assert!(fetcher.targets.get(&MockKey(3)).unwrap().contains(&peer1));
});
}
#[test]
fn test_targets_cleanup() {
let runner = Runner::default();
runner.start(|context| async {
let mut fetcher = create_test_fetcher(context);
let peer1 = PrivateKey::from_seed(1).public_key();
let peer2 = PrivateKey::from_seed(2).public_key();
fetcher.add_targets(MockKey(1), [peer1.clone()]);
fetcher.add_targets(MockKey(2), [peer1.clone()]);
fetcher.add_retry(MockKey(1));
fetcher.add_retry(MockKey(2));
assert_eq!(fetcher.targets.len(), 2);
fetcher.retain(|key| *key != MockKey(1));
assert!(!fetcher.targets.contains_key(&MockKey(1)));
assert!(fetcher.targets.contains_key(&MockKey(2)));
fetcher.retain(|key| *key != MockKey(2));
assert!(fetcher.targets.is_empty());
fetcher.add_targets(MockKey(1), [peer1.clone(), peer2.clone()]);
fetcher.add_targets(MockKey(2), [peer1.clone()]);
fetcher.add_targets(MockKey(3), [peer2]);
assert_eq!(fetcher.targets.len(), 3);
fetcher.retain(|_| false);
assert!(fetcher.targets.is_empty());
fetcher.add_targets(MockKey(1), [peer1.clone()]);
fetcher.add_targets(MockKey(2), [peer1.clone()]);
fetcher.add_targets(MockKey(10), [peer1.clone()]);
fetcher.add_targets(MockKey(20), [peer1]);
assert_eq!(fetcher.targets.len(), 4);
fetcher.retain(|key| key.0 <= 5);
assert_eq!(fetcher.targets.len(), 2);
assert!(fetcher.targets.contains_key(&MockKey(1)));
assert!(fetcher.targets.contains_key(&MockKey(2)));
assert!(!fetcher.targets.contains_key(&MockKey(10)));
assert!(!fetcher.targets.contains_key(&MockKey(20)));
});
}
#[test]
fn test_blocked_peers_keep_targets_until_unblocked() {
let runner = Runner::default();
runner.start(|context| async move {
let mut fetcher = create_test_fetcher(context.child("fetcher"));
let public_key = PrivateKey::from_seed(0).public_key();
let peer1 = PrivateKey::from_seed(1).public_key();
let peer2 = PrivateKey::from_seed(2).public_key();
fetcher.reconcile(&[public_key, peer1.clone(), peer2.clone()]);
fetcher.add_targets(MockKey(1), [peer1.clone()]);
fetcher.add_targets(MockKey(2), [peer1.clone(), peer2.clone()]);
fetcher.set_blocked([peer1.clone()]);
assert!(fetcher.targets.get(&MockKey(1)).unwrap().contains(&peer1));
assert!(fetcher.get_eligible_peers(&MockKey(1)).is_empty());
assert_eq!(fetcher.get_eligible_peers(&MockKey(2)), vec![peer2]);
fetcher.waiter = Some(context.current());
fetcher.set_blocked([]);
assert!(fetcher.waiter.is_none());
assert_eq!(fetcher.get_eligible_peers(&MockKey(1)), vec![peer1]);
});
}
#[test]
fn test_target_behavior_on_send_failure() {
let runner = Runner::default();
runner.start(|context| async move {
let mut fetcher = create_test_fetcher(context.child("fetcher"));
let public_key = PrivateKey::from_seed(0).public_key();
let peer1 = PrivateKey::from_seed(1).public_key();
let peer2 = PrivateKey::from_seed(2).public_key();
let peer3 = PrivateKey::from_seed(3).public_key();
fetcher.reconcile(&[public_key, peer1.clone(), peer2.clone(), peer3]);
let mut sender = WrappedSender::new(FailMockSender::default());
fetcher.add_targets(MockKey(2), [peer1, peer2]);
fetcher.add_ready(MockKey(2));
assert_eq!(fetcher.targets.get(&MockKey(2)).unwrap().len(), 2);
fetcher.fetch(&mut sender);
assert_eq!(fetcher.targets.get(&MockKey(2)).unwrap().len(), 2);
assert!(fetcher.pending.contains(&MockKey(2)));
});
}
#[test]
fn test_target_retention_on_pop() {
let runner = Runner::default();
runner.start(|context| async move {
let mut fetcher = create_test_fetcher(context.child("fetcher"));
let public_key = PrivateKey::from_seed(0).public_key();
let peer1 = PrivateKey::from_seed(1).public_key();
let peer2 = PrivateKey::from_seed(2).public_key();
fetcher.reconcile(&[public_key, peer1.clone(), peer2.clone()]);
let mut sender = WrappedSender::new(SuccessMockSender::default());
fetcher.add_targets(MockKey(1), [peer1.clone(), peer2.clone()]);
fetcher.add_ready(MockKey(1));
assert_eq!(fetcher.targets.get(&MockKey(1)).unwrap().len(), 2);
fetcher.fetch(&mut sender);
context.sleep(Duration::from_millis(200)).await;
assert_eq!(fetcher.pop_active(), Some(MockKey(1)));
assert_eq!(fetcher.targets.get(&MockKey(1)).unwrap().len(), 2);
fetcher.targets.clear();
fetcher.add_targets(MockKey(2), [peer1.clone()]);
fetcher.add_ready(MockKey(2));
fetcher.fetch(&mut sender);
let id = *fetcher.active.iter().next().unwrap().0;
assert_eq!(fetcher.pop_missing(id, &peer1), Some(MockKey(2)));
assert!(fetcher.targets.get(&MockKey(2)).unwrap().contains(&peer1));
fetcher.targets.clear();
fetcher.add_targets(MockKey(3), [peer1.clone()]);
fetcher.add_ready(MockKey(3));
fetcher.fetch(&mut sender);
let id = *fetcher.active.iter().next().unwrap().0;
assert_eq!(
fetcher.pop_response(id, &peer1).map(|(key, _)| key),
Some(MockKey(3))
);
assert!(fetcher.targets.get(&MockKey(3)).unwrap().contains(&peer1));
});
}
#[test]
fn test_no_fallback_when_targets_unavailable() {
let runner = Runner::default();
runner.start(|context| async move {
let mut fetcher = create_test_fetcher(context.child("fetcher"));
let public_key = PrivateKey::from_seed(0).public_key();
let peer1 = PrivateKey::from_seed(1).public_key();
let peer2 = PrivateKey::from_seed(2).public_key();
let peer3 = PrivateKey::from_seed(3).public_key();
fetcher.reconcile(&[public_key, peer1, peer2]);
fetcher.add_targets(MockKey(1), [peer3]);
assert!(fetcher.targets.contains_key(&MockKey(1)));
fetcher.add_ready(MockKey(1));
let mut sender = WrappedSender::new(SuccessMockSender::default());
fetcher.fetch(&mut sender);
assert!(fetcher.targets.contains_key(&MockKey(1)));
assert_eq!(fetcher.len_pending(), 1);
assert_eq!(fetcher.len_active(), 0);
assert!(fetcher.waiter.is_some());
let waiter_time = fetcher.waiter.unwrap();
assert!(waiter_time > context.current() + Duration::from_secs(1000));
});
}
#[test]
fn test_clear_targets() {
let runner = Runner::default();
runner.start(|context| async {
let mut fetcher = create_test_fetcher(context);
let peer1 = PrivateKey::from_seed(1).public_key();
let peer2 = PrivateKey::from_seed(2).public_key();
fetcher.add_targets(MockKey(1), [peer1.clone(), peer2]);
fetcher.add_targets(MockKey(2), [peer1]);
assert_eq!(fetcher.targets.len(), 2);
fetcher.clear_targets(&MockKey(1));
assert!(!fetcher.targets.contains_key(&MockKey(1)));
assert!(fetcher.targets.contains_key(&MockKey(2)));
fetcher.clear_targets(&MockKey(99));
assert_eq!(fetcher.targets.len(), 1);
fetcher.clear_targets(&MockKey(2));
assert!(fetcher.targets.is_empty());
});
}
#[test]
fn test_skips_keys_with_rate_limited_targets() {
let runner = Runner::default();
runner.start(|context| async move {
let public_key = PrivateKey::from_seed(0).public_key();
let peer1 = PrivateKey::from_seed(1).public_key();
let peer2 = PrivateKey::from_seed(2).public_key();
let config = Config {
me: Some(public_key.clone()),
fetch_timeout: NZDuration!(Duration::from_secs(5)),
retry_delay: Duration::from_millis(100),
priority_requests: false,
};
let mut fetcher: Fetcher<_, _, MockKey> =
Fetcher::new(context.child("fetcher"), config);
fetcher.reconcile(&[public_key, peer1.clone(), peer2.clone()]);
let quota = Quota::per_second(NZU32!(1));
let mut sender = WrappedSender::new(LimitedMockSender::new(
SuccessMockSenderInner,
quota,
context.child("rate_limiter"),
));
fetcher.add_targets(MockKey(1), [peer1.clone()]);
fetcher.add_targets(MockKey(2), [peer1.clone()]);
fetcher.add_targets(MockKey(3), [peer2.clone()]);
fetcher.add_ready(MockKey(1));
context.sleep(Duration::from_millis(1)).await;
fetcher.add_ready(MockKey(2));
context.sleep(Duration::from_millis(1)).await;
fetcher.add_ready(MockKey(3));
fetcher.fetch(&mut sender);
assert_eq!(fetcher.len_active(), 1);
assert_eq!(fetcher.len_pending(), 2);
assert!(!fetcher.pending.contains(&MockKey(1)));
fetcher.fetch(&mut sender);
assert_eq!(fetcher.len_active(), 2);
assert_eq!(fetcher.len_pending(), 1);
assert!(fetcher.pending.contains(&MockKey(2))); assert!(!fetcher.pending.contains(&MockKey(3)));
fetcher.fetch(&mut sender);
assert_eq!(fetcher.len_active(), 2); assert_eq!(fetcher.len_pending(), 1); assert_eq!(
fetcher.get_pending_deadline(),
Some(context.current() + Duration::from_secs(1))
);
context.sleep(Duration::from_secs(1)).await;
fetcher.fetch(&mut sender);
assert_eq!(fetcher.len_active(), 3);
assert_eq!(fetcher.len_pending(), 0);
});
}
#[test]
fn test_selection_weighted_by_service() {
Runner::default().start(|context| async {
let mut fetcher = create_test_fetcher(context);
let local = PrivateKey::from_seed(0).public_key();
let fast = PrivateKey::from_seed(1).public_key();
let slow = PrivateKey::from_seed(2).public_key();
let medium = PrivateKey::from_seed(3).public_key();
fetcher.reconcile(&[local, fast.clone(), slow.clone(), medium.clone()]);
fetcher.record_response(&fast, Duration::from_millis(10), 1024);
fetcher.record_response(&slow, Duration::from_millis(500), 1024);
fetcher.record_response(&medium, Duration::from_millis(200), 1024);
let mut peers = fetcher.get_eligible_peers(&MockKey(1));
peers.sort();
let mut expected = vec![fast.clone(), slow.clone(), medium.clone()];
expected.sort();
assert_eq!(peers, expected);
let picks = first_picks(&mut fetcher, &MockKey(1));
assert!(picks[&fast] > DRAWS * 8 / 10, "{picks:?}");
assert!(picks[&medium] > picks[&slow], "{picks:?}");
assert!(picks[&slow] > 0, "{picks:?}");
});
}
#[test]
fn test_selection_uniform_without_scores() {
Runner::default().start(|context| async {
let mut fetcher = create_test_fetcher(context);
let peers: Vec<_> = (1..=3)
.map(|seed| PrivateKey::from_seed(seed).public_key())
.collect();
fetcher.reconcile(&peers);
let picks = first_picks(&mut fetcher, &MockKey(1));
for peer in &peers {
let count = picks[peer];
assert!((DRAWS / 4..DRAWS * 5 / 12).contains(&count), "{picks:?}");
}
});
}
#[test]
fn test_retries_flatten_sampling() {
for targeted in [false, true] {
Runner::seeded(0).start(|context| async move {
let mut fetcher = create_test_fetcher(context);
let peers: Vec<_> = (0..7)
.map(|seed| PrivateKey::from_seed(seed).public_key())
.collect();
fetcher.reconcile(&peers[..6]);
fetcher.set_blocked([peers[5].clone()]);
let key = MockKey(1);
if targeted {
fetcher.add_targets(key.clone(), [1, 2, 3, 5, 6].map(|i| peers[i].clone()));
} else {
fetcher.set_blocked([peers[4].clone(), peers[5].clone()]);
}
fetcher.record_response(&peers[1], Duration::from_millis(1), 1024);
fetcher.record_response(&peers[2], Duration::from_millis(100), 1024);
fetcher.record_response(&peers[3], Duration::from_millis(100), 1024);
fetcher.record_response(&peers[4], Duration::from_nanos(1), 1024);
let weighted = peer_weight(&fetcher, &peers[1]) as f64
/ (1..=3)
.map(|i| peer_weight(&fetcher, &peers[i]))
.sum::<u64>() as f64;
let mut previous = DRAWS + 1;
for (retries, weighted_share) in [(0, 1.0), (1, 0.8), (4, 0.5), (u32::MAX, 0.0)] {
fetcher.retries.insert(key.clone(), retries);
let picks = first_picks(&mut fetcher, &key);
assert_eq!(picks.len(), 3, "{picks:?}");
for peer in &peers[1..=3] {
assert!(picks.contains_key(peer), "{picks:?}");
}
let expected =
DRAWS as f64 * (weighted_share * weighted + (1.0 - weighted_share) / 3.0);
let count = picks[&peers[1]];
assert!(
(count as f64 - expected).abs() < DRAWS as f64 / 12.0,
"retries={retries} count={count} expected={expected}"
);
assert!(count < previous, "retries={retries}: {picks:?}");
previous = count;
assert_cached_weights(&fetcher);
}
});
}
}
#[test]
fn test_retry_revisits_peers_after_each_outcome() {
for targeted in [false, true] {
for outcome in ["timeout", "missing", "partial"] {
Runner::seeded(0).start(|context| async move {
let mut fetcher = create_test_fetcher(context.child("fetcher"));
let peers: Vec<_> = (1..=3)
.map(|seed| PrivateKey::from_seed(seed).public_key())
.collect();
fetcher.reconcile(&peers);
let key = MockKey(1);
let eligible = if targeted { &peers[..2] } else { &peers[..] };
if targeted {
fetcher.add_targets(key.clone(), eligible.iter().cloned());
}
fetcher.record_response(&peers[0], Duration::from_nanos(1), 1024);
fetcher.add_ready(key.clone());
let mut sender = WrappedSender::new(SuccessMockSender::default());
let (id, peer) = send_next(&mut fetcher, &mut sender);
assert_eq!(peer, peers[0]);
match outcome {
"timeout" => {
context
.sleep_until(fetcher.get_active_deadline().unwrap())
.await;
assert_eq!(fetcher.pop_active(), Some(key.clone()));
}
"missing" => assert_eq!(fetcher.pop_missing(id, &peer), Some(key.clone())),
"partial" => {
let (_, elapsed) = fetcher.pop_response(id, &peer).unwrap();
fetcher.record_response(&peer, elapsed, 1024);
}
_ => unreachable!(),
}
fetcher.add_retry(key.clone());
assert_eq!(fetcher.retries[&key], 1);
let candidates: AHashSet<_> =
fetcher.get_eligible_peers(&key).into_iter().collect();
assert_eq!(candidates, eligible.iter().cloned().collect());
assert!(first_picks(&mut fetcher, &key).contains_key(&peer));
assert_cached_weights(&fetcher);
});
}
}
}
#[test]
fn test_single_silent_target_is_retried() {
Runner::default().start(|context| async move {
let mut fetcher = create_test_fetcher(context.child("fetcher"));
let target = PrivateKey::from_seed(1).public_key();
let other = PrivateKey::from_seed(2).public_key();
fetcher.reconcile(&[target.clone(), other]);
fetcher.add_targets(MockKey(1), [target.clone()]);
let mut sender = WrappedSender::new(SuccessMockSender::default());
fetcher.add_ready(MockKey(1));
for _ in 0..3 {
context
.sleep_until(fetcher.get_pending_deadline().unwrap())
.await;
let (_, peer) = send_next(&mut fetcher, &mut sender);
assert_eq!(peer, target);
context
.sleep_until(fetcher.get_active_deadline().unwrap())
.await;
assert_eq!(fetcher.pop_active(), Some(MockKey(1)));
fetcher.add_retry(MockKey(1));
}
});
}
#[test]
fn test_retry_history_lifecycle() {
Runner::default().start(|context| async move {
let mut fetcher = create_test_fetcher(context.child("fetcher"));
let peers: Vec<_> = (1..=2)
.map(|seed| PrivateKey::from_seed(seed).public_key())
.collect();
fetcher.reconcile(&peers);
let mut sender = WrappedSender::new(SuccessMockSender::default());
let key = MockKey(1);
fetcher.add_targets(key.clone(), [peers[0].clone()]);
fetcher.add_ready(key.clone());
let (id, peer) = send_next(&mut fetcher, &mut sender);
assert_eq!(fetcher.pop_missing(id, &peer), Some(key.clone()));
fetcher.add_retry(key.clone());
assert_eq!(fetcher.retries[&key], 1);
fetcher.clear_targets(&key);
fetcher.add_targets(key.clone(), [peers[1].clone()]);
fetcher.reconcile(&peers[1..]);
assert_eq!(fetcher.retries[&key], 1);
context
.sleep_until(fetcher.get_pending_deadline().unwrap())
.await;
let (id, peer) = send_next(&mut fetcher, &mut sender);
assert!(fetcher.pop_response(id, &peer).is_some());
fetcher.clear(&key);
assert!(!fetcher.targets.contains_key(&key));
assert!(!fetcher.retries.contains_key(&key));
fetcher.add_ready(key.clone());
assert!(!fetcher.retries.contains_key(&key));
send_next(&mut fetcher, &mut sender);
context
.sleep_until(fetcher.get_active_deadline().unwrap())
.await;
assert_eq!(fetcher.pop_active(), Some(key.clone()));
fetcher.add_retry(key.clone());
assert_eq!(fetcher.retries[&key], 1);
fetcher.retries.insert(key.clone(), u32::MAX);
assert!(fetcher.pending.remove(&key));
fetcher.add_retry(key.clone());
assert_eq!(fetcher.retries[&key], u32::MAX);
fetcher.add_retry(MockKey(2));
fetcher.retain(|key| *key == MockKey(2));
assert!(!fetcher.retries.contains_key(&key));
assert_eq!(fetcher.retries[&MockKey(2)], 1);
});
}
fn assert_cached_weights(fetcher: &Fetcher<Context, PublicKey, MockKey>) {
let mut total = 0;
for (index, (peer, score)) in fetcher.participants.peers.iter().enumerate() {
let expected = if fetcher.me.as_ref() == Some(peer) || fetcher.blocked.contains(peer) {
0
} else {
score.weight()
};
assert_eq!(fetcher.participants.weights.weight(index), expected);
total += u128::from(expected);
}
assert_eq!(fetcher.participants.weights.total(), total);
}
#[test]
fn test_cached_weights_follow_scores_and_blocks() {
Runner::seeded(0).start(|context| async move {
let mut fetcher = create_test_fetcher(context);
let peers: Vec<_> = (0..4)
.map(|seed| PrivateKey::from_seed(seed).public_key())
.collect();
fetcher.reconcile(&[
peers[3].clone(),
peers[1].clone(),
peers[0].clone(),
peers[2].clone(),
peers[1].clone(),
]);
assert_eq!(fetcher.participants.peers.len(), 4);
for peer in &peers {
fetcher.record_response(peer, Duration::from_millis(10), 1024);
}
assert_cached_weights(&fetcher);
let roster = fetcher.participants.peers.as_ptr();
let peer = fetcher.get_eligible_peers(&MockKey(1))[0].clone();
assert_cached_weights(&fetcher);
let index = fetcher.participants.index(&peer).unwrap();
let before = peer_weight(&fetcher, &peer);
fetcher.update_performance(&peer, fetcher.fetch_timeout.get(), 0);
assert!(fetcher.participants.weights.weight(index) < before);
assert_cached_weights(&fetcher);
fetcher.set_blocked(peers.clone());
fetcher.record_response(&peer, Duration::ZERO, 1024);
assert!(fetcher.candidates(&MockKey(2)).is_none());
assert_cached_weights(&fetcher);
fetcher.set_blocked([peers[2].clone()]);
assert_cached_weights(&fetcher);
fetcher.set_blocked([]);
assert_cached_weights(&fetcher);
assert_eq!(fetcher.participants.peers.as_ptr(), roster);
let mut sender = WrappedSender::new(SuccessMockSender::default());
for key in 3..8 {
fetcher.add_ready(MockKey(key));
let (id, peer) = send_next(&mut fetcher, &mut sender);
assert!(fetcher.pop_response(id, &peer).is_some());
fetcher.clear(&MockKey(key));
assert_cached_weights(&fetcher);
assert_eq!(fetcher.participants.peers.as_ptr(), roster);
}
});
}
#[test]
fn test_cached_targets_use_only_current_members() {
Runner::seeded(0).start(|context| async move {
let mut fetcher = create_test_fetcher(context);
let peers: Vec<_> = (0..32)
.map(|seed| PrivateKey::from_seed(seed).public_key())
.collect();
fetcher.reconcile(&peers);
let target = peers[1].clone();
fetcher.add_targets(MockKey(1), [target.clone()]);
assert_eq!(
fetcher.get_eligible_peers(&MockKey(1)),
vec![target.clone()]
);
fetcher.add_targets(
MockKey(1),
(32..160).map(|seed| PrivateKey::from_seed(seed).public_key()),
);
assert_eq!(
fetcher.get_eligible_peers(&MockKey(1)),
vec![target.clone()]
);
assert_cached_weights(&fetcher);
fetcher.record_response(&peers[2], Duration::from_millis(10), 1024);
let surviving_weight = peer_weight(&fetcher, &peers[2]);
let mut sender = WrappedSender::new(SuccessMockSender::default());
fetcher.add_ready(MockKey(1));
let (id, peer) = send_next(&mut fetcher, &mut sender);
fetcher.reconcile(&[peers[2].clone()]);
assert!(fetcher.pop_response(id, &peer).is_some());
fetcher.record_response(&peer, Duration::from_nanos(1), 1024);
assert_eq!(peer_weight(&fetcher, &peers[2]), surviving_weight);
assert!(fetcher.get_eligible_peers(&MockKey(1)).is_empty());
assert_cached_weights(&fetcher);
fetcher.reconcile(&[target.clone(), peers[2].clone()]);
assert_eq!(peer_weight(&fetcher, &target), Score::default().weight());
assert!(fetcher.targets[&MockKey(1)].contains(&target));
assert_eq!(fetcher.get_eligible_peers(&MockKey(1)), vec![target]);
assert_cached_weights(&fetcher);
});
}
#[test]
fn test_cached_selection_is_independent_of_set_insertion_order() {
fn run(reverse: bool) -> Vec<PublicKey> {
Runner::seeded(9).start(|context| async move {
let mut fetcher = create_test_fetcher(context.child("fetcher"));
let peers: Vec<_> = (0..12)
.map(|seed| PrivateKey::from_seed(seed).public_key())
.collect();
fetcher.reconcile(&peers);
for (index, peer) in peers.iter().enumerate() {
fetcher.record_response(peer, Duration::from_millis(index as u64 + 1), 1024);
}
let mut blocked = vec![peers[10].clone(), peers[11].clone()];
let mut targets = vec![peers[1].clone(), peers[2].clone()];
if reverse {
blocked.reverse();
targets.reverse();
}
fetcher.set_blocked(blocked);
fetcher.add_targets(MockKey(1), targets);
fetcher.add_ready(MockKey(1));
let mut sender = WrappedSender::new(SuccessMockSender::default());
let mut selected = Vec::new();
for round in 0..10 {
if round == 4 {
fetcher.clear_targets(&MockKey(1));
}
context
.sleep_until(fetcher.get_pending_deadline().unwrap())
.await;
let (id, peer) = send_next(&mut fetcher, &mut sender);
assert_eq!(fetcher.pop_missing(id, &peer), Some(MockKey(1)));
fetcher.add_retry(MockKey(1));
selected.push(peer);
assert_cached_weights(&fetcher);
}
selected
})
}
assert_eq!(run(false), run(true));
}
#[test]
fn test_retry_reuses_available_peer_while_another_is_rate_limited() {
for (seed, targeted) in (0..8).flat_map(|seed| [(seed, false), (seed, true)]) {
Runner::seeded(seed).start(|context| async move {
let mut fetcher = create_test_fetcher(context.child("fetcher"));
let a = PrivateKey::from_seed(1).public_key();
let b = PrivateKey::from_seed(2).public_key();
fetcher.reconcile(&[a.clone(), b.clone()]);
let mut sender = WrappedSender::new(LimitedMockSender::new(
SuccessMockSenderInner,
Quota::per_second(NZU32!(1)),
context.child("rate_limiter"),
));
fetcher.add_targets(MockKey(1), [b.clone()]);
fetcher.add_ready(MockKey(1));
let (id, peer) = send_next(&mut fetcher, &mut sender);
assert_eq!(fetcher.pop_missing(id, &peer), Some(MockKey(1)));
fetcher.add_retry(MockKey(1));
context.sleep(Duration::from_secs(1)).await;
fetcher.add_targets(MockKey(2), [a.clone()]);
fetcher.add_ready(MockKey(2));
fetcher.fetch(&mut sender);
let weights = [peer_weight(&fetcher, &a), peer_weight(&fetcher, &b)];
if targeted {
fetcher.add_targets(MockKey(1), [a.clone()]);
} else {
fetcher.clear_targets(&MockKey(1));
}
fetcher.fetch(&mut sender);
assert_eq!(fetcher.len_pending(), 0);
assert_eq!(fetcher.len_active(), 2);
for req in fetcher.requests.values() {
assert_eq!(
req.peer,
if req.key == MockKey(1) {
b.clone()
} else {
a.clone()
}
);
}
assert_eq!(fetcher.retries[&MockKey(1)], 1);
assert_eq!(
[peer_weight(&fetcher, &a), peer_weight(&fetcher, &b)],
weights
);
assert_cached_weights(&fetcher);
});
}
}
#[derive(Debug)]
struct SendState {
ready: AHashSet<PublicKey>,
attempts: Vec<PublicKey>,
closed: bool,
}
#[derive(Clone, Debug)]
struct RecordingSender(Arc<RwLock<SendState>>);
impl UnlimitedSender for RecordingSender {
type PublicKey = PublicKey;
fn send(
&mut self,
recipients: Recipients<Self::PublicKey>,
_message: impl Into<IoBufs> + Send,
_priority: bool,
) -> Unreliable<Feedback> {
let Recipients::One(peer) = recipients else {
unreachable!()
};
let mut state = self.0.write();
state.attempts.push(peer.clone());
if state.ready.contains(&peer) {
Unreliable::new(Feedback::Ok)
} else if state.closed {
Unreliable::new(Feedback::Closed)
} else {
Unreliable::Rejected
}
}
}
impl LimitedSender for RecordingSender {
type PublicKey = PublicKey;
type Checked<'a> = CheckedSender<'a, Self>;
fn check(
&mut self,
recipients: Recipients<Self::PublicKey>,
) -> Result<Self::Checked<'_>, SystemTime> {
Ok(CheckedSender {
sender: self,
recipients,
})
}
}
#[test]
fn test_local_send_failures_are_bounded_and_preserve_scores() {
for targeted in [false, true] {
for closed in [false, true] {
for retries in [None, Some(u32::MAX)] {
Runner::seeded(0).start(|context| async move {
let mut fetcher = create_test_fetcher(context.child("fetcher"));
let peers: Vec<_> = (1..=3)
.map(|seed| PrivateKey::from_seed(seed).public_key())
.collect();
fetcher.reconcile(&peers);
let state = Arc::new(RwLock::new(SendState {
ready: [peers[0].clone()].into_iter().collect(),
attempts: Vec::new(),
closed,
}));
let mut sender = WrappedSender::new(RecordingSender(state.clone()));
fetcher.add_targets(MockKey(0), [peers[0].clone()]);
fetcher.add_ready(MockKey(0));
let (id, peer) = send_next(&mut fetcher, &mut sender);
context.sleep(Duration::from_millis(1)).await;
let (_, elapsed) = fetcher.pop_response(id, &peer).unwrap();
fetcher.record_response(&peer, elapsed, 1024);
fetcher.clear(&MockKey(0));
let before = peer_weight(&fetcher, &peer);
state.write().ready.clear();
state.write().attempts.clear();
let key = MockKey(1);
let eligible = if targeted { &peers[..2] } else { &peers[..] };
if targeted {
fetcher.add_targets(key.clone(), eligible.iter().cloned());
}
if let Some(retries) = retries {
fetcher.retries.insert(key.clone(), retries);
}
fetcher.add_ready(key.clone());
fetcher.fetch(&mut sender);
assert_eq!(fetcher.len_active(), 0);
assert_eq!(fetcher.len_pending(), 1);
assert_eq!(state.read().attempts.len(), eligible.len());
assert_eq!(
state
.read()
.attempts
.iter()
.cloned()
.collect::<AHashSet<_>>(),
eligible.iter().cloned().collect()
);
assert_eq!(peer_weight(&fetcher, &peer), before);
assert_eq!(fetcher.retries.get(&key).copied(), retries);
assert_eq!(
fetcher.get_pending_deadline(),
Some(context.current() + fetcher.retry_delay)
);
assert_cached_weights(&fetcher);
state.write().ready.insert(peer.clone());
state.write().attempts.clear();
context
.sleep_until(fetcher.get_pending_deadline().unwrap())
.await;
let (_, retried) = send_next(&mut fetcher, &mut sender);
assert_eq!(retried, peer);
assert!(state.read().attempts.len() <= eligible.len());
assert_eq!(peer_weight(&fetcher, &peer), before);
assert_eq!(fetcher.retries.get(&key).copied(), retries);
assert_cached_weights(&fetcher);
});
}
}
}
}
}