diff --git a/protocols/gossipsub/src/behaviour.rs b/protocols/gossipsub/src/behaviour.rs index aa9190be6fb..baacbdf12a7 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,10 @@ where // remove topic from the peer_topics mapping subscribed_topics.remove(topic_hash); + if let Some(peers) = self.topic_peers.get_mut(topic_hash) { + peers.remove(propagation_source); + }; + unsubscribed_peers.push((*propagation_source, topic_hash.clone())); // generate an unsubscribe event to be polled application_event.push(ToSwarm::GenerateEvent(Event::Unsubscribed { @@ -2080,6 +2082,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 +3343,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