Speed up remove_stale_channels_and_tracking nontrivially

lightning/src/routing/gossip.rs

@@ -37,7 +37,9 @@ use crate::ln::types::ChannelId;
 use crate::routing::utxo::{self, UtxoLookup, UtxoResolver};
 use crate::types::features::{ChannelFeatures, InitFeatures, NodeFeatures};
 use crate::types::string::PrintableString;
-use crate::util::indexed_map::{Entry as IndexedMapEntry, IndexedMap};
+use crate::util::indexed_map::{
+	Entry as IndexedMapEntry, IndexedMap, OccupiedEntry as IndexedMapOccupiedEntry,
+};
 use crate::util::logger::{Level, Logger};
 use crate::util::scid_utils::{block_from_scid, scid_from_parts, MAX_SCID_BLOCK};
 use crate::util::ser::{MaybeReadable, Readable, ReadableArgs, RequiredWrapper, Writeable, Writer};
@@ -2356,9 +2358,7 @@ where
 			return;
 		}
 		let min_time_unix: u32 = (current_time_unix - STALE_CHANNEL_UPDATE_AGE_LIMIT_SECS) as u32;
-		// Sadly BTreeMap::retain was only stabilized in 1.53 so we can't switch to it for some
-		// time.
-		let mut scids_to_remove = Vec::new();
+		let mut scids_to_remove = new_hash_set();
 		for (scid, info) in channels.unordered_iter_mut() {
 			if info.one_to_two.is_some()
 				&& info.one_to_two.as_ref().unwrap().last_update < min_time_unix
@@ -2382,19 +2382,24 @@ where
 				if announcement_received_timestamp < min_time_unix as u64 {
 					log_gossip!(self.logger, "Removing channel {} because both directional updates are missing and its announcement timestamp {} being below {}",
 						scid, announcement_received_timestamp, min_time_unix);
-					scids_to_remove.push(*scid);
+					scids_to_remove.insert(*scid);
 				}
 			}
 		}
 		if !scids_to_remove.is_empty() {
 			let mut nodes = self.nodes.write().unwrap();
-			for scid in scids_to_remove {
-				let info = channels
-					.remove(&scid)
-					.expect("We just accessed this scid, it should be present");
-				self.remove_channel_in_nodes(&mut nodes, &info, scid);
-				self.removed_channels.lock().unwrap().insert(scid, Some(current_time_unix));
+			let mut removed_channels_lck = self.removed_channels.lock().unwrap();
+
+			let channels_removed_bulk = channels.remove_fetch_bulk(&scids_to_remove);
+			self.removed_node_counters.lock().unwrap().reserve(channels_removed_bulk.len());
+			let mut nodes_to_remove = hash_set_with_capacity(channels_removed_bulk.len());
+			for (scid, info) in channels_removed_bulk {
+				self.remove_channel_in_nodes_callback(&mut nodes, &info, scid, |e| {
+					nodes_to_remove.insert(*e.key());
+				});
+				removed_channels_lck.insert(scid, Some(current_time_unix));
 			}
+			nodes.remove_bulk(&nodes_to_remove);
 		}
 
 		let should_keep_tracking = |time: &mut Option<u64>| {
@@ -2633,16 +2638,17 @@ where
 		Ok(())
 	}
 
-	fn remove_channel_in_nodes(
+	fn remove_channel_in_nodes_callback<RM: FnMut(IndexedMapOccupiedEntry<NodeId, NodeInfo>)>(
 		&self, nodes: &mut IndexedMap<NodeId, NodeInfo>, chan: &ChannelInfo, short_channel_id: u64,
+		mut remove_node: RM,
 	) {
 		macro_rules! remove_from_node {
 			($node_id: expr) => {
 				if let IndexedMapEntry::Occupied(mut entry) = nodes.entry($node_id) {
 					entry.get_mut().channels.retain(|chan_id| short_channel_id != *chan_id);
 					if entry.get().channels.is_empty() {
 						self.removed_node_counters.lock().unwrap().push(entry.get().node_counter);
-						entry.remove_entry();
+						remove_node(entry);
 					}
 				} else {
 					panic!(
@@ -2655,6 +2661,14 @@ where
 		remove_from_node!(chan.node_one);
 		remove_from_node!(chan.node_two);
 	}
+
+	fn remove_channel_in_nodes(
+		&self, nodes: &mut IndexedMap<NodeId, NodeInfo>, chan: &ChannelInfo, short_channel_id: u64,
+	) {
+		self.remove_channel_in_nodes_callback(nodes, chan, short_channel_id, |e| {
+			e.remove_entry();
+		});
+	}
 }
 
 impl ReadOnlyNetworkGraph<'_> {

lightning/src/util/indexed_map.rs

@@ -72,6 +72,26 @@ impl<K: Clone + Hash + Ord, V> IndexedMap<K, V> {
 		ret
 	}
 
+	/// Removes elements with the given `keys` in bulk, returning the set of removed elements.
+	pub fn remove_fetch_bulk(&mut self, keys: &HashSet<K>) -> Vec<(K, V)> {
+		let mut res = Vec::with_capacity(keys.len());
+		for key in keys.iter() {
+			if let Some((k, v)) = self.map.remove_entry(key) {
+				res.push((k, v));
+			}
+		}
+		self.keys.retain(|k| !keys.contains(k));
+		res
+	}
+
+	/// Removes elements with the given `keys` in bulk.
+	pub fn remove_bulk(&mut self, keys: &HashSet<K>) {
+		for key in keys.iter() {
+			self.map.remove(key);
+		}
+		self.keys.retain(|k| !keys.contains(k));
+	}
+
 	/// Inserts the given `key`/`value` pair into the map, returning the element that was
 	/// previously stored at the given `key`, if one exists.
 	pub fn insert(&mut self, key: K, value: V) -> Option<V> {
@@ -210,6 +230,11 @@ impl<'a, K: Hash + Ord, V> OccupiedEntry<'a, K, V> {
 		res
 	}
 
+	/// Get a reference to the key at the position described by this entry.
+	pub fn key(&self) -> &K {
+		self.underlying_entry.key()
+	}
+
 	/// Get a reference to the value at the position described by this entry.
 	pub fn get(&self) -> &V {
 		self.underlying_entry.get()