diff --git a/p2p/src/store.rs b/p2p/src/store.rs index 9dae5387f..a8dae3017 100644 --- a/p2p/src/store.rs +++ b/p2p/src/store.rs @@ -16,8 +16,7 @@ use chrono::Utc; use num::FromPrimitive; -use rand::seq::SliceRandom; -use rand::thread_rng; +use rand::prelude::*; use crate::core::ser::{self, Readable, Reader, Writeable, Writer}; use crate::types::{Capabilities, PeerAddr, ReasonForBan}; @@ -146,31 +145,28 @@ impl PeerStore { batch.commit() } + /// Find some peers in our local db. pub fn find_peers( &self, state: State, cap: Capabilities, count: usize, ) -> Result, Error> { - let mut peers = self + let key = to_key(PEER_PREFIX, ""); + let peers = self .db - .iter::(&to_key(PEER_PREFIX, ""))? + .iter::(&key)? .map(|(_, v)| v) .filter(|p| p.flags == state && p.capabilities.contains(cap)) - .collect::>(); - peers[..].shuffle(&mut thread_rng()); - Ok(peers.iter().take(count).cloned().collect()) + .choose_multiple(&mut thread_rng(), count); + Ok(peers) } /// List all known peers /// Used for /v1/peers/all api endpoint pub fn all_peers(&self) -> Result, Error> { let key = to_key(PEER_PREFIX, ""); - Ok(self - .db - .iter::(&key)? - .map(|(_, v)| v) - .collect::>()) + Ok(self.db.iter::(&key)?.map(|(_, v)| v).collect()) } /// Convenience method to load a peer data, update its status and save it