chore: add artificial delay to peer-exchange (#1038)

waku/v2/discv5/discover.go

@@ -317,24 +317,7 @@ func (d *DiscoveryV5) Iterate(ctx context.Context, iterator enode.Iterator, onNo
 	defer iterator.Close()
 
 	peerCnt := 0
-	for {
+	for DelayedHasNext(ctx, iterator, &peerCnt) {
-
-		if !delayedHasNext(ctx, iterator) {
-			return
-		}
-
-		peerCnt++
-		if peerCnt == bucketSize { // Delay every bucketSize peers discovered
-			peerCnt = 0
-			t := time.NewTimer(delayBetweenDiscoveredPeerCnt)
-			select {
-			case <-ctx.Done():
-				return
-			case <-t.C:
-				t.Stop()
-			}
-		}
-
 		_, addresses, err := wenr.Multiaddress(iterator.Node())
 		if err != nil {
 			d.metrics.RecordError(peerInfoFailure)
@@ -364,7 +347,7 @@ func (d *DiscoveryV5) Iterate(ctx context.Context, iterator enode.Iterator, onNo
 	}
 }
 
-func delayedHasNext(ctx context.Context, iterator enode.Iterator) bool {
+func DelayedHasNext(ctx context.Context, iterator enode.Iterator, peerCnt *int) bool {
 	// Delay if .Next() is too fast
 	start := time.Now()
 	hasNext := iterator.Next()
@@ -383,6 +366,18 @@ func delayedHasNext(ctx context.Context, iterator enode.Iterator) bool {
 		}
 	}
 
+	*peerCnt++
+	if *peerCnt == bucketSize { // Delay every bucketSize peers discovered
+		*peerCnt = 0
+		t := time.NewTimer(delayBetweenDiscoveredPeerCnt)
+		select {
+		case <-ctx.Done():
+			return false
+		case <-t.C:
+			t.Stop()
+		}
+	}
+
 	return true
 }
 

waku/v2/protocol/peer_exchange/protocol.go

@@ -144,7 +144,8 @@ func (wakuPX *WakuPeerExchange) iterate(ctx context.Context) error {
 	// Closing iterator
 	defer iterator.Close()
 
-	for iterator.Next() {
+	peerCnt := 0
+	for discv5.DelayedHasNext(ctx, iterator, &peerCnt) {
 		_, addresses, err := enr.Multiaddress(iterator.Node())
 		if err != nil {
 			wakuPX.log.Error("extracting multiaddrs from enr", zap.Error(err))
@@ -183,12 +184,15 @@ func (wakuPX *WakuPeerExchange) runPeerExchangeDiscv5Loop(ctx context.Context) {
 		err := wakuPX.iterate(ctx)
 		if err != nil {
 			wakuPX.log.Debug("iterating peer exchange", zap.Error(err))
-			time.Sleep(2 * time.Second)
 		}
+
+		t := time.NewTimer(5 * time.Second)
 		select {
+		case <-t.C:
+			t.Stop()
 		case <-ctx.Done():
+			t.Stop()
 			return
-		default:
 		}
 	}
 }

waku/v2/protocol/peer_exchange/waku_peer_exchange_test.go

@@ -57,8 +57,6 @@ func TestRetrieveProvidePeerExchangePeers(t *testing.T) {
 	err = d2.Start(context.Background())
 	require.NoError(t, err)
 
-	time.Sleep(3 * time.Second) // Wait some time for peers to be discovered
-
 	// mount peer exchange
 	pxPeerConn1 := discv5.NewTestPeerDiscoverer()
 	px1, err := NewWakuPeerExchange(d1, pxPeerConn1, nil, prometheus.DefaultRegisterer, utils.Logger())
@@ -80,6 +78,8 @@ func TestRetrieveProvidePeerExchangePeers(t *testing.T) {
 	err = host3.Peerstore().AddProtocols(host1.ID(), PeerExchangeID_v20alpha1)
 	require.NoError(t, err)
 
+	time.Sleep(3 * time.Second) // Wait some time for peers to be discovered
+
 	err = px3.Request(context.Background(), 1, WithPeer(host1.ID()))
 	require.NoError(t, err)