YGGverse/aquatic
1
0
Fork 0
mirror of https://github.com/YGGverse/aquatic.git synced 2026-03-31 09:45:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Move common/extract_response_peers to ws since it is only user

Browse source
This commit is contained in:
Joakim Frostegård 2024-01-25 19:05:27 +01:00
parent 09c61b884c
commit 238cce9b16
2 changed files with 137 additions and 140 deletions
Download patch file Download diff file Expand all files Collapse all files

137
crates/common/src/lib.rs
View file

@ -139,140 +139,3 @@ impl CanonicalSocketAddr {
self.0.is_ipv4()
}
}
/// Extract response peers
///
/// If there are more peers in map than `max_num_peers_to_take`, do a random
/// selection of peers from first and second halves of map in order to avoid
/// returning too homogeneous peers.
#[inline]
pub fn extract_response_peers<K, V, R, F>(
rng: &mut impl Rng,
peer_map: &IndexMap<K, V>,
max_num_peers_to_take: usize,
sender_peer_map_key: K,
peer_conversion_function: F,
) -> Vec<R>
where
K: Eq + ::std::hash::Hash,
F: Fn(&K, &V) -> R,
{
if peer_map.len() <= max_num_peers_to_take + 1 {
// This branch: number of peers in map (minus sender peer) is less than
// or equal to number of peers to take, so return all except sender
// peer.
let mut peers = Vec::with_capacity(peer_map.len());
peers.extend(peer_map.iter().filter_map(|(k, v)| {
(*k != sender_peer_map_key).then_some(peer_conversion_function(k, v))
}));
// Handle the case when sender peer is not in peer list. Typically,
// this function will not be called when this is the case.
if peers.len() > max_num_peers_to_take {
peers.pop();
}
peers
} else {
// Note: if this branch is taken, the peer map contains at least two
// more peers than max_num_peers_to_take
let middle_index = peer_map.len() / 2;
// Add one to take two extra peers in case sender peer is among
// selected peers and will need to be filtered out
let num_to_take_per_half = (max_num_peers_to_take / 2) + 1;
let offset_half_one = {
let from = 0;
let to = usize::max(1, middle_index - num_to_take_per_half);
rng.gen_range(from..to)
};
let offset_half_two = {
let from = middle_index;
let to = usize::max(middle_index + 1, peer_map.len() - num_to_take_per_half);
rng.gen_range(from..to)
};
let end_half_one = offset_half_one + num_to_take_per_half;
let end_half_two = offset_half_two + num_to_take_per_half;
let mut peers = Vec::with_capacity(max_num_peers_to_take + 2);
if let Some(slice) = peer_map.get_range(offset_half_one..end_half_one) {
peers.extend(slice.iter().filter_map(|(k, v)| {
(*k != sender_peer_map_key).then_some(peer_conversion_function(k, v))
}));
}
if let Some(slice) = peer_map.get_range(offset_half_two..end_half_two) {
peers.extend(slice.iter().filter_map(|(k, v)| {
(*k != sender_peer_map_key).then_some(peer_conversion_function(k, v))
}));
}
while peers.len() > max_num_peers_to_take {
peers.pop();
}
peers
}
}
#[cfg(test)]
mod tests {
use ahash::HashSet;
use rand::{rngs::SmallRng, SeedableRng};
use super::*;
#[test]
fn test_extract_response_peers() {
let mut rng = SmallRng::from_entropy();
for num_peers_in_map in 0..50 {
for max_num_peers_to_take in 0..50 {
for sender_peer_map_key in 0..50 {
test_extract_response_peers_helper(
&mut rng,
num_peers_in_map,
max_num_peers_to_take,
sender_peer_map_key,
);
}
}
}
}
fn test_extract_response_peers_helper(
rng: &mut SmallRng,
num_peers_in_map: usize,
max_num_peers_to_take: usize,
sender_peer_map_key: usize,
) {
let peer_map = IndexMap::from_iter((0..num_peers_in_map).map(|i| (i, i)));
let response_peers = extract_response_peers(
rng,
&peer_map,
max_num_peers_to_take,
sender_peer_map_key,
|_, p| *p,
);
if num_peers_in_map > max_num_peers_to_take + 1 {
assert_eq!(response_peers.len(), max_num_peers_to_take);
} else {
assert!(response_peers.len() <= max_num_peers_to_take);
}
assert!(!response_peers.contains(&sender_peer_map_key));
assert_eq!(
response_peers.len(),
HashSet::from_iter(response_peers.iter().copied()).len()
);
}
}

