Minor cleanup:

- remove outdated comment
- add new peer at the end
- turn peers to send from list to set
This commit is contained in:
NIC619 2019-12-17 17:17:03 +08:00
parent ef31f7f6d6
commit 7d6daa8e10
No known key found for this signature in database
GPG Key ID: 570C35F5C2D51B17
3 changed files with 8 additions and 12 deletions

View File

@ -77,10 +77,12 @@ class FloodSub(IPubsubRouter):
:param pubsub_msg: pubsub message in protobuf. :param pubsub_msg: pubsub message in protobuf.
""" """
peers_gen = self._get_peers_to_send( peers_gen = set(
pubsub_msg.topicIDs, self._get_peers_to_send(
msg_forwarder=msg_forwarder, pubsub_msg.topicIDs,
origin=ID(pubsub_msg.from_id), msg_forwarder=msg_forwarder,
origin=ID(pubsub_msg.from_id),
)
) )
rpc_msg = rpc_pb2.RPC(publish=[pubsub_msg]) rpc_msg = rpc_pb2.RPC(publish=[pubsub_msg])

View File

@ -148,11 +148,9 @@ class GossipSub(IPubsubRouter):
for topic in self.mesh: for topic in self.mesh:
if peer_id in self.mesh[topic]: if peer_id in self.mesh[topic]:
# Delete the entry if no other peers left
self.mesh[topic].remove(peer_id) self.mesh[topic].remove(peer_id)
for topic in self.fanout: for topic in self.fanout:
if peer_id in self.fanout[topic]: if peer_id in self.fanout[topic]:
# Delete the entry if no other peers left
self.fanout[topic].remove(peer_id) self.fanout[topic].remove(peer_id)
self.peers_to_protocol.pop(peer_id, None) self.peers_to_protocol.pop(peer_id, None)

View File

@ -289,24 +289,22 @@ class Pubsub:
logger.debug("fail to add new peer %s, error %s", peer_id, error) logger.debug("fail to add new peer %s, error %s", peer_id, error)
return return
self.peers[peer_id] = stream
# Send hello packet # Send hello packet
hello = self.get_hello_packet() hello = self.get_hello_packet()
try: try:
await stream.write(encode_varint_prefixed(hello.SerializeToString())) await stream.write(encode_varint_prefixed(hello.SerializeToString()))
except StreamClosed: except StreamClosed:
logger.debug("Fail to add new peer %s: stream closed", peer_id) logger.debug("Fail to add new peer %s: stream closed", peer_id)
del self.peers[peer_id]
return return
# TODO: Check if the peer in black list. # TODO: Check if the peer in black list.
try: try:
self.router.add_peer(peer_id, stream.get_protocol()) self.router.add_peer(peer_id, stream.get_protocol())
except Exception as error: except Exception as error:
logger.debug("fail to add new peer %s, error %s", peer_id, error) logger.debug("fail to add new peer %s, error %s", peer_id, error)
del self.peers[peer_id]
return return
self.peers[peer_id] = stream
logger.debug("added new peer %s", peer_id) logger.debug("added new peer %s", peer_id)
def _handle_dead_peer(self, peer_id: ID) -> None: def _handle_dead_peer(self, peer_id: ID) -> None:
@ -316,7 +314,6 @@ class Pubsub:
for topic in self.peer_topics: for topic in self.peer_topics:
if peer_id in self.peer_topics[topic]: if peer_id in self.peer_topics[topic]:
# Delete the entry if no other peers left
self.peer_topics[topic].remove(peer_id) self.peer_topics[topic].remove(peer_id)
self.router.remove_peer(peer_id) self.router.remove_peer(peer_id)
@ -360,7 +357,6 @@ class Pubsub:
else: else:
if sub_message.topicid in self.peer_topics: if sub_message.topicid in self.peer_topics:
if origin_id in self.peer_topics[sub_message.topicid]: if origin_id in self.peer_topics[sub_message.topicid]:
# Delete the entry if no other peers left
self.peer_topics[sub_message.topicid].remove(origin_id) self.peer_topics[sub_message.topicid].remove(origin_id)
# FIXME(mhchia): Change the function name? # FIXME(mhchia): Change the function name?