From 0bfd57def23ccdbe9d794c52634b412c24d8d164 Mon Sep 17 00:00:00 2001
From: Henry de Valence <hdevalence@hdevalence.ca>
Date: Thu, 17 Oct 2019 16:25:24 -0700
Subject: [PATCH] Add iteration functions to `AddressBook`.

The disconnected_peers() function allows us to prevent duplicate
connections without maintaining shared state between the peerset and the
dial-additional-peers task.
---
 zebra-network/src/address_book.rs | 48 ++++++++++++++++++++++++-------
 zebrad/src/commands/connect.rs    |  2 +-
 2 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/zebra-network/src/address_book.rs b/zebra-network/src/address_book.rs
index 4e75cb3e4..4f712dcce 100644
--- a/zebra-network/src/address_book.rs
+++ b/zebra-network/src/address_book.rs
@@ -28,21 +28,12 @@ impl AddressBook {
     /// Update the address book with `event`, a [`MetaAddr`] representing
     /// observation of a peer.
     pub fn update(&mut self, event: MetaAddr) {
-        use chrono::Duration as CD;
         use std::collections::hash_map::Entry;
 
-        trace!(
+        debug!(
             ?event,
             data.total = self.by_time.len(),
-            // This would be cleaner if it used "variables" but keeping
-            // it inside the trace! invocation prevents running the range
-            // query unless the output will actually be used.
-            data.recent = self
-                .by_time
-                .range(
-                    (Utc::now() - CD::from_std(constants::LIVE_PEER_DURATION).unwrap())..Utc::now()
-                )
-                .count()
+            data.recent = (self.by_time.len() - self.disconnected_peers().count()),
         );
 
         let MetaAddr {
@@ -73,4 +64,39 @@ impl AddressBook {
             }
         }
     }
+
+    /// Return an iterator over all peers, ordered from most recently seen to
+    /// least recently seen.
+    pub fn peers<'a>(&'a self) -> impl Iterator<Item = MetaAddr> + 'a {
+        self.by_time.iter().rev().map(from_by_time_kv)
+    }
+
+    /// Return an iterator over peers known to be disconnected, ordered from most
+    /// recently seen to least recently seen.
+    pub fn disconnected_peers<'a>(&'a self) -> impl Iterator<Item = MetaAddr> + 'a {
+        use chrono::Duration as CD;
+        use std::ops::Bound::{Excluded, Unbounded};
+
+        // LIVE_PEER_DURATION represents the time interval in which we are
+        // guaranteed to receive at least one message from a peer or close the
+        // connection. Therefore, if the last-seen timestamp is older than
+        // LIVE_PEER_DURATION ago, we know we must have disconnected from it.
+        let cutoff = Utc::now() - CD::from_std(constants::LIVE_PEER_DURATION).unwrap();
+
+        self.by_time
+            .range((Unbounded, Excluded(cutoff)))
+            .rev()
+            .map(from_by_time_kv)
+    }
+}
+
+// Helper impl to convert by_time Iterator Items back to MetaAddrs
+// This could easily be a From impl, but trait impls are public, and this shouldn't be.
+fn from_by_time_kv(by_time_kv: (&DateTime<Utc>, &(SocketAddr, PeerServices))) -> MetaAddr {
+    let (last_seen, (addr, services)) = by_time_kv;
+    MetaAddr {
+        last_seen: last_seen.clone(),
+        addr: addr.clone(),
+        services: services.clone(),
+    }
 }
diff --git a/zebrad/src/commands/connect.rs b/zebrad/src/commands/connect.rs
index ea14ad74f..4c30c6d04 100644
--- a/zebrad/src/commands/connect.rs
+++ b/zebrad/src/commands/connect.rs
@@ -108,7 +108,7 @@ impl ConnectCmd {
             addrs.into_iter().map(|meta| meta.addr).collect::<Vec<_>>()
         };
 
-        let (mut peer_set, _address_book) = zebra_network::init(config, node);
+        let (mut peer_set, address_book) = zebra_network::init(config, node);
 
         info!("waiting for peer_set ready");
         peer_set.ready().await.map_err(Error::from_boxed_compat)?;