diff --git a/p2p/node/p2p_services.go b/p2p/node/p2p_services.go index 3fe97d479c..74f8837164 100644 --- a/p2p/node/p2p_services.go +++ b/p2p/node/p2p_services.go @@ -56,7 +56,7 @@ func (p *P2PNode) requestFromPeer(peerID peer.ID, location common.Location, data } // Send the request to the peer - err = common.WriteMessageToStream(stream, requestBytes) + err = p.GetPeerManager().WriteMessageToStream(peerID, stream, requestBytes) if err != nil { return nil, err } @@ -69,12 +69,14 @@ func (p *P2PNode) requestFromPeer(peerID peer.ID, location common.Location, data var recvdType interface{} select { case recvdType = <-dataChan: + p.GetPeerManager().ClosePendingRequest(peerID) break case <-time.After(requestManager.C_requestTimeout): log.Global.WithFields(log.Fields{ "peerId": peerID, }).Warn("Peer did not respond in time") p.peerManager.MarkUnresponsivePeer(peerID, location) + p.GetPeerManager().ClosePendingRequest(peerID) return nil, errors.New("peer did not respond in time") } diff --git a/p2p/node/peerManager/peerManager.go b/p2p/node/peerManager/peerManager.go index b5f6019093..2aae1f1aa7 100644 --- a/p2p/node/peerManager/peerManager.go +++ b/p2p/node/peerManager/peerManager.go @@ -604,3 +604,11 @@ func (pm *BasicPeerManager) GetStream(peerID p2p.PeerID) (network.Stream, error) func (pm *BasicPeerManager) CloseStream(peerID p2p.PeerID) error { return pm.streamManager.CloseStream(peerID) } + +func (pm *BasicPeerManager) WriteMessageToStream(peerId p2p.PeerID, stream network.Stream, msg []byte) error { + return pm.streamManager.WriteMessageToStream(peerId, stream, msg) +} + +func (pm *BasicPeerManager) ClosePendingRequest(peer p2p.PeerID) error { + return pm.streamManager.ClosePendingRequest(peer) +} diff --git a/p2p/node/streamManager/streamManager.go b/p2p/node/streamManager/streamManager.go index ee7cdc3245..0297616e68 100644 --- a/p2p/node/streamManager/streamManager.go +++ b/p2p/node/streamManager/streamManager.go @@ -47,6 +47,12 @@ type StreamManager interface { // CloseStream goes through all the steps to properly close and remove a stream's resources CloseStream(peer.ID) error + + // WriteMessageToStream writes the given message into the given stream + WriteMessageToStream(peerID p2p.PeerID, stream network.Stream, msg []byte) error + + // Releases a semaphore slot for the given peerID + ClosePendingRequest(peerID p2p.PeerID) error } type basicStreamManager struct { @@ -184,5 +190,15 @@ func (sm *basicStreamManager) WriteMessageToStream(peerID p2p.PeerID, stream net if err != nil { return errors.Wrap(err, "failed to write message to stream") } + + return nil +} + +func (sm *basicStreamManager) ClosePendingRequest(peerID p2p.PeerID) error { + wrappedStream, found := sm.streamCache.Get(peerID) + if !found { + return errors.New("stream not found") + } + <-wrappedStream.(*streamWrapper).semaphore return nil }