From 772805a5d7f3faeb58b93055e3696866218cc340 Mon Sep 17 00:00:00 2001 From: onur-ozkan Date: Tue, 17 Dec 2024 10:14:33 +0000 Subject: [PATCH] fix mem-leak on topic peers Signed-off-by: onur-ozkan --- protocols/gossipsub/src/behaviour.rs | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/protocols/gossipsub/src/behaviour.rs b/protocols/gossipsub/src/behaviour.rs index aa9190be6fb..49dfae14b97 100644 --- a/protocols/gossipsub/src/behaviour.rs +++ b/protocols/gossipsub/src/behaviour.rs @@ -20,9 +20,7 @@ use std::{ cmp::{max, Ordering}, - collections::HashSet, - collections::{hash_map::Entry, VecDeque}, - collections::{BTreeSet, HashMap}, + collections::{hash_map::Entry, BTreeSet, HashMap, HashSet, VecDeque}, fmt, net::IpAddr, task::{Context, Poll}, @@ -2070,6 +2068,8 @@ where // remove topic from the peer_topics mapping subscribed_topics.remove(topic_hash); + self.topic_peers.remove(topic_hash); + unsubscribed_peers.push((*propagation_source, topic_hash.clone())); // generate an unsubscribe event to be polled application_event.push(ToSwarm::GenerateEvent(Event::Unsubscribed { @@ -2080,6 +2080,10 @@ where } if let Some(m) = self.metrics.as_mut() { + let peer_list = self + .topic_peers + .entry(topic_hash.clone()) + .or_insert_with(Default::default); m.set_topic_peers(topic_hash, peer_list.len()); } } @@ -3337,6 +3341,11 @@ where // support the protocol. self.peer_topics.remove(&peer_id); + self.topic_peers.retain(|_, peers| { + peers.remove(&peer_id); + !peers.is_empty() + }); + // If metrics are enabled, register the disconnection of a peer based on its protocol. if let Some(metrics) = self.metrics.as_mut() { let peer_kind = &self