From 238cce9b16fca774a7cbdb8a0dcf843ff4d298fa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= Date: Thu, 25 Jan 2024 19:05:27 +0100 Subject: [PATCH] Move common/extract_response_peers to ws since it is only user --- crates/common/src/lib.rs | 137 ------------------------ crates/ws/src/workers/swarm/storage.rs | 140 ++++++++++++++++++++++++- 2 files changed, 137 insertions(+), 140 deletions(-) diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index 6d0dbb6..c8f4243 100644 --- a/crates/common/src/lib.rs +++ b/crates/common/src/lib.rs @@ -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( - rng: &mut impl Rng, - peer_map: &IndexMap, - max_num_peers_to_take: usize, - sender_peer_map_key: K, - peer_conversion_function: F, -) -> Vec -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() - ); - } -} diff --git a/crates/ws/src/workers/swarm/storage.rs b/crates/ws/src/workers/swarm/storage.rs index d579ad9..9eff053 100644 --- a/crates/ws/src/workers/swarm/storage.rs +++ b/crates/ws/src/workers/swarm/storage.rs @@ -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( + rng: &mut impl Rng, + peer_map: &IndexMap, + max_num_peers_to_take: usize, + sender_peer_map_key: K, + peer_conversion_function: F, +) -> Vec +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(),); + } +}