Fix several tests

This commit is contained in:
mhchia
2019-07-26 18:35:25 +08:00
parent 035d08b8bd
commit 65aedcb25a
7 changed files with 245 additions and 213 deletions

View File

@ -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
"""

View File

@ -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):