Minor StreamMgr improvements

Author: swift1337

keysign/signature_notifier.go

@@ -42,21 +42,19 @@ type SignatureNotifier struct {
 
 // NewSignatureNotifier create a new instance of SignatureNotifier
 func NewSignatureNotifier(host host.Host, logger zerolog.Logger) *SignatureNotifier {
-	logger = logger.With().
-		Str(logs.Component, "signature_notifier").
-		Stringer(logs.Host, host.ID()).
-		Logger()
-
 	ctx, cancel := context.WithCancel(context.Background())
 
 	s := &SignatureNotifier{
 		ctx:          ctx,
 		cancel:       cancel,
-		logger:       logger,
 		host:         host,
 		notifierLock: &sync.Mutex{},
 		notifiers:    make(map[string]*notifier),
-		streamMgr:    p2p.NewStreamMgr(),
+		streamMgr:    p2p.NewStreamMgr(logger),
+		logger: logger.With().
+			Str(logs.Component, "signature_notifier").
+			Stringer(logs.Host, host.ID()).
+			Logger(),
 	}
 
 	host.SetStreamHandler(signatureNotifierProtocol, s.handleStream)

p2p/communication.go

@@ -95,7 +95,7 @@ func NewCommunication(
 		streamCount:      0,
 		BroadcastMsgChan: make(chan *messages.BroadcastMsgChan, 1024),
 		externalAddr:     externalAddr,
-		streamMgr:        NewStreamMgr(),
+		streamMgr:        NewStreamMgr(logger),
 		whitelistedPeers: whitelistedPeers,
 	}, nil
 }

p2p/party_coordinator.go

@@ -45,16 +45,19 @@ func NewPartyCoordinator(host host.Host, timeout time.Duration, logger zerolog.L
 	if timeout.Nanoseconds() == 0 {
 		timeout = 10 * time.Second
 	}
+
 	pc := &PartyCoordinator{
 		logger:             logger,
 		host:               host,
 		stopChan:           make(chan struct{}),
 		timeout:            timeout,
 		peersGroup:         make(map[string]*peerStatus),
 		joinPartyGroupLock: &sync.Mutex{},
-		streamMgr:          NewStreamMgr(),
+		streamMgr:          NewStreamMgr(logger),
 	}
+
 	host.SetStreamHandler(joinPartyProtocolWithLeader, pc.HandleStreamWithLeader)
+
 	return pc
 }
 

p2p/stream_helper.go

@@ -12,7 +12,8 @@ import (
 	"github.com/libp2p/go-libp2p/core/network"
 	"github.com/pkg/errors"
 	"github.com/rs/zerolog"
-	"github.com/rs/zerolog/log"
+
+	"github.com/zeta-chain/go-tss/logs"
 )
 
 const (
@@ -32,52 +33,60 @@ func init() {
 
 type StreamMgr struct {
 	unusedStreams map[string][]network.Stream
-	streamLocker  *sync.RWMutex
+	mu            *sync.RWMutex
 	logger        zerolog.Logger
 }
 
-func NewStreamMgr() *StreamMgr {
+func NewStreamMgr(logger zerolog.Logger) *StreamMgr {
 	return &StreamMgr{
 		unusedStreams: make(map[string][]network.Stream),
-		streamLocker:  &sync.RWMutex{},
-		logger:        log.With().Str("module", "communication").Logger(),
+		mu:            &sync.RWMutex{},
+		logger:        logger.With().Str(logs.Component, "stream_manager").Logger(),
 	}
 }
 
 func (sm *StreamMgr) ReleaseStream(msgID string) {
-	sm.streamLocker.RLock()
+	sm.mu.RLock()
 	usedStreams, okStream := sm.unusedStreams[msgID]
 	unknownStreams, okUnknown := sm.unusedStreams["UNKNOWN"]
 	streams := append(usedStreams, unknownStreams...)
-	sm.streamLocker.RUnlock()
-	if okStream || okUnknown {
-		for _, el := range streams {
-			err := el.Reset()
-			if err != nil {
-				sm.logger.Error().Err(err).Msg("fail to reset the stream,skip it")
-			}
+	sm.mu.RUnlock()
+
+	// noop
+	if !(okStream || okUnknown) {
+		return
+	}
+
+	for _, stream := range streams {
+		if err := stream.Reset(); err != nil {
+			sm.logger.Error().Err(err).
+				Str(logs.MsgID, msgID).
+				Str("stream_id", stream.ID()).
+				Msg("Failed to reset the stream")
 		}
-		sm.streamLocker.Lock()
-		delete(sm.unusedStreams, msgID)
-		delete(sm.unusedStreams, "UNKNOWN")
-		sm.streamLocker.Unlock()
 	}
+
+	sm.mu.Lock()
+	delete(sm.unusedStreams, msgID)
+	delete(sm.unusedStreams, "UNKNOWN")
+	sm.mu.Unlock()
 }
 
 func (sm *StreamMgr) AddStream(msgID string, stream network.Stream) {
 	if stream == nil {
 		return
 	}
-	sm.streamLocker.Lock()
-	defer sm.streamLocker.Unlock()
+
+	sm.mu.Lock()
+	defer sm.mu.Unlock()
+
 	entries, ok := sm.unusedStreams[msgID]
 	if !ok {
-		entries := []network.Stream{stream}
-		sm.unusedStreams[msgID] = entries
-	} else {
-		entries = append(entries, stream)
-		sm.unusedStreams[msgID] = entries
+		sm.unusedStreams[msgID] = []network.Stream{stream}
+		return
 	}
+
+	sm.unusedStreams[msgID] = append(entries, stream)
 }
 
 // ReadStreamWithBuffer read data from the given stream

p2p/stream_helper_test.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/pkg/errors"
+	"github.com/rs/zerolog"
 
 	"github.com/libp2p/go-libp2p/core/network"
 	"github.com/libp2p/go-libp2p/core/protocol"
@@ -234,7 +235,7 @@ func TestReadPayload(t *testing.T) {
 }
 
 func TestStreamManager(t *testing.T) {
-	streamMgr := NewStreamMgr()
+	streamMgr := NewStreamMgr(zerolog.New(zerolog.NewTestWriter(t)))
 	stream := NewMockNetworkStream()
 
 	streamMgr.AddStream("1", nil)