mirror of
https://github.com/varun-r-mallya/py-libp2p.git
synced 2025-12-31 20:36:24 +00:00
Fix several tests
This commit is contained in:
@ -80,6 +80,21 @@ class FloodSub(IPubsubRouter):
|
||||
# Ref: https://github.com/libp2p/go-libp2p-pubsub/blob/master/comm.go#L107
|
||||
await stream.write(rpc_msg.SerializeToString())
|
||||
|
||||
async def join(self, topic):
|
||||
"""
|
||||
Join notifies the router that we want to receive and
|
||||
forward messages in a topic. It is invoked after the
|
||||
subscription announcement
|
||||
:param topic: topic to join
|
||||
"""
|
||||
|
||||
async def leave(self, topic):
|
||||
"""
|
||||
Leave notifies the router that we are no longer interested in a topic.
|
||||
It is invoked after the unsubscription announcement.
|
||||
:param topic: topic to leave
|
||||
"""
|
||||
|
||||
def _get_peers_to_send(
|
||||
self,
|
||||
topic_ids: Iterable[str],
|
||||
@ -102,18 +117,3 @@ class FloodSub(IPubsubRouter):
|
||||
if str(peer_id) not in self.pubsub.peers:
|
||||
continue
|
||||
yield peer_id
|
||||
|
||||
async def join(self, topic):
|
||||
"""
|
||||
Join notifies the router that we want to receive and
|
||||
forward messages in a topic. It is invoked after the
|
||||
subscription announcement
|
||||
:param topic: topic to join
|
||||
"""
|
||||
|
||||
async def leave(self, topic):
|
||||
"""
|
||||
Leave notifies the router that we are no longer interested in a topic.
|
||||
It is invoked after the unsubscription announcement.
|
||||
:param topic: topic to leave
|
||||
"""
|
||||
|
||||
@ -1,10 +1,22 @@
|
||||
import random
|
||||
import asyncio
|
||||
import random
|
||||
from typing import (
|
||||
Iterable,
|
||||
List,
|
||||
MutableSet,
|
||||
Sequence,
|
||||
)
|
||||
|
||||
from ast import literal_eval
|
||||
|
||||
from libp2p.peer.id import (
|
||||
ID,
|
||||
id_b58_decode,
|
||||
)
|
||||
|
||||
from .mcache import MessageCache
|
||||
from .pb import rpc_pb2
|
||||
from .pubsub_router_interface import IPubsubRouter
|
||||
from .mcache import MessageCache
|
||||
|
||||
|
||||
class GossipSub(IPubsubRouter):
|
||||
@ -107,70 +119,73 @@ class GossipSub(IPubsubRouter):
|
||||
for prune in control_message.prune:
|
||||
await self.handle_prune(prune, sender_peer_id)
|
||||
|
||||
async def publish(self, sender_peer_id, rpc_message):
|
||||
async def publish(self, src: ID, pubsub_msg: rpc_pb2.Message) -> None:
|
||||
# pylint: disable=too-many-locals
|
||||
"""
|
||||
Invoked to forward a new message that has been validated.
|
||||
"""
|
||||
self.mcache.put(pubsub_msg)
|
||||
|
||||
packet = rpc_pb2.RPC()
|
||||
packet.ParseFromString(rpc_message)
|
||||
msg_sender = str(sender_peer_id)
|
||||
peers_gen = self._get_peers_to_send(
|
||||
pubsub_msg.topicIDs,
|
||||
src=src,
|
||||
origin=ID(pubsub_msg.from_id),
|
||||
)
|
||||
rpc_msg = rpc_pb2.RPC(
|
||||
publish=[pubsub_msg],
|
||||
)
|
||||
for peer_id in peers_gen:
|
||||
stream = self.pubsub.peers[str(peer_id)]
|
||||
# FIXME: We should add a `WriteMsg` similar to write delimited messages.
|
||||
# Ref: https://github.com/libp2p/go-libp2p-pubsub/blob/master/comm.go#L107
|
||||
# TODO: Go use `sendRPC`, which possibly piggybacks gossip/control messages.
|
||||
await stream.write(rpc_msg.SerializeToString())
|
||||
|
||||
# Deliver to self if self was origin
|
||||
# Note: handle_talk checks if self is subscribed to topics in message
|
||||
for message in packet.publish:
|
||||
# Add RPC message to cache
|
||||
self.mcache.put(message)
|
||||
def _get_peers_to_send(
|
||||
self,
|
||||
topic_ids: Iterable[str],
|
||||
src: ID,
|
||||
origin: ID) -> Iterable[ID]:
|
||||
"""
|
||||
Get the eligible peers to send the data to.
|
||||
:param src: the peer id of the peer who forwards the message to me.
|
||||
:param origin: the peer id of the peer who originally broadcast the message.
|
||||
:return: a generator of the peer ids who we send data to.
|
||||
"""
|
||||
to_send: MutableSet[ID] = set()
|
||||
for topic in topic_ids:
|
||||
if topic not in self.pubsub.peer_topics:
|
||||
continue
|
||||
|
||||
decoded_from_id = message.from_id.decode('utf-8')
|
||||
new_packet = rpc_pb2.RPC()
|
||||
new_packet.publish.extend([message])
|
||||
new_packet_serialized = new_packet.SerializeToString()
|
||||
# floodsub peers
|
||||
for peer_id_str in self.pubsub.peer_topics[topic]:
|
||||
peer_id = id_b58_decode(peer_id_str)
|
||||
# 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.
|
||||
if peer_id_str in self.peers_floodsub:
|
||||
to_send.add(peer_id)
|
||||
|
||||
# Deliver to self if needed
|
||||
if msg_sender == decoded_from_id and msg_sender == str(self.pubsub.host.get_id()):
|
||||
id_in_seen_msgs = (message.seqno, message.from_id)
|
||||
# gossipsub peers
|
||||
# FIXME: Change `str` to `ID`
|
||||
gossipsub_peers: List[str] = None
|
||||
# TODO: Do we need to check `topic in self.pubsub.my_topics`?
|
||||
if topic in self.mesh:
|
||||
gossipsub_peers = self.mesh[topic]
|
||||
else:
|
||||
# TODO(robzajac): Is topic DEFINITELY supposed to be in fanout if we are not
|
||||
# subscribed?
|
||||
# I assume there could be short periods between heartbeats where topic may not
|
||||
# be but we should check that this path gets hit appropriately
|
||||
|
||||
if id_in_seen_msgs not in self.pubsub.seen_messages:
|
||||
self.pubsub.seen_messages[id_in_seen_msgs] = 1
|
||||
|
||||
await self.pubsub.handle_talk(message)
|
||||
|
||||
# Deliver to peers
|
||||
for topic in message.topicIDs:
|
||||
# If topic has floodsub peers, deliver to floodsub peers
|
||||
# TODO: This can be done more efficiently. Do it more efficiently.
|
||||
floodsub_peers_in_topic = []
|
||||
if topic in self.pubsub.peer_topics:
|
||||
for peer in self.pubsub.peer_topics[topic]:
|
||||
if str(peer) in self.peers_floodsub:
|
||||
floodsub_peers_in_topic.append(peer)
|
||||
|
||||
await self.deliver_messages_to_peers(floodsub_peers_in_topic, msg_sender,
|
||||
decoded_from_id, new_packet_serialized)
|
||||
|
||||
# If you are subscribed to topic, send to mesh, otherwise send to fanout
|
||||
if topic in self.pubsub.my_topics and topic in self.mesh:
|
||||
await self.deliver_messages_to_peers(self.mesh[topic], msg_sender,
|
||||
decoded_from_id, new_packet_serialized)
|
||||
else:
|
||||
# Send to fanout peers
|
||||
if topic not in self.fanout:
|
||||
# If no peers in fanout, choose some peers from gossipsub peers in topic
|
||||
gossipsub_peers_in_topic = [peer for peer in self.pubsub.peer_topics[topic]
|
||||
if peer in self.peers_gossipsub]
|
||||
|
||||
selected = \
|
||||
GossipSub.select_from_minus(self.degree, gossipsub_peers_in_topic, [])
|
||||
self.fanout[topic] = selected
|
||||
|
||||
# TODO: Is topic DEFINITELY supposed to be in fanout if we are not subscribed?
|
||||
# I assume there could be short periods between heartbeats where topic may not
|
||||
# be but we should check that this path gets hit appropriately
|
||||
|
||||
await self.deliver_messages_to_peers(self.fanout[topic], msg_sender,
|
||||
decoded_from_id, new_packet_serialized)
|
||||
# pylint: disable=len-as-condition
|
||||
if (topic not in self.fanout) or (len(self.fanout[topic]) == 0):
|
||||
# If no peers in fanout, choose some peers from gossipsub peers in topic.
|
||||
self.fanout[topic] = self._get_peers_from_minus(topic, self.degree, [])
|
||||
gossipsub_peers = self.fanout[topic]
|
||||
for peer_id_str in gossipsub_peers:
|
||||
to_send.add(id_b58_decode(peer_id_str))
|
||||
# Excludes `src` and `origin`
|
||||
yield from to_send.difference([src, origin])
|
||||
|
||||
async def join(self, topic):
|
||||
# Note: the comments here are the near-exact algorithm description from the spec
|
||||
@ -401,6 +416,22 @@ class GossipSub(IPubsubRouter):
|
||||
|
||||
return selection
|
||||
|
||||
def _get_peers_from_minus(
|
||||
self,
|
||||
topic: str,
|
||||
num_to_select: int,
|
||||
minus: Sequence[ID]) -> List[ID]:
|
||||
gossipsub_peers_in_topic = [
|
||||
peer_str
|
||||
for peer_str in self.pubsub.peer_topics[topic]
|
||||
if peer_str in self.peers_gossipsub
|
||||
]
|
||||
return self.select_from_minus(
|
||||
num_to_select,
|
||||
gossipsub_peers_in_topic,
|
||||
list(minus),
|
||||
)
|
||||
|
||||
# RPC handlers
|
||||
|
||||
async def handle_ihave(self, ihave_msg, sender_peer_id):
|
||||
|
||||
Reference in New Issue
Block a user