Reset connection after write stream fails (#65)

Author: mpetrun5

comm/p2p/libp2p.go

@@ -39,7 +39,7 @@ func NewCommunication(h host.Host, protocolID protocol.ID) Libp2pCommunication {
 		h:                          h,
 		protocolID:                 protocolID,
 		logger:                     logger,
-		streamManager:              NewStreamManager(),
+		streamManager:              NewStreamManager(h),
 	}
 
 	// start processing incoming messages
@@ -179,21 +179,27 @@ func (c Libp2pCommunication) sendMessage(
 	}
 
 	var stream network.Stream
-	stream, err = c.streamManager.Stream(sessionID, to)
+	stream, err = c.streamManager.Stream(sessionID, to, c.protocolID)
 	if err != nil {
-		// try to open the stream again if it failed the first time
-		stream, err = c.h.NewStream(context.TODO(), to, c.protocolID)
-		if err != nil {
-			return err
-		}
-		c.streamManager.AddStream(sessionID, to, stream)
+		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)
+
+		stream, err = c.streamManager.Stream(sessionID, to, c.protocolID)
+		if err != nil {
+			return err
+		}
+
+		err = WriteStream(msg, bufio.NewWriterSize(stream, defaultBufferSize))
+	}
+	if err != nil {
 		return err
 	}
+
 	c.logger.Trace().Str(
 		"To", to.String()).Str(
 		"MsgType", msgType.String()).Str(

comm/p2p/manager.go

@@ -4,11 +4,13 @@
 package p2p
 
 import (
-	"fmt"
+	"context"
 	"sync"
 
+	"github.com/libp2p/go-libp2p/core/host"
 	"github.com/libp2p/go-libp2p/core/network"
 	"github.com/libp2p/go-libp2p/core/peer"
+	"github.com/libp2p/go-libp2p/core/protocol"
 	"github.com/rs/zerolog/log"
 )
 
@@ -18,13 +20,15 @@ import (
 type StreamManager struct {
 	streamsBySessionID map[string]map[peer.ID]network.Stream
 	streamLocker       *sync.Mutex
+	host               host.Host
 }
 
 // NewStreamManager creates new StreamManager
-func NewStreamManager() *StreamManager {
+func NewStreamManager(host host.Host) *StreamManager {
 	return &StreamManager{
 		streamsBySessionID: make(map[string]map[peer.ID]network.Stream),
 		streamLocker:       &sync.Mutex{},
+		host:               host,
 	}
 }
 
@@ -39,6 +43,10 @@ func (sm *StreamManager) ReleaseStreams(sessionID string) {
 	}
 
 	for peer, stream := range streams {
+		if stream.Conn() != nil {
+			_ = stream.Conn().Close()
+		}
+
 		err := stream.Close()
 		if err != nil {
 			log.Err(err).Msgf("Cannot close stream to peer %s", peer.String())
@@ -67,14 +75,21 @@ func (sm *StreamManager) AddStream(sessionID string, peerID peer.ID, stream netw
 }
 
 // Stream fetches stream by peer and session ID
-func (sm *StreamManager) Stream(sessionID string, peerID peer.ID) (network.Stream, error) {
+func (sm *StreamManager) Stream(sessionID string, peerID peer.ID, protocolID protocol.ID) (network.Stream, error) {
 	sm.streamLocker.Lock()
-	defer sm.streamLocker.Unlock()
 
 	stream, ok := sm.streamsBySessionID[sessionID][peerID]
 	if !ok {
-		return nil, fmt.Errorf("no stream for peerID %s", peerID)
+		stream, err := sm.host.NewStream(context.TODO(), peerID, protocolID)
+		if err != nil {
+			return nil, err
+		}
+
+		sm.streamLocker.Unlock()
+		sm.AddStream(sessionID, peerID, stream)
+		return stream, nil
 	}
 
+	sm.streamLocker.Unlock()
 	return stream, nil
 }

comm/p2p/manager_test.go

@@ -7,8 +7,10 @@ import (
 	"testing"
 
 	"github.com/libp2p/go-libp2p/core/peer"
+	"github.com/libp2p/go-libp2p/core/protocol"
 	"github.com/sprintertech/sprinter-signing/comm/p2p"
 
+	mock_host "github.com/sprintertech/sprinter-signing/comm/p2p/mock/host"
 	mock_network "github.com/sprintertech/sprinter-signing/comm/p2p/mock/stream"
 	"github.com/stretchr/testify/suite"
 	"go.uber.org/mock/gomock"
@@ -17,6 +19,7 @@ import (
 type StreamManagerTestSuite struct {
 	suite.Suite
 	mockController *gomock.Controller
+	mockHost       *mock_host.MockHost
 }
 
 func TestRunStreamManagerTestSuite(t *testing.T) {
@@ -25,14 +28,21 @@ func TestRunStreamManagerTestSuite(t *testing.T) {
 
 func (s *StreamManagerTestSuite) SetupTest() {
 	s.mockController = gomock.NewController(s.T())
+	s.mockHost = mock_host.NewMockHost(s.mockController)
 }
 
 func (s *StreamManagerTestSuite) Test_ManagingSubscriptions_Success() {
-	streamManager := p2p.NewStreamManager()
+	streamManager := p2p.NewStreamManager(s.mockHost)
+
+	mockConn := mock_network.NewMockConn(s.mockController)
+	mockConn.EXPECT().Close().Return(nil).Times(2)
 
 	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()
 
 	peerID1, _ := peer.Decode("QmcW3oMdSqoEcjbyd51auqC23vhKX6BqfcZcY2HJ3sKAZR")
 	peerID2, _ := peer.Decode("QmZHPnN3CKiTAp8VaJqszbf8m7v4mPh15M421KpVdYHF54")
@@ -49,36 +59,40 @@ func (s *StreamManagerTestSuite) Test_ManagingSubscriptions_Success() {
 }
 
 func (s *StreamManagerTestSuite) Test_FetchStream_NoStream() {
-	streamManager := p2p.NewStreamManager()
+	streamManager := p2p.NewStreamManager(s.mockHost)
+
+	expectedStream := mock_network.NewMockStream(s.mockController)
+	s.mockHost.EXPECT().NewStream(gomock.Any(), gomock.Any(), gomock.Any()).Return(expectedStream, nil)
 
-	_, err := streamManager.Stream("1", peer.ID(""))
+	stream, err := streamManager.Stream("1", peer.ID(""), protocol.ID(""))
 
-	s.NotNil(err)
+	s.Nil(err)
+	s.Equal(stream, expectedStream)
 }
 
 func (s *StreamManagerTestSuite) Test_FetchStream_ValidStream() {
-	streamManager := p2p.NewStreamManager()
+	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)
+	expectedStream, err := streamManager.Stream("1", peerID1, protocol.ID(""))
 
 	s.Nil(err)
 	s.Equal(stream, expectedStream)
 }
 
 func (s *StreamManagerTestSuite) Test_AddStream_IgnoresExistingPeer() {
-	streamManager := p2p.NewStreamManager()
+	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)
+	expectedStream, err := streamManager.Stream("1", peerID1, protocol.ID(""))
 
 	s.Nil(err)
 	s.Equal(stream1, expectedStream)