varun-r-mallya/py-libp2p
1
0
Fork 0
mirror of https://github.com/varun-r-mallya/py-libp2p.git synced 2026-02-12 16:10:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

Apply PR feedback:

Browse Source 
use defaultdict and init control message
This commit is contained in:
NIC619
2019-12-06 23:42:31 +08:00
parent db0017ddbb
commit 2d3bfc8184
1 changed files with 27 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
libp2p/pubsub/gossipsub.py
View File

@ -1,8 +1,9 @@
from ast import literal_eval from ast import literal_eval
import asyncio import asyncio
from collections import defaultdict
import logging import logging
import random import random
from typing import Any, Dict, Iterable, List, Sequence, Tuple from typing import Any, DefaultDict, Dict, Iterable, List, Sequence, Tuple
from libp2p.network.stream.exceptions import StreamClosed from libp2p.network.stream.exceptions import StreamClosed
from libp2p.peer.id import ID from libp2p.peer.id import ID
@ -319,24 +320,24 @@ class GossipSub(IPubsubRouter):
# Starting with GRAFT messages # Starting with GRAFT messages
for peer, topics in peers_to_graft.items(): for peer, topics in peers_to_graft.items():
for topic in topics: for topic in topics:
graft_msg: rpc_pb2.ControlGraft = rpc_pb2.ControlGraft() graft_msg: rpc_pb2.ControlGraft = rpc_pb2.ControlGraft(topicID=topic)
graft_msg.topicID = topic
graft_msgs.append(graft_msg) graft_msgs.append(graft_msg)
# If there are also PRUNE messages to send to this peer # If there are also PRUNE messages to send to this peer
if peer in peers_to_prune: if peer in peers_to_prune:
for topic in peers_to_prune[peer]: for topic in peers_to_prune[peer]:
prune_msg: rpc_pb2.ControlPrune = rpc_pb2.ControlPrune() prune_msg: rpc_pb2.ControlPrune = rpc_pb2.ControlPrune(
prune_msg.topicID = topic topicID=topic
)
prune_msgs.append(prune_msg) prune_msgs.append(prune_msg)
del peers_to_prune[peer] del peers_to_prune[peer]
# If there are also IHAVE messages to send to this peer # If there are also IHAVE messages to send to this peer
if peer in peers_to_gossip: if peer in peers_to_gossip:
for topic in peers_to_gossip[peer]: for topic in peers_to_gossip[peer]:
ihave_msg: rpc_pb2.ControlIHave = rpc_pb2.ControlIHave() ihave_msg: rpc_pb2.ControlIHave = rpc_pb2.ControlIHave(
ihave_msg.messageIDs.extend(peers_to_gossip[peer][topic]) messageIDs=peers_to_gossip[peer][topic], topicID=topic
ihave_msg.topicID = topic )
ihave_msgs.append(ihave_msg) ihave_msgs.append(ihave_msg)
del peers_to_gossip[peer] del peers_to_gossip[peer]
@ -347,17 +348,16 @@ class GossipSub(IPubsubRouter):
for peer, topics in peers_to_prune.items(): for peer, topics in peers_to_prune.items():
prune_msgs = [] prune_msgs = []
for topic in topics: for topic in topics:
prune_msg = rpc_pb2.ControlPrune() prune_msg = rpc_pb2.ControlPrune(topicID=topic)
prune_msg.topicID = topic
prune_msgs.append(prune_msg) prune_msgs.append(prune_msg)
# If there are also IHAVE messages to send to this peer # If there are also IHAVE messages to send to this peer
if peer in peers_to_gossip: if peer in peers_to_gossip:
ihave_msgs = [] ihave_msgs = []
for topic in peers_to_gossip[peer]: for topic in peers_to_gossip[peer]:
ihave_msg = rpc_pb2.ControlIHave() ihave_msg = rpc_pb2.ControlIHave(
ihave_msg.messageIDs.extend(peers_to_gossip[peer][topic]) messageIDs=peers_to_gossip[peer][topic], topicID=topic
ihave_msg.topicID = topic )
ihave_msgs.append(ihave_msg) ihave_msgs.append(ihave_msg)
del peers_to_gossip[peer] del peers_to_gossip[peer]
@ -368,9 +368,9 @@ class GossipSub(IPubsubRouter):
for peer in peers_to_gossip: for peer in peers_to_gossip:
ihave_msgs = [] ihave_msgs = []
for topic in peers_to_gossip[peer]: for topic in peers_to_gossip[peer]:
ihave_msg = rpc_pb2.ControlIHave() ihave_msg = rpc_pb2.ControlIHave(
ihave_msg.messageIDs.extend(peers_to_gossip[peer][topic]) messageIDs=peers_to_gossip[peer][topic], topicID=topic
ihave_msg.topicID = topic )
ihave_msgs.append(ihave_msg) ihave_msgs.append(ihave_msg)
control_msg = self.pack_control_msgs(ihave_msgs, None, None) control_msg = self.pack_control_msgs(ihave_msgs, None, None)
@ -402,9 +402,11 @@ class GossipSub(IPubsubRouter):
await asyncio.sleep(self.heartbeat_interval) await asyncio.sleep(self.heartbeat_interval)
def mesh_heartbeat(self) -> Tuple[Dict[ID, List[str]], Dict[ID, List[str]]]: def mesh_heartbeat(
peers_to_graft: Dict[ID, List[str]] = {} self
peers_to_prune: Dict[ID, List[str]] = {} ) -> Tuple[DefaultDict[ID, List[str]], DefaultDict[ID, List[str]]]:
peers_to_graft: DefaultDict[ID, List[str]] = defaultdict(list)
peers_to_prune: DefaultDict[ID, List[str]] = defaultdict(list)
for topic in self.mesh: for topic in self.mesh:
# Skip if no peers have subscribed to the topic # Skip if no peers have subscribed to the topic
if topic not in self.pubsub.peer_topics: if topic not in self.pubsub.peer_topics:
@ -422,10 +424,7 @@ class GossipSub(IPubsubRouter):
self.mesh[topic].append(peer) self.mesh[topic].append(peer)
# Emit GRAFT(topic) control message to peer # Emit GRAFT(topic) control message to peer
if peer not in peers_to_graft: peers_to_graft[peer].append(topic)
peers_to_graft[peer] = [topic]
else:
peers_to_graft[peer].append(topic)
if num_mesh_peers_in_topic > self.degree_high: if num_mesh_peers_in_topic > self.degree_high:
# Select |mesh[topic]| - D peers from mesh[topic] # Select |mesh[topic]| - D peers from mesh[topic]
@ -437,10 +436,7 @@ class GossipSub(IPubsubRouter):
self.mesh[topic].remove(peer) self.mesh[topic].remove(peer)
# Emit PRUNE(topic) control message to peer # Emit PRUNE(topic) control message to peer
if peer not in peers_to_prune: peers_to_prune[peer].append(topic)
peers_to_prune[peer] = [topic]
else:
peers_to_prune[peer].append(topic)
return peers_to_graft, peers_to_prune return peers_to_graft, peers_to_prune
def fanout_heartbeat(self) -> None: def fanout_heartbeat(self) -> None:
@ -478,8 +474,8 @@ class GossipSub(IPubsubRouter):
# Add the peers to fanout[topic] # Add the peers to fanout[topic]
self.fanout[topic].extend(selected_peers) self.fanout[topic].extend(selected_peers)
def gossip_heartbeat(self) -> Dict[ID, Dict[str, List[str]]]: def gossip_heartbeat(self) -> DefaultDict[ID, Dict[str, List[str]]]:
peers_to_gossip: Dict[ID, Dict[str, List[str]]] = {} peers_to_gossip: DefaultDict[ID, Dict[str, List[str]]] = defaultdict(dict)
for topic in self.mesh: for topic in self.mesh:
msg_ids = self.mcache.window(topic) msg_ids = self.mcache.window(topic)
if msg_ids: if msg_ids:
@ -492,10 +488,7 @@ class GossipSub(IPubsubRouter):
msg_id_strs = [str(msg_id) for msg_id in msg_ids] msg_id_strs = [str(msg_id) for msg_id in msg_ids]
for peer in peers_to_emit_ihave_to: for peer in peers_to_emit_ihave_to:
if peer not in peers_to_gossip: peers_to_gossip[peer][topic] = msg_id_strs
peers_to_gossip[peer] = {topic: msg_id_strs}
else:
peers_to_gossip[peer][topic] = msg_id_strs
# TODO: Refactor and Dedup. This section is the roughly the same as the above. # TODO: Refactor and Dedup. This section is the roughly the same as the above.
# Do the same for fanout, for all topics not already hit in mesh # Do the same for fanout, for all topics not already hit in mesh
@ -510,10 +503,7 @@ class GossipSub(IPubsubRouter):
) )
msg_id_strs = [str(msg) for msg in msg_ids] msg_id_strs = [str(msg) for msg in msg_ids]
for peer in peers_to_emit_ihave_to: for peer in peers_to_emit_ihave_to:
if peer not in peers_to_gossip: peers_to_gossip[peer][topic] = msg_id_strs
peers_to_gossip[peer] = {topic: msg_id_strs}
else:
peers_to_gossip[peer][topic] = msg_id_strs
return peers_to_gossip return peers_to_gossip
@staticmethod @staticmethod