Add pubsub test for gossipsub

This commit is contained in:
mhchia
2019-09-03 16:07:44 +08:00
parent fd1f466002
commit 33dae87c35
3 changed files with 155 additions and 37 deletions

View File

@ -5,6 +5,7 @@ from typing import Any, Dict, Iterable, List, Sequence, Set
from libp2p.peer.id import ID
from libp2p.typing import TProtocol
from libp2p.utils import encode_varint_prefixed
from .mcache import MessageCache
from .pb import rpc_pb2
@ -169,7 +170,7 @@ class GossipSub(IPubsubRouter):
# 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())
await stream.write(encode_varint_prefixed(rpc_msg.SerializeToString()))
def _get_peers_to_send(
self, topic_ids: Iterable[str], msg_forwarder: ID, origin: ID
@ -276,19 +277,6 @@ class GossipSub(IPubsubRouter):
return "flood"
return "unknown"
async def deliver_messages_to_peers(
self, peers: List[ID], msg_sender: ID, origin_id: ID, serialized_packet: bytes
) -> None:
for peer_id_in_topic in peers:
# Forward to all peers that are not the
# message sender and are not the message origin
if peer_id_in_topic not in (msg_sender, origin_id):
stream = self.pubsub.peers[peer_id_in_topic]
# Publish the packet
await stream.write(serialized_packet)
# Heartbeat
async def heartbeat(self) -> None:
"""
@ -511,7 +499,7 @@ class GossipSub(IPubsubRouter):
peer_stream = self.pubsub.peers[sender_peer_id]
# 4) And write the packet to the stream
await peer_stream.write(rpc_msg)
await peer_stream.write(encode_varint_prefixed(rpc_msg))
async def handle_graft(
self, graft_msg: rpc_pb2.ControlGraft, sender_peer_id: ID
@ -603,4 +591,4 @@ class GossipSub(IPubsubRouter):
peer_stream = self.pubsub.peers[to_peer]
# Write rpc to stream
await peer_stream.write(rpc_msg)
await peer_stream.write(encode_varint_prefixed(rpc_msg))

View File

@ -72,7 +72,7 @@ class Pubsub:
topic_validators: Dict[str, TopicValidator]
# NOTE: Be sure it is increased atomically everytime.
# TODO: Be sure it is increased atomically everytime.
counter: int # uint64
def __init__(
@ -165,8 +165,6 @@ class Pubsub:
for msg in rpc_incoming.publish:
if not self._is_subscribed_to_msg(msg):
continue
# TODO(mhchia): This will block this read_stream loop until all data are pushed.
# Should investigate further if this is an issue.
asyncio.ensure_future(self.push_msg(msg_forwarder=peer_id, msg=msg))
if rpc_incoming.subscriptions: