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 }