mirror of
https://github.com/varun-r-mallya/py-libp2p.git
synced 2026-03-17 02:31:27 +00:00
Fix all modules except for security
This commit is contained in:
@ -64,7 +64,7 @@ class FloodSub(IPubsubRouter):
|
||||
:param rpc: rpc message
|
||||
"""
|
||||
# Checkpoint
|
||||
await trio.sleep(0)
|
||||
await trio.hazmat.checkpoint()
|
||||
|
||||
async def publish(self, msg_forwarder: ID, pubsub_msg: rpc_pb2.Message) -> None:
|
||||
"""
|
||||
@ -107,7 +107,7 @@ class FloodSub(IPubsubRouter):
|
||||
:param topic: topic to join
|
||||
"""
|
||||
# Checkpoint
|
||||
await trio.sleep(0)
|
||||
await trio.hazmat.checkpoint()
|
||||
|
||||
async def leave(self, topic: str) -> None:
|
||||
"""
|
||||
@ -117,7 +117,7 @@ class FloodSub(IPubsubRouter):
|
||||
:param topic: topic to leave
|
||||
"""
|
||||
# Checkpoint
|
||||
await trio.sleep(0)
|
||||
await trio.hazmat.checkpoint()
|
||||
|
||||
def _get_peers_to_send(
|
||||
self, topic_ids: Iterable[str], msg_forwarder: ID, origin: ID
|
||||
|
||||
@ -1,15 +1,18 @@
|
||||
from ast import literal_eval
|
||||
import asyncio
|
||||
import logging
|
||||
import random
|
||||
from typing import Any, Dict, Iterable, List, Sequence, Set
|
||||
|
||||
from async_service import Service
|
||||
import trio
|
||||
|
||||
from libp2p.network.stream.exceptions import StreamClosed
|
||||
from libp2p.peer.id import ID
|
||||
from libp2p.pubsub import floodsub
|
||||
from libp2p.typing import TProtocol
|
||||
from libp2p.utils import encode_varint_prefixed
|
||||
|
||||
from .exceptions import NoPubsubAttached
|
||||
from .mcache import MessageCache
|
||||
from .pb import rpc_pb2
|
||||
from .pubsub import Pubsub
|
||||
@ -20,8 +23,7 @@ PROTOCOL_ID = TProtocol("/meshsub/1.0.0")
|
||||
logger = logging.getLogger("libp2p.pubsub.gossipsub")
|
||||
|
||||
|
||||
class GossipSub(IPubsubRouter):
|
||||
|
||||
class GossipSub(IPubsubRouter, Service):
|
||||
protocols: List[TProtocol]
|
||||
pubsub: Pubsub
|
||||
|
||||
@ -86,6 +88,12 @@ class GossipSub(IPubsubRouter):
|
||||
# Create heartbeat timer
|
||||
self.heartbeat_interval = heartbeat_interval
|
||||
|
||||
async def run(self) -> None:
|
||||
if self.pubsub is None:
|
||||
raise NoPubsubAttached
|
||||
self.manager.run_task(self.heartbeat)
|
||||
await self.manager.wait_finished()
|
||||
|
||||
# Interface functions
|
||||
|
||||
def get_protocols(self) -> List[TProtocol]:
|
||||
@ -105,10 +113,6 @@ class GossipSub(IPubsubRouter):
|
||||
|
||||
logger.debug("attached to pusub")
|
||||
|
||||
# Start heartbeat now that we have a pubsub instance
|
||||
# TODO: Start after delay
|
||||
asyncio.ensure_future(self.heartbeat())
|
||||
|
||||
def add_peer(self, peer_id: ID, protocol_id: TProtocol) -> None:
|
||||
"""
|
||||
Notifies the router that a new peer has been connected.
|
||||
@ -310,7 +314,7 @@ class GossipSub(IPubsubRouter):
|
||||
await self.fanout_heartbeat()
|
||||
await self.gossip_heartbeat()
|
||||
|
||||
await asyncio.sleep(self.heartbeat_interval)
|
||||
await trio.sleep(self.heartbeat_interval)
|
||||
|
||||
async def mesh_heartbeat(self) -> None:
|
||||
# Note: the comments here are the exact pseudocode from the spec
|
||||
@ -338,7 +342,7 @@ class GossipSub(IPubsubRouter):
|
||||
|
||||
if num_mesh_peers_in_topic > self.degree_high:
|
||||
# Select |mesh[topic]| - D peers from mesh[topic]
|
||||
selected_peers = GossipSub.select_from_minus(
|
||||
selected_peers = self.select_from_minus(
|
||||
num_mesh_peers_in_topic - self.degree, self.mesh[topic], []
|
||||
)
|
||||
for peer in selected_peers:
|
||||
@ -353,7 +357,10 @@ class GossipSub(IPubsubRouter):
|
||||
for topic in self.fanout:
|
||||
# If time since last published > ttl
|
||||
# TODO: there's no way time_since_last_publish gets set anywhere yet
|
||||
if self.time_since_last_publish[topic] > self.time_to_live:
|
||||
if (
|
||||
topic in self.time_since_last_publish
|
||||
and self.time_since_last_publish[topic] > self.time_to_live
|
||||
):
|
||||
# Remove topic from fanout
|
||||
del self.fanout[topic]
|
||||
del self.time_since_last_publish[topic]
|
||||
@ -407,11 +414,7 @@ class GossipSub(IPubsubRouter):
|
||||
topic, self.degree, []
|
||||
)
|
||||
for peer in peers_to_emit_ihave_to:
|
||||
if (
|
||||
peer not in self.mesh[topic]
|
||||
and peer not in self.fanout[topic]
|
||||
):
|
||||
|
||||
if peer not in self.fanout[topic]:
|
||||
msg_id_strs = [str(msg) for msg in msg_ids]
|
||||
await self.emit_ihave(topic, msg_id_strs, peer)
|
||||
|
||||
|
||||
@ -1,4 +1,4 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from abc import ABC
|
||||
import logging
|
||||
import math
|
||||
import time
|
||||
@ -57,6 +57,7 @@ class TopicValidator(NamedTuple):
|
||||
is_async: bool
|
||||
|
||||
|
||||
# TODO: Add interface for Pubsub
|
||||
class BasePubsub(ABC):
|
||||
pass
|
||||
|
||||
@ -103,20 +104,24 @@ class Pubsub(BasePubsub, Service):
|
||||
# Attach this new Pubsub object to the router
|
||||
self.router.attach(self)
|
||||
|
||||
peer_send_channel, peer_receive_channel = trio.open_memory_channel(0)
|
||||
dead_peer_send_channel, dead_peer_receive_channel = trio.open_memory_channel(0)
|
||||
peer_channels: Tuple[
|
||||
"trio.MemorySendChannel[ID]", "trio.MemoryReceiveChannel[ID]"
|
||||
] = trio.open_memory_channel(0)
|
||||
dead_peer_channels: Tuple[
|
||||
"trio.MemorySendChannel[ID]", "trio.MemoryReceiveChannel[ID]"
|
||||
] = trio.open_memory_channel(0)
|
||||
# Only keep the receive channels in `Pubsub`.
|
||||
# Therefore, we can only close from the receive side.
|
||||
self.peer_receive_channel = peer_receive_channel
|
||||
self.dead_peer_receive_channel = dead_peer_receive_channel
|
||||
self.peer_receive_channel = peer_channels[1]
|
||||
self.dead_peer_receive_channel = dead_peer_channels[1]
|
||||
# Register a notifee
|
||||
self.host.get_network().register_notifee(
|
||||
PubsubNotifee(peer_send_channel, dead_peer_send_channel)
|
||||
PubsubNotifee(peer_channels[0], dead_peer_channels[0])
|
||||
)
|
||||
|
||||
# Register stream handlers for each pubsub router protocol to handle
|
||||
# the pubsub streams opened on those protocols
|
||||
for protocol in router.protocols:
|
||||
for protocol in router.get_protocols():
|
||||
self.host.set_stream_handler(protocol, self.stream_handler)
|
||||
|
||||
# keeps track of seen messages as LRU cache
|
||||
@ -328,8 +333,9 @@ class Pubsub(BasePubsub, Service):
|
||||
self.manager.run_task(self._handle_new_peer, peer_id)
|
||||
|
||||
async def handle_dead_peer_queue(self) -> None:
|
||||
"""Continuously read from dead peer channel and close the stream between
|
||||
that peer and remove peer info from pubsub and pubsub router."""
|
||||
"""Continuously read from dead peer channel and close the stream
|
||||
between that peer and remove peer info from pubsub and pubsub
|
||||
router."""
|
||||
async with self.dead_peer_receive_channel:
|
||||
while self.manager.is_running:
|
||||
peer_id: ID = await self.dead_peer_receive_channel.receive()
|
||||
@ -391,7 +397,11 @@ class Pubsub(BasePubsub, Service):
|
||||
return self.subscribed_topics_receive[topic_id]
|
||||
|
||||
# Map topic_id to a blocking channel
|
||||
send_channel, receive_channel = trio.open_memory_channel(math.inf)
|
||||
channels: Tuple[
|
||||
"trio.MemorySendChannel[rpc_pb2.Message]",
|
||||
"trio.MemoryReceiveChannel[rpc_pb2.Message]",
|
||||
] = trio.open_memory_channel(math.inf)
|
||||
send_channel, receive_channel = channels
|
||||
self.subscribed_topics_send[topic_id] = send_channel
|
||||
self.subscribed_topics_receive[topic_id] = receive_channel
|
||||
|
||||
@ -506,7 +516,7 @@ class Pubsub(BasePubsub, Service):
|
||||
|
||||
if len(async_topic_validators) > 0:
|
||||
# TODO: Use a better pattern
|
||||
final_result = True
|
||||
final_result: bool = True
|
||||
|
||||
async def run_async_validator(func: AsyncValidatorFn) -> None:
|
||||
nonlocal final_result
|
||||
@ -514,8 +524,8 @@ class Pubsub(BasePubsub, Service):
|
||||
final_result = final_result and result
|
||||
|
||||
async with trio.open_nursery() as nursery:
|
||||
for validator in async_topic_validators:
|
||||
nursery.start_soon(run_async_validator, validator)
|
||||
for async_validator in async_topic_validators:
|
||||
nursery.start_soon(run_async_validator, async_validator)
|
||||
|
||||
if not final_result:
|
||||
raise ValidationError(f"Validation failed for msg={msg}")
|
||||
|
||||
Reference in New Issue
Block a user