From 9e129db7661f8da6d605466a2dc1671b1ea33a64 Mon Sep 17 00:00:00 2001
From: vyzo <vyzo@hackzen.org>
Date: Tue, 12 Jan 2021 16:26:38 +0200
Subject: [PATCH] defer stream removal instead of doing it inline.

---
 comm.go | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/comm.go b/comm.go
index aefeed3..29f3cea 100644
--- a/comm.go
+++ b/comm.go
@@ -52,6 +52,14 @@ func (p *PubSub) handleNewStream(s network.Stream) {
 	p.inboundStreams[peer] = s
 	p.inboundStreamsMx.Unlock()
 
+	defer func() {
+		p.inboundStreamsMx.Lock()
+		if p.inboundStreams[peer] == s {
+			delete(p.inboundStreams, peer)
+		}
+		p.inboundStreamsMx.Unlock()
+	}()
+
 	r := protoio.NewDelimitedReader(s, p.maxMessageSize)
 	for {
 		rpc := new(RPC)
@@ -66,12 +74,6 @@ func (p *PubSub) handleNewStream(s network.Stream) {
 				s.Close()
 			}
 
-			p.inboundStreamsMx.Lock()
-			if p.inboundStreams[peer] == s {
-				delete(p.inboundStreams, peer)
-			}
-			p.inboundStreamsMx.Unlock()
-
 			return
 		}