From 5e2af9feb680b0d96d6321fd4f3bc83432d8083c Mon Sep 17 00:00:00 2001
From: Manu NALEPA <enalepa@offchainlabs.com>
Date: Tue, 3 Sep 2024 12:59:17 +0200
Subject: [PATCH] `FindPeersWithSubnet`: Address Nishant's comment.

---
 beacon-chain/p2p/subnets.go | 42 ++++++++++++++++++++++---------------
 1 file changed, 25 insertions(+), 17 deletions(-)

diff --git a/beacon-chain/p2p/subnets.go b/beacon-chain/p2p/subnets.go
index 9870c9204924..31f5b7c31f16 100644
--- a/beacon-chain/p2p/subnets.go
+++ b/beacon-chain/p2p/subnets.go
@@ -162,29 +162,35 @@ func (s *Service) FindPeersWithSubnet(
 		return false, errors.Wrap(err, "node filter")
 	}
 
-	firstLoop := true
-
-	wg := new(sync.WaitGroup)
-	for {
+	peersSummary := func(topic string, threshold int) (int, int) {
 		// Retrieve how many peers we have for this topic.
 		peerCountForTopic := len(s.pubsub.ListPeers(topic))
 
 		// Compute how many peers we are missing to reach the threshold.
-		missingPeersCount := max(0, threshold-peerCountForTopic)
+		missingPeerCountForTopic := max(0, threshold-peerCountForTopic)
 
-		// If we have enough peers, we can exit the loop. This is the happy path.
-		if missingPeersCount == 0 {
-			break
-		}
+		return peerCountForTopic, missingPeerCountForTopic
+	}
 
-		if firstLoop {
-			firstLoop = false
+	// Compute how many peers we are missing to reach the threshold.
+	peerCountForTopic, missingPeerCountForTopic := peersSummary(topic, threshold)
 
-			log.WithFields(logrus.Fields{
-				"topic":            topic,
-				"currentPeerCount": peerCountForTopic,
-				"targetPeerCount":  threshold,
-			}).Debug("Searching for new peers in the network - Start")
+	// Exit early if we have enough peers.
+	if missingPeerCountForTopic == 0 {
+		return true, nil
+	}
+
+	log.WithFields(logrus.Fields{
+		"topic":            topic,
+		"currentPeerCount": peerCountForTopic,
+		"targetPeerCount":  threshold,
+	}).Debug("Searching for new peers in the network - Start")
+
+	wg := new(sync.WaitGroup)
+	for {
+		// If we have enough peers, we can exit the loop. This is the happy path.
+		if missingPeerCountForTopic == 0 {
+			break
 		}
 
 		// If the context is done, we can exit the loop. This is the unhappy path.
@@ -196,7 +202,7 @@ func (s *Service) FindPeersWithSubnet(
 		}
 
 		// Search for new peers in the network.
-		nodes := searchForPeers(iterator, batchSize, missingPeersCount, filter)
+		nodes := searchForPeers(iterator, batchSize, missingPeerCountForTopic, filter)
 
 		// Restrict dials if limit is applied.
 		maxConcurrentDials := math.MaxInt
@@ -214,6 +220,8 @@ func (s *Service) FindPeersWithSubnet(
 			// Wait for all dials to be completed.
 			wg.Wait()
 		}
+
+		_, missingPeerCountForTopic = peersSummary(topic, threshold)
 	}
 
 	log.WithField("topic", topic).Debug("Searching for new peers in the network - Success")