defer stream removal instead of doing it inline.

9e129db7 · vyzo · b080f595 · 9e129db7
Commit 9e129db7 authored Jan 12, 2021 by vyzo
Hide whitespace changes
Inline Side-by-side

Showing with 8 additions and 6 deletions

comm.go comm.go +8 -6

No files found.
--- 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
 		}