Refactor _get_peers_to_send

This commit is contained in:
NIC619 2019-12-05 14:35:34 +08:00
parent 8e591229fd
commit e6813da5f5
No known key found for this signature in database
GPG Key ID: 570C35F5C2D51B17

View File

@ -2,7 +2,7 @@ from ast import literal_eval
import asyncio import asyncio
import logging import logging
import random import random
from typing import Any, Dict, Iterable, List, Sequence, Set, Tuple from typing import Any, 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
@ -216,38 +216,40 @@ class GossipSub(IPubsubRouter):
:param origin: peer id of the peer the message originate from. :param origin: peer id of the peer the message originate from.
:return: a generator of the peer ids who we send data to. :return: a generator of the peer ids who we send data to.
""" """
send_to: Set[ID] = set()
for topic in topic_ids: for topic in topic_ids:
if topic not in self.pubsub.peer_topics: if topic not in self.pubsub.peer_topics:
continue continue
# floodsub peers # floodsub peers
for peer_id in self.pubsub.peer_topics[topic]: # FIXME: `gossipsub.peers_floodsub` can be changed to `gossipsub.peers` in go.
# FIXME: `gossipsub.peers_floodsub` can be changed to `gossipsub.peers` in go. # This will improve the efficiency when searching for a peer's protocol id.
# This will improve the efficiency when searching for a peer's protocol id. floodsub_peers: List[ID] = [
if peer_id in self.peers_floodsub: peer_id
send_to.add(peer_id) for peer_id in self.pubsub.peer_topics[topic]
if peer_id in self.peers_floodsub
]
# gossipsub peers # gossipsub peers
in_topic_gossipsub_peers: List[ID] = None gossipsub_peers: List[ID] = []
if topic in self.mesh: if topic in self.mesh:
in_topic_gossipsub_peers = self.mesh[topic] gossipsub_peers = self.mesh[topic]
else: else:
# It could be the case that we publish to a topic that we have not subscribe # When we publish to a topic that we have not subscribe to, we randomly pick
# and the topic is not yet added to our `fanout`. # `self.degree` number of peers who have subscribe to the topic and add them
if (topic not in self.fanout) or (len(self.fanout[topic]) == 0): # as our `fanout` peers.
# If no peers in fanout, choose some peers from gossipsub peers in topic. topic_in_fanout: bool = topic in self.fanout
self.fanout[topic] = self._get_in_topic_gossipsub_peers_from_minus( fanout_peers: List[ID] = self.fanout[topic] if topic_in_fanout else []
topic, self.degree, [] fanout_size = len(fanout_peers)
if not topic_in_fanout or (topic_in_fanout and fanout_size < self.degree):
# Combine fanout peers with selected peers
self.fanout[topic] += self._get_in_topic_gossipsub_peers_from_minus(
topic, self.degree - fanout_size, fanout_peers
) )
in_topic_gossipsub_peers = self.fanout[topic] gossipsub_peers = fanout_peers
for peer_id in in_topic_gossipsub_peers:
send_to.add(peer_id)
# Excludes `msg_forwarder` and `origin` # Excludes `msg_forwarder` and `origin`
yield from send_to.difference([msg_forwarder, origin]) yield from set(floodsub_peers + gossipsub_peers).difference([msg_forwarder, origin])
async def join(self, topic: str) -> None: async def join(self, topic: str) -> None:
# Note: the comments here are the near-exact algorithm description from the spec
""" """
Join notifies the router that we want to receive and forward messages Join notifies the router that we want to receive and forward messages
in a topic. It is invoked after the subscription announcement. in a topic. It is invoked after the subscription announcement.
@ -277,9 +279,8 @@ class GossipSub(IPubsubRouter):
# Add fanout peers to mesh and notifies them with a GRAFT(topic) control message. # Add fanout peers to mesh and notifies them with a GRAFT(topic) control message.
for peer in fanout_peers: for peer in fanout_peers:
if peer not in self.mesh[topic]: self.mesh[topic].append(peer)
self.mesh[topic].append(peer) await self.emit_graft(topic, peer)
await self.emit_graft(topic, peer)
self.fanout.pop(topic, None) self.fanout.pop(topic, None)
@ -300,7 +301,7 @@ class GossipSub(IPubsubRouter):
await self.emit_prune(topic, peer) await self.emit_prune(topic, peer)
# Forget mesh[topic] # Forget mesh[topic]
del self.mesh[topic] self.mesh.pop(topic, None)
async def _emit_control_msgs( async def _emit_control_msgs(
self, self,
@ -452,7 +453,7 @@ class GossipSub(IPubsubRouter):
del self.fanout[topic] del self.fanout[topic]
del self.time_since_last_publish[topic] del self.time_since_last_publish[topic]
else: else:
# Check whether our peers are still in the topic # Check if fanout peers are still in the topic and remove the ones that are not
# ref: https://github.com/libp2p/go-libp2p-pubsub/blob/01b9825fbee1848751d90a8469e3f5f43bac8466/gossipsub.go#L498-L504 # noqa: E501 # ref: https://github.com/libp2p/go-libp2p-pubsub/blob/01b9825fbee1848751d90a8469e3f5f43bac8466/gossipsub.go#L498-L504 # noqa: E501
in_topic_fanout_peers = [ in_topic_fanout_peers = [
peer peer