f Switch to using a HashMap for PeerInfoStore

Author: tnull

src/lib.rs

@@ -106,6 +106,9 @@ const BDK_CLIENT_CONCURRENCY: u8 = 8;
 // The timeout after which we abandon retrying failed payments.
 const LDK_PAYMENT_RETRY_TIMEOUT: Duration = Duration::from_secs(10);
 
+// The time in between peer reconnection attempts.
+const PEER_RECONNECTION_INTERVAL: Duration = Duration::from_secs(10);
+
 #[derive(Debug, Clone)]
 /// Represents the configuration of an [`Node`] instance.
 pub struct Config {
@@ -622,12 +625,11 @@ impl Node {
 		let connect_peer_store = Arc::clone(&self.peer_store);
 		let stop_connect = Arc::clone(&stop_networking);
 		tokio_runtime.spawn(async move {
-			let mut interval = tokio::time::interval(Duration::from_secs(1));
+			let mut interval = tokio::time::interval(PEER_RECONNECTION_INTERVAL);
 			loop {
 				if stop_connect.load(Ordering::Acquire) {
 					return;
 				}
-				interval.tick().await;
 				let pm_peers = connect_pm
 					.get_peer_node_ids()
 					.iter()
@@ -639,18 +641,17 @@ impl Node {
 					.map(|chan| chan.counterparty.node_id)
 					.filter(|id| !pm_peers.contains(id))
 				{
-					for peer_info in connect_peer_store.peers() {
-						if peer_info.pubkey == node_id {
-							let _ = do_connect_peer(
-								peer_info.pubkey,
-								peer_info.address,
-								Arc::clone(&connect_pm),
-								Arc::clone(&connect_logger),
-							)
-							.await;
-						}
+					if let Some(peer_info) = connect_peer_store.get_peer(&node_id) {
+						let _ = do_connect_peer(
+							peer_info.pubkey,
+							peer_info.address,
+							Arc::clone(&connect_pm),
+							Arc::clone(&connect_logger),
+						)
+						.await;
 					}
 				}
+				interval.tick().await;
 			}
 		});
 

src/peer_store.rs

@@ -6,6 +6,7 @@ use lightning::util::ser::{Readable, ReadableArgs, Writeable, Writer};
 
 use bitcoin::secp256k1::PublicKey;
 
+use std::collections::HashMap;
 use std::convert::TryFrom;
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, ToSocketAddrs};
 use std::sync::{Arc, RwLock};
@@ -14,30 +15,21 @@ use std::sync::{Arc, RwLock};
 pub(crate) const PEER_INFO_PERSISTENCE_KEY: &str = "peers";
 
 pub(crate) struct PeerInfoStorage<K: KVStorePersister> {
-	peers: RwLock<Vec<PeerInfo>>,
+	peers: RwLock<HashMap<PublicKey, PeerInfo>>,
 	persister: Arc<K>,
 }
 
 impl<K: KVStorePersister> PeerInfoStorage<K> {
 	pub(crate) fn new(persister: Arc<K>) -> Self {
-		let peers = RwLock::new(Vec::new());
+		let peers = RwLock::new(HashMap::new());
 		Self { peers, persister }
 	}
 
 	pub(crate) fn add_peer(&self, peer_info: PeerInfo) -> Result<(), Error> {
 		let mut locked_peers = self.peers.write().unwrap();
 
-		// Check if we have the peer. If so, either update it or do nothing.
-		for stored_info in locked_peers.iter_mut() {
-			if stored_info.pubkey == peer_info.pubkey {
-				if stored_info.address != peer_info.address {
-					stored_info.address = peer_info.address;
-				}
-				return Ok(());
-			}
-		}
+		locked_peers.insert(peer_info.pubkey, peer_info);
 
-		locked_peers.push(peer_info);
 		self.persister
 			.persist(PEER_INFO_PERSISTENCE_KEY, &PeerInfoStorageSerWrapper(&*locked_peers))
 			.map_err(|_| Error::PersistenceFailed)?;
@@ -48,7 +40,7 @@ impl<K: KVStorePersister> PeerInfoStorage<K> {
 	pub(crate) fn remove_peer(&self, peer_pubkey: &PublicKey) -> Result<(), Error> {
 		let mut locked_peers = self.peers.write().unwrap();
 
-		locked_peers.retain(|info| info.pubkey != *peer_pubkey);
+		locked_peers.remove(peer_pubkey);
 
 		self.persister
 			.persist(PEER_INFO_PERSISTENCE_KEY, &PeerInfoStorageSerWrapper(&*locked_peers))
@@ -57,8 +49,12 @@ impl<K: KVStorePersister> PeerInfoStorage<K> {
 		Ok(())
 	}
 
-	pub(crate) fn peers(&self) -> Vec<PeerInfo> {
-		self.peers.read().unwrap().clone()
+	pub(crate) fn list_peers(&self) -> Vec<PeerInfo> {
+		self.peers.read().unwrap().values().cloned().collect()
+	}
+
+	pub(crate) fn get_peer(&self, peer_pubkey: &PublicKey) -> Option<PeerInfo> {
+		self.peers.read().unwrap().get(peer_pubkey).cloned()
 	}
 }
 
@@ -68,34 +64,37 @@ impl<K: KVStorePersister> ReadableArgs<Arc<K>> for PeerInfoStorage<K> {
 		reader: &mut R, persister: Arc<K>,
 	) -> Result<Self, lightning::ln::msgs::DecodeError> {
 		let read_peers: PeerInfoStorageDeserWrapper = Readable::read(reader)?;
-		let peers: RwLock<Vec<PeerInfo>> = RwLock::new(read_peers.0);
+		let peers: RwLock<HashMap<PublicKey, PeerInfo>> = RwLock::new(read_peers.0);
 		Ok(Self { peers, persister })
 	}
 }
 
 #[derive(Clone, Debug, PartialEq, Eq)]
-pub(crate) struct PeerInfoStorageDeserWrapper(Vec<PeerInfo>);
+pub(crate) struct PeerInfoStorageDeserWrapper(HashMap<PublicKey, PeerInfo>);
 
 impl Readable for PeerInfoStorageDeserWrapper {
 	fn read<R: lightning::io::Read>(
 		reader: &mut R,
 	) -> Result<Self, lightning::ln::msgs::DecodeError> {
 		let len: u16 = Readable::read(reader)?;
-		let mut peers = Vec::with_capacity(len as usize);
+		let mut peers = HashMap::with_capacity(len as usize);
 		for _ in 0..len {
-			peers.push(Readable::read(reader)?);
+			let k: PublicKey = Readable::read(reader)?;
+			let v: PeerInfo = Readable::read(reader)?;
+			peers.insert(k, v);
 		}
 		Ok(Self(peers))
 	}
 }
 
-pub(crate) struct PeerInfoStorageSerWrapper<'a>(&'a Vec<PeerInfo>);
+pub(crate) struct PeerInfoStorageSerWrapper<'a>(&'a HashMap<PublicKey, PeerInfo>);
 
 impl Writeable for PeerInfoStorageSerWrapper<'_> {
 	fn write<W: Writer>(&self, writer: &mut W) -> Result<(), lightning::io::Error> {
 		(self.0.len() as u16).write(writer)?;
-		for e in self.0.iter() {
-			e.write(writer)?;
+		for (k, v) in self.0.iter() {
+			k.write(writer)?;
+			v.write(writer)?;
 		}
 		Ok(())
 	}
@@ -196,9 +195,10 @@ mod tests {
 		let deser_peer_store =
 			PeerInfoStorage::read(&mut &persisted_bytes[..], Arc::clone(&persister)).unwrap();
 
-		let peers = deser_peer_store.peers();
+		let peers = deser_peer_store.list_peers();
 		assert_eq!(peers.len(), 1);
 		assert_eq!(peers[0], expected_peer_info);
+		assert_eq!(deser_peer_store.get_peer(&pubkey), Some(expected_peer_info));
 		assert!(!persister.get_and_clear_did_persist());
 	}