From 9be9b4bbfcc73cedcf3b4ac25e92a1f5d5de82d4 Mon Sep 17 00:00:00 2001
From: NIC619 <twedusuck@gmail.com>
Date: Tue, 12 Nov 2019 18:10:41 +0800
Subject: [PATCH] Handle `StreamClosed` in pub/gossip/flood-sub

---
 libp2p/pubsub/floodsub.py  |  6 +++++-
 libp2p/pubsub/gossipsub.py |  8 +++++++-
 libp2p/pubsub/pubsub.py    | 18 +++++++++++++++---
 3 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/libp2p/pubsub/floodsub.py b/libp2p/pubsub/floodsub.py
index 15ca6b0..bac0bd7 100644
--- a/libp2p/pubsub/floodsub.py
+++ b/libp2p/pubsub/floodsub.py
@@ -1,6 +1,7 @@
 import logging
 from typing import Iterable, List, Sequence
 
+from libp2p.network.stream.exceptions import StreamClosed
 from libp2p.peer.id import ID
 from libp2p.typing import TProtocol
 from libp2p.utils import encode_varint_prefixed
@@ -89,7 +90,10 @@ class FloodSub(IPubsubRouter):
             stream = self.pubsub.peers[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
-            await stream.write(encode_varint_prefixed(rpc_msg.SerializeToString()))
+            try:
+                await stream.write(encode_varint_prefixed(rpc_msg.SerializeToString()))
+            except StreamClosed:
+                logger.debug("Fail to publish message to %s: stream closed", peer_id)
 
     async def join(self, topic: str) -> None:
         """
diff --git a/libp2p/pubsub/gossipsub.py b/libp2p/pubsub/gossipsub.py
index 864189e..564cf44 100644
--- a/libp2p/pubsub/gossipsub.py
+++ b/libp2p/pubsub/gossipsub.py
@@ -4,6 +4,7 @@ import logging
 import random
 from typing import Any, Dict, Iterable, List, Sequence, Set
 
+from libp2p.network.stream.exceptions import StreamClosed
 from libp2p.peer.id import ID
 from libp2p.pubsub import floodsub
 from libp2p.typing import TProtocol
@@ -188,7 +189,12 @@ 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(encode_varint_prefixed(rpc_msg.SerializeToString()))
+            try:
+                await stream.write(encode_varint_prefixed(rpc_msg.SerializeToString()))
+            except StreamClosed:
+                logger.debug("Fail to publish message to %s: stream closed", peer_id)
+                # TODO: also remove peer info from pubsub
+                self.remove_peer(peer_id)
 
     def _get_peers_to_send(
         self, topic_ids: Iterable[str], msg_forwarder: ID, origin: ID
diff --git a/libp2p/pubsub/pubsub.py b/libp2p/pubsub/pubsub.py
index d5e0267..8cc6539 100644
--- a/libp2p/pubsub/pubsub.py
+++ b/libp2p/pubsub/pubsub.py
@@ -20,7 +20,7 @@ from libp2p.exceptions import ParseError, ValidationError
 from libp2p.host.host_interface import IHost
 from libp2p.io.exceptions import IncompleteReadError
 from libp2p.network.exceptions import SwarmException
-from libp2p.network.stream.exceptions import StreamEOF, StreamReset
+from libp2p.network.stream.exceptions import StreamClosed, StreamEOF, StreamReset
 from libp2p.network.stream.net_stream_interface import INetStream
 from libp2p.peer.id import ID
 from libp2p.typing import TProtocol
@@ -279,13 +279,19 @@ class Pubsub:
 
         # Send hello packet
         hello = self.get_hello_packet()
-        await stream.write(encode_varint_prefixed(hello.SerializeToString()))
+        try:
+            await stream.write(encode_varint_prefixed(hello.SerializeToString()))
+        except StreamClosed:
+            logger.debug("Fail to add new peer %s: stream closed", peer_id)
+            del self.peers[peer_id]
+            return
         # TODO: Check EOF of this stream.
         # TODO: Check if the peer in black list.
         try:
             self.router.add_peer(peer_id, stream.get_protocol())
         except Exception as error:
             logger.debug("fail to add new peer %s, error %s", peer_id, error)
+            del self.peers[peer_id]
             return
 
         logger.debug("added new peer %s", peer_id)
@@ -429,7 +435,13 @@ class Pubsub:
         # Broadcast message
         for stream in self.peers.values():
             # Write message to stream
-            await stream.write(encode_varint_prefixed(raw_msg))
+            try:
+                await stream.write(encode_varint_prefixed(raw_msg))
+            except StreamClosed:
+                peer_id = stream.muxed_conn.peer_id
+                logger.debug("Fail to message peer %s: stream closed", peer_id)
+                del self.peers[peer_id]
+                self.router.remove_peer(peer_id)
 
     async def publish(self, topic_id: str, data: bytes) -> None:
         """