From 88557ddd0a7ba95a8f6dae9ce346876ff5727a2e Mon Sep 17 00:00:00 2001
From: Jane Lusby <jane@zfnd.org>
Date: Tue, 1 Sep 2020 16:37:06 -0700
Subject: [PATCH] address more comments

---
 zebra-network/src/peer/handshake.rs           |  2 +-
 .../src/peer_set/inventory_registry.rs        | 26 ++++++++++++++++++-
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/zebra-network/src/peer/handshake.rs b/zebra-network/src/peer/handshake.rs
index f96d4a1a..10225f75 100644
--- a/zebra-network/src/peer/handshake.rs
+++ b/zebra-network/src/peer/handshake.rs
@@ -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.
                             //
diff --git a/zebra-network/src/peer_set/inventory_registry.rs b/zebra-network/src/peer_set/inventory_registry.rs
index 9adc2a0b..4a6f27e4 100644
--- a/zebra-network/src/peer_set/inventory_registry.rs
+++ b/zebra-network/src/peer_set/inventory_registry.rs
@@ -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);
     }