p2p: make registries thread safer

License: MIT
Signed-off-by: Łukasz Magiera <magik6k@gmail.com>

p2p/listener.go

 package p2p
 
+import (
+	"sync"
+)
+
 type Listener interface {
 	Protocol() string
 	ListenAddress() string
@@ -18,16 +22,23 @@ type listenerKey struct {
 // ListenerRegistry is a collection of local application proto listeners.
 type ListenerRegistry struct {
 	Listeners map[listenerKey]Listener
+	lk        *sync.Mutex
 }
 
 // Register registers listenerInfo in this registry
-func (c *ListenerRegistry) Register(l Listener) {
-	c.Listeners[getListenerKey(l)] = l
+func (r *ListenerRegistry) Register(l Listener) {
+	r.lk.Lock()
+	defer r.lk.Unlock()
+
+	r.Listeners[getListenerKey(l)] = l
 }
 
 // Deregister removes p2p listener from this registry
-func (c *ListenerRegistry) Deregister(k listenerKey) {
-	delete(c.Listeners, k)
+func (r *ListenerRegistry) Deregister(k listenerKey) {
+	r.lk.Lock()
+	defer r.lk.Unlock()
+
+	delete(r.Listeners, k)
 }
 
 func getListenerKey(l Listener) listenerKey {

p2p/p2p.go

 package p2p
 
 import (
+	"sync"
+
 	pstore "gx/ipfs/QmZR2XWVVBCtbgBWnQhWk2xcQfaR3W8faQPriAiaaj7rsr/go-libp2p-peerstore"
 	p2phost "gx/ipfs/Qmb8T6YBBsjYsVGfrihQLfCJveczZnneSBqBKkYEBWDjge/go-libp2p-host"
 	peer "gx/ipfs/QmdVrMn1LhB4ybb8hMVaMLXnA8XRSewMnK6YqXKXoTcRvN/go-libp2p-peer"
@@ -25,9 +27,11 @@ func NewP2P(identity peer.ID, peerHost p2phost.Host, peerstore pstore.Peerstore)
 
 		Listeners: &ListenerRegistry{
 			Listeners: map[listenerKey]Listener{},
+			lk:        &sync.Mutex{},
 		},
 		Streams: &StreamRegistry{
 			Streams: map[uint64]*Stream{},
+			lk:      &sync.Mutex{},
 		},
 	}
 }

p2p/remote.go

@@ -38,7 +38,7 @@ func (p2p *P2P) ForwardRemote(ctx context.Context, proto string, addr ma.Multiad
 			return
 		}
 
-		stream := Stream{
+		stream := &Stream{
 			Protocol: proto,
 
 			OriginAddr: remote.Conn().RemoteMultiaddr(),
@@ -50,7 +50,7 @@ func (p2p *P2P) ForwardRemote(ctx context.Context, proto string, addr ma.Multiad
 			Registry: p2p.Streams,
 		}
 
-		p2p.Streams.Register(&stream)
+		p2p.Streams.Register(stream)
 		stream.startStreaming()
 	})
 

p2p/stream.go

@@ -2,6 +2,7 @@ package p2p
 
 import (
 	"io"
+	"sync"
 
 	ma "gx/ipfs/QmWWQ2Txc2c6tqjsBpzg5Ar652cHPGNsQQp2SejkNmkUMb/go-multiaddr"
 	net "gx/ipfs/QmYj8wdn5sZEHX2XMDWGBvcXJNdzVbaVpHmXvhHBVZepen/go-libp2p-net"
@@ -58,18 +59,25 @@ func (s *Stream) startStreaming() {
 // StreamRegistry is a collection of active incoming and outgoing proto app streams.
 type StreamRegistry struct {
 	Streams map[uint64]*Stream
+	lk      *sync.Mutex
 
 	nextId uint64
 }
 
 // Register registers a stream to the registry
-func (c *StreamRegistry) Register(streamInfo *Stream) {
-	streamInfo.Id = c.nextId
-	c.Streams[c.nextId] = streamInfo
-	c.nextId++
+func (r *StreamRegistry) Register(streamInfo *Stream) {
+	r.lk.Lock()
+	defer r.lk.Unlock()
+
+	streamInfo.Id = r.nextId
+	r.Streams[r.nextId] = streamInfo
+	r.nextId++
 }
 
 // Deregister deregisters stream from the registry
-func (c *StreamRegistry) Deregister(streamId uint64) {
-	delete(c.Streams, streamId)
+func (r *StreamRegistry) Deregister(streamId uint64) {
+	r.lk.Lock()
+	defer r.lk.Unlock()
+
+	delete(r.Streams, streamId)
 }