fix: reuse streams across session IDs

Author: mpetrun5

comm/health_test.go

@@ -48,7 +48,7 @@ func (s *CommunicationHealthTestSuite) TearDownTest() {
 	}
 }
 
-func (s *CommunicationHealthTestSuite) TestCommHealth_AllPearsAvailable() {
+func (s *CommunicationHealthTestSuite) TestCommHealth_AllPeersAvailable() {
 	errors := comm.ExecuteCommHealthCheck(
 		s.testCommunications[0], peer.IDSlice{s.testHosts[1].ID(), s.testHosts[2].ID()},
 	)

comm/p2p/libp2p.go

@@ -27,7 +27,6 @@ const (
 type Libp2pCommunication struct {
 	SessionSubscriptionManager
 	h             host.Host
-	protocolID    protocol.ID
 	logger        zerolog.Logger
 	streamManager *StreamManager
 }
@@ -37,21 +36,18 @@ func NewCommunication(h host.Host, protocolID protocol.ID) Libp2pCommunication {
 	c := Libp2pCommunication{
 		SessionSubscriptionManager: NewSessionSubscriptionManager(),
 		h:                          h,
-		protocolID:                 protocolID,
 		logger:                     logger,
-		streamManager:              NewStreamManager(h),
+		streamManager:              NewStreamManager(h, protocolID),
 	}
 
 	// start processing incoming messages
-	c.h.SetStreamHandler(c.protocolID, c.StreamHandlerFunc)
+	c.h.SetStreamHandler(protocolID, c.StreamHandlerFunc)
 	return c
 }
 
 /** Communication interface methods **/
 
-func (c Libp2pCommunication) CloseSession(sessionID string) {
-	c.streamManager.ReleaseStreams(sessionID)
-}
+func (c Libp2pCommunication) CloseSession(sessionID string) {}
 
 func (c Libp2pCommunication) Broadcast(
 	peers peer.IDSlice,
@@ -179,15 +175,15 @@ func (c Libp2pCommunication) sendMessage(
 	}
 
 	var stream network.Stream
-	stream, err = c.streamManager.Stream(sessionID, to, c.protocolID)
+	stream, err = c.streamManager.Stream(to)
 	if err != nil {
 		return err
 	}
 
 	err = WriteStream(msg, bufio.NewWriterSize(stream, defaultBufferSize))
 	if err != nil {
 		c.logger.Error().Str("To", to.String()).Err(err).Msg("unable to send message")
-		c.streamManager.ReleaseStreams(sessionID)
+		c.streamManager.CloseStream(to, stream)
 		return err
 	}
 

comm/p2p/manager.go

@@ -18,75 +18,50 @@ import (
 //
 // Each stream is mapped to a specific session, by sessionID
 type StreamManager struct {
-	streamsBySessionID map[string]map[peer.ID]network.Stream
-	streamLocker       *sync.Mutex
-	host               host.Host
+	streamsByPeer map[peer.ID]network.Stream
+	streamLocker  *sync.Mutex
+	host          host.Host
+	protocolID    protocol.ID
 }
 
 // NewStreamManager creates new StreamManager
-func NewStreamManager(host host.Host) *StreamManager {
+func NewStreamManager(host host.Host, protocolID protocol.ID) *StreamManager {
 	return &StreamManager{
-		streamsBySessionID: make(map[string]map[peer.ID]network.Stream),
-		streamLocker:       &sync.Mutex{},
-		host:               host,
+		streamsByPeer: make(map[peer.ID]network.Stream),
+		streamLocker:  &sync.Mutex{},
+		host:          host,
+		protocolID:    protocolID,
 	}
 }
 
-// ReleaseStream removes reference on streams mapped to provided sessionID and closes them
-func (sm *StreamManager) ReleaseStreams(sessionID string) {
+// CloseStream closes stream to the peer
+func (sm *StreamManager) CloseStream(peerID peer.ID, stream network.Stream) {
 	sm.streamLocker.Lock()
-	defer sm.streamLocker.Unlock()
+	delete(sm.streamsByPeer, peerID)
+	sm.streamLocker.Unlock()
 
-	streams, ok := sm.streamsBySessionID[sessionID]
-	if !ok {
+	err := stream.Close()
+	if err != nil {
+		log.Warn().Err(err).Msgf("Failed to close stream")
 		return
 	}
-
-	for peer, stream := range streams {
-		err := stream.Close()
-		if err != nil {
-			log.Debug().Msgf("Cannot close stream to peer %s, err: %s", peer.String(), err.Error())
-			_ = stream.Reset()
-		}
-	}
-
-	delete(sm.streamsBySessionID, sessionID)
 }
 
-// AddStream saves and maps provided stream to sessionID
-func (sm *StreamManager) AddStream(sessionID string, peerID peer.ID, stream network.Stream) {
+// Stream fetches stream by peer
+func (sm *StreamManager) Stream(peerID peer.ID) (network.Stream, error) {
 	sm.streamLocker.Lock()
 	defer sm.streamLocker.Unlock()
 
-	_, ok := sm.streamsBySessionID[sessionID]
+	stream, ok := sm.streamsByPeer[peerID]
 	if !ok {
-		sm.streamsBySessionID[sessionID] = make(map[peer.ID]network.Stream)
-	}
-
-	_, ok = sm.streamsBySessionID[sessionID][peerID]
-	if ok {
-		return
-	}
-
-	sm.streamsBySessionID[sessionID][peerID] = stream
-}
-
-// Stream fetches stream by peer and session ID
-func (sm *StreamManager) Stream(sessionID string, peerID peer.ID, protocolID protocol.ID) (network.Stream, error) {
-	sm.streamLocker.Lock()
-
-	stream, ok := sm.streamsBySessionID[sessionID][peerID]
-	if !ok {
-		stream, err := sm.host.NewStream(context.TODO(), peerID, protocolID)
+		stream, err := sm.host.NewStream(context.TODO(), peerID, sm.protocolID)
 		if err != nil {
 			return nil, err
 		}
 
-		sm.streamLocker.Unlock()
-		sm.AddStream(sessionID, peerID, stream)
+		sm.streamsByPeer[peerID] = stream
 		return stream, nil
 	}
 
-	sm.streamLocker.Unlock()
 	return stream, nil
 }

comm/p2p/manager_test.go

@@ -32,66 +32,22 @@ func (s *StreamManagerTestSuite) SetupTest() {
 }
 
 func (s *StreamManagerTestSuite) Test_ManagingSubscriptions_Success() {
-	streamManager := p2p.NewStreamManager(s.mockHost)
+	streamManager := p2p.NewStreamManager(s.mockHost, protocol.ID("1"))
 
 	mockConn := mock_network.NewMockConn(s.mockController)
 	stream1 := mock_network.NewMockStream(s.mockController)
 	stream1.EXPECT().Conn().Return(mockConn).AnyTimes()
-	stream2 := mock_network.NewMockStream(s.mockController)
-	stream2.EXPECT().Conn().Return(mockConn).AnyTimes()
-	stream3 := mock_network.NewMockStream(s.mockController)
-	stream3.EXPECT().Conn().Return(mockConn).AnyTimes()
+	s.mockHost.EXPECT().NewStream(gomock.Any(), gomock.Any(), gomock.Any()).Return(stream1, nil)
 
 	peerID1, _ := peer.Decode("QmcW3oMdSqoEcjbyd51auqC23vhKX6BqfcZcY2HJ3sKAZR")
-	peerID2, _ := peer.Decode("QmZHPnN3CKiTAp8VaJqszbf8m7v4mPh15M421KpVdYHF54")
-
-	streamManager.AddStream("1", peerID1, stream1)
-	streamManager.AddStream("1", peerID1, stream1)
-	streamManager.AddStream("1", peerID2, stream2)
-	streamManager.AddStream("2", peerID1, stream3)
-
-	stream1.EXPECT().Close().Times(1).Return(nil)
-	stream2.EXPECT().Close().Times(1).Return(nil)
-
-	streamManager.ReleaseStreams("1")
-}
-
-func (s *StreamManagerTestSuite) Test_FetchStream_NoStream() {
-	streamManager := p2p.NewStreamManager(s.mockHost)
-
-	expectedStream := mock_network.NewMockStream(s.mockController)
-	s.mockHost.EXPECT().NewStream(gomock.Any(), gomock.Any(), gomock.Any()).Return(expectedStream, nil)
-
-	stream, err := streamManager.Stream("1", peer.ID(""), protocol.ID(""))
 
+	s1, err := streamManager.Stream(peerID1)
 	s.Nil(err)
-	s.Equal(stream, expectedStream)
-}
-
-func (s *StreamManagerTestSuite) Test_FetchStream_ValidStream() {
-	streamManager := p2p.NewStreamManager(s.mockHost)
-
-	stream := mock_network.NewMockStream(s.mockController)
-	peerID1, _ := peer.Decode("QmcW3oMdSqoEcjbyd51auqC23vhKX6BqfcZcY2HJ3sKAZR")
-	streamManager.AddStream("1", peerID1, stream)
-
-	expectedStream, err := streamManager.Stream("1", peerID1, protocol.ID(""))
-
+	s2, err := streamManager.Stream(peerID1)
 	s.Nil(err)
-	s.Equal(stream, expectedStream)
-}
-
-func (s *StreamManagerTestSuite) Test_AddStream_IgnoresExistingPeer() {
-	streamManager := p2p.NewStreamManager(s.mockHost)
 
-	stream1 := mock_network.NewMockStream(s.mockController)
-	stream2 := mock_network.NewMockStream(s.mockController)
-	peerID1, _ := peer.Decode("QmcW3oMdSqoEcjbyd51auqC23vhKX6BqfcZcY2HJ3sKAZR")
-	streamManager.AddStream("1", peerID1, stream1)
-	streamManager.AddStream("1", peerID1, stream2)
-
-	expectedStream, err := streamManager.Stream("1", peerID1, protocol.ID(""))
+	s.Equal(s1, s2)
 
-	s.Nil(err)
-	s.Equal(stream1, expectedStream)
+	stream1.EXPECT().Close().Times(1).Return(nil)
+	streamManager.CloseStream(peerID1, s1)
 }