diff --git a/lightning/src/routing/gossip.rs b/lightning/src/routing/gossip.rs index 5be09d7932b..80ffbf9fb6c 100644 --- a/lightning/src/routing/gossip.rs +++ b/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| { @@ -2633,8 +2638,9 @@ where Ok(()) } - fn remove_channel_in_nodes( + fn remove_channel_in_nodes_callback)>( &self, nodes: &mut IndexedMap, chan: &ChannelInfo, short_channel_id: u64, + mut remove_node: RM, ) { macro_rules! remove_from_node { ($node_id: expr) => { @@ -2642,7 +2648,7 @@ where 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, chan: &ChannelInfo, short_channel_id: u64, + ) { + self.remove_channel_in_nodes_callback(nodes, chan, short_channel_id, |e| { + e.remove_entry(); + }); + } } impl ReadOnlyNetworkGraph<'_> { diff --git a/lightning/src/util/indexed_map.rs b/lightning/src/util/indexed_map.rs index 3f8b3578086..f8113955370 100644 --- a/lightning/src/util/indexed_map.rs +++ b/lightning/src/util/indexed_map.rs @@ -72,6 +72,26 @@ impl IndexedMap { ret } + /// Removes elements with the given `keys` in bulk, returning the set of removed elements. + pub fn remove_fetch_bulk(&mut self, keys: &HashSet) -> 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) { + 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 { @@ -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()