defer stream removal instead of doing it inline.

sstanculeanu · Jan 12, 2021 · 9e129db · 9e129db
1 parent b080f59
commit 9e129db
Showing 1 changed file with 8 additions and 6 deletions.
diff --git 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
 		}