From e7ac09cb9444000c6ab1f638b86c6a88f8ae1b75 Mon Sep 17 00:00:00 2001 From: NIC619 Date: Sun, 28 Jul 2019 14:52:02 +0800 Subject: [PATCH] Fix: Add Gossipsub attribute `peers_protocol` and do cleanup when peer removed --- libp2p/pubsub/gossipsub.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/libp2p/pubsub/gossipsub.py b/libp2p/pubsub/gossipsub.py index 91517a6..d7fe84a 100644 --- a/libp2p/pubsub/gossipsub.py +++ b/libp2p/pubsub/gossipsub.py @@ -41,6 +41,9 @@ class GossipSub(IPubsubRouter): # FIXME: Should be changed to `Dict[str, List[ID]]` fanout: Dict[str, List[str]] + # FIXME: Should be changed to `Dict[ID, str]` + peers_to_protocol: Dict[str, str] + time_since_last_publish: Dict[str, int] #FIXME: Should be changed to List[ID] @@ -119,6 +122,9 @@ class GossipSub(IPubsubRouter): # Add peer to the correct peer list peer_type = GossipSub.get_peer_type(protocol_id) peer_id_str = str(peer_id) + + self.peers_to_protocol[peer_id_str] = protocol_id + if peer_type == "gossip": self.peers_gossipsub.append(peer_id_str) elif peer_type == "flood": @@ -130,7 +136,12 @@ class GossipSub(IPubsubRouter): :param peer_id: id of peer to remove """ peer_id_str = str(peer_id) - self.peers_to_protocol.remove(peer_id_str) + del self.peers_to_protocol[peer_id_str] + + if peer_id_str in self.peers_gossipsub: + self.peers_gossipsub.remove(peer_id_str) + if peer_id_str in self.peers_gossipsub: + self.peers_floodsub.remove(peer_id_str) # FIXME: type of `sender_peer_id` should be changed to `ID` async def handle_rpc(self, rpc: rpc_pb2.Message, sender_peer_id: str) -> None: