mirror of
https://github.com/varun-r-mallya/py-libp2p.git
synced 2025-12-31 20:36:24 +00:00
Add pubsub test for gossipsub
This commit is contained in:
@ -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))
|
||||
|
||||
@ -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:
|
||||
|
||||
Reference in New Issue
Block a user