Refactor _get_peers_to_send
This commit is contained in:
parent
8e591229fd
commit
e6813da5f5
|
@ -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
|
||||||
|
|
Loading…
Reference in New Issue
Block a user