140
crates/ws/src/workers/swarm/storage.rs
View file

@ -10,10 +10,9 @@ use hashbrown::HashMap;
use metrics::Gauge;
use rand::rngs::SmallRng;
use aquatic_common::{
extract_response_peers, IndexMap, SecondsSinceServerStart, ServerStartInstant,
};
use aquatic_common::{IndexMap, SecondsSinceServerStart, ServerStartInstant};
use aquatic_ws_protocol::common::*;
use rand::Rng;
use crate::common::*;
use crate::config::Config;
@ -505,3 +504,138 @@ impl PeerStatus {
}
}
}
/// Extract response peers
///
/// If there are more peers in map than `max_num_peers_to_take`, do a random
/// selection of peers from first and second halves of map in order to avoid
/// returning too homogeneous peers.
#[inline]
pub fn extract_response_peers<K, V, R, F>(
rng: &mut impl Rng,
peer_map: &IndexMap<K, V>,
max_num_peers_to_take: usize,
sender_peer_map_key: K,
peer_conversion_function: F,
) -> Vec<R>
where
K: Eq + ::std::hash::Hash,
F: Fn(&K, &V) -> R,
{
if peer_map.len() <= max_num_peers_to_take + 1 {
// This branch: number of peers in map (minus sender peer) is less than
// or equal to number of peers to take, so return all except sender
// peer.
let mut peers = Vec::with_capacity(peer_map.len());
peers.extend(peer_map.iter().filter_map(|(k, v)| {
(*k != sender_peer_map_key).then_some(peer_conversion_function(k, v))
}));
// Handle the case when sender peer is not in peer list. Typically,
// this function will not be called when this is the case.
if peers.len() > max_num_peers_to_take {
peers.pop();
}
peers
} else {
// Note: if this branch is taken, the peer map contains at least two
// more peers than max_num_peers_to_take
let middle_index = peer_map.len() / 2;
// Add one to take two extra peers in case sender peer is among
// selected peers and will need to be filtered out
let num_to_take_per_half = (max_num_peers_to_take / 2) + 1;
let offset_half_one = {
let from = 0;
let to = usize::max(1, middle_index - num_to_take_per_half);
rng.gen_range(from..to)
};
let offset_half_two = {
let from = middle_index;
let to = usize::max(middle_index + 1, peer_map.len() - num_to_take_per_half);
rng.gen_range(from..to)
};
let end_half_one = offset_half_one + num_to_take_per_half;
let end_half_two = offset_half_two + num_to_take_per_half;
let mut peers = Vec::with_capacity(max_num_peers_to_take + 2);
if let Some(slice) = peer_map.get_range(offset_half_one..end_half_one) {
peers.extend(slice.iter().filter_map(|(k, v)| {
(*k != sender_peer_map_key).then_some(peer_conversion_function(k, v))
}));
}
if let Some(slice) = peer_map.get_range(offset_half_two..end_half_two) {
peers.extend(slice.iter().filter_map(|(k, v)| {
(*k != sender_peer_map_key).then_some(peer_conversion_function(k, v))
}));
}
while peers.len() > max_num_peers_to_take {
peers.pop();
}
peers
}
}
#[cfg(test)]
mod tests {
use hashbrown::HashSet;
use rand::{rngs::SmallRng, SeedableRng};
use super::*;
#[test]
fn test_extract_response_peers() {
let mut rng = SmallRng::from_entropy();
for num_peers_in_map in 0..50 {
for max_num_peers_to_take in 0..50 {
for sender_peer_map_key in 0..50 {
test_extract_response_peers_helper(
&mut rng,
num_peers_in_map,
max_num_peers_to_take,
sender_peer_map_key,
);
}
}
}
}
fn test_extract_response_peers_helper(
rng: &mut SmallRng,
num_peers_in_map: usize,
max_num_peers_to_take: usize,
sender_peer_map_key: usize,
) {
let peer_map = IndexMap::from_iter((0..num_peers_in_map).map(|i| (i, i)));
let response_peers = extract_response_peers(
rng,
&peer_map,
max_num_peers_to_take,
sender_peer_map_key,
|_, p| *p,
);
if num_peers_in_map > max_num_peers_to_take + 1 {
assert_eq!(response_peers.len(), max_num_peers_to_take);
} else {
assert!(response_peers.len() <= max_num_peers_to_take);
}
assert!(!response_peers.contains(&sender_peer_map_key));
let unique: HashSet<_> = response_peers.iter().copied().collect();
assert_eq!(response_peers.len(), unique.len(),);
}
}