address more comments

This commit is contained in:
Jane Lusby 2020-09-01 16:37:06 -07:00 committed by Deirdre Connolly
parent d933abeebf
commit 88557ddd0a
2 changed files with 26 additions and 2 deletions

View File

@ -396,7 +396,7 @@ where
let inv_collector = inv_collector.clone();
async move {
if let Ok(Message::Inv(hashes)) = &msg {
// We reject inventory messages with more than one
// We ignore inventory messages with more than one
// item because they are most likely replies to a
// query rather than a newly gossiped block.
//

View File

@ -1,3 +1,6 @@
//! Inventory Registry Implementation
//!
//! [RFC]: https://zebra.zfnd.org/dev/rfcs/0003-inventory-tracking.html
use crate::{protocol::external::InventoryHash, BoxedStdError};
use futures::Stream;
use std::{
@ -12,16 +15,26 @@ use tokio::{
time::{self, Interval},
};
/// An Inventory Registry for tracking recent inventory advertisements by peer.
///
/// For more details please refer to the [RFC].
///
/// [RFC]: https://zebra.zfnd.org/dev/rfcs/0003-inventory-tracking.html
#[derive(Debug)]
pub struct InventoryRegistry {
/// Map tracking the inventory advertisements from the current interval
/// period
current: HashMap<InventoryHash, HashSet<SocketAddr>>,
/// Map tracking inventory advertisements from the previous interval period
prev: HashMap<InventoryHash, HashSet<SocketAddr>>,
/// Stream of incoming inventory hashes to
/// Stream of incoming inventory hashes to register
inv_stream: broadcast::Receiver<(InventoryHash, SocketAddr)>,
/// Interval tracking how frequently we should rotate our maps
interval: Interval,
}
impl InventoryRegistry {
/// Returns an Inventory Registry
pub fn new(inv_stream: broadcast::Receiver<(InventoryHash, SocketAddr)>) -> Self {
Self {
current: Default::default(),
@ -31,6 +44,8 @@ impl InventoryRegistry {
}
}
/// Returns an iterator over addrs of peers that have recently advertised
/// having `hash` in their inventory.
pub fn peers(&self, hash: &InventoryHash) -> impl Iterator<Item = &SocketAddr> {
let prev = self.prev.get(hash).into_iter();
let current = self.current.get(hash).into_iter();
@ -38,6 +53,12 @@ impl InventoryRegistry {
prev.chain(current).flatten()
}
/// Drive periodic inventory tasks
///
/// # Details
///
/// - rotates HashMaps based on interval events
/// - drains the inv_stream channel and registers all advertised inventory
pub fn poll_inventory(&mut self, cx: &mut Context<'_>) -> Result<(), BoxedStdError> {
while let Poll::Ready(_) = self.interval.poll_tick(cx) {
self.rotate();
@ -50,10 +71,13 @@ impl InventoryRegistry {
Ok(())
}
/// Record that the given inventory `hash` is available from the peer `addr`
fn register(&mut self, hash: InventoryHash, addr: SocketAddr) {
self.current.entry(hash).or_default().insert(addr);
}
/// Replace the prev HashMap with current's and replace current with an empty
/// HashMap
fn rotate(&mut self) {
self.prev = std::mem::take(&mut self.current);
}