Using the rate limiter WriteMessageToStream method

dominant-strategies · May 16, 2024 · 716bf96 · 716bf96
1 parent e8d6196
commit 716bf96
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 1 deletion.
diff --git 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
 	}

diff --git a/p2p/node/peerManager/peerManager.go b/p2p/node/peerManager/peerManager.go
@@ -604,3 +604,7 @@ 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)
+}
diff --git a/p2p/node/streamManager/streamManager.go b/p2p/node/streamManager/streamManager.go
@@ -47,6 +47,9 @@ 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
 }
 
 type basicStreamManager struct {