use choose_multiple to clean get_peers() up (#3457)

This commit is contained in:
Antioch Peverell 2020-10-01 19:42:57 +01:00 committed by GitHub
parent cc4e608d1b
commit f25b7ae032
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

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