diff --git a/zebra-network/src/address_book.rs b/zebra-network/src/address_book.rs index 4f712dcce..1a162f2a5 100644 --- a/zebra-network/src/address_book.rs +++ b/zebra-network/src/address_book.rs @@ -3,6 +3,7 @@ use std::{ collections::{BTreeMap, HashMap}, + iter::Extend, net::SocketAddr, sync::{Arc, Mutex}, }; @@ -88,6 +89,12 @@ impl AddressBook { .rev() .map(from_by_time_kv) } + + /// Returns an iterator that drains entries from the address book, removing + /// them in order from most recent to least recent. + pub fn drain_recent<'a>(&'a mut self) -> impl Iterator + 'a { + Drain { book: self } + } } // Helper impl to convert by_time Iterator Items back to MetaAddrs @@ -100,3 +107,37 @@ fn from_by_time_kv(by_time_kv: (&DateTime, &(SocketAddr, PeerServices))) -> services: services.clone(), } } + +impl Extend for AddressBook { + fn extend(&mut self, iter: T) + where + T: IntoIterator, + { + for meta in iter.into_iter() { + self.update(meta); + } + } +} + +struct Drain<'a> { + book: &'a mut AddressBook, +} + +impl<'a> Iterator for Drain<'a> { + type Item = MetaAddr; + + fn next(&mut self) -> Option { + let most_recent = self.book.by_time.keys().rev().next()?.clone(); + let (addr, services) = self + .book + .by_time + .remove(&most_recent) + .expect("key from keys() must be present in btreemap"); + self.book.by_addr.remove(&addr); + Some(MetaAddr { + addr, + services, + last_seen: most_recent, + }) + } +} \ No newline at end of file