p2p/discover: optimize findnodeByID (#33348)

This PR adds an optimization to the `findNodeByID` function in
`p2p/discover`. There is already an open PR (#33205) for similar
improvements, and I have further optimized the function to get better
performance. I have attached the benchmark results comparing the current
`main` branch with my `optimized version`, and the results show clear
improvements.

---------

Co-authored-by: Csaba Kiraly <csaba.kiraly@gmail.com>

p2p/discover/table.go

@@ -290,28 +290,33 @@ func (tab *Table) refresh() <-chan struct{} {
 // preferLive is true and the table contains any verified nodes, the result will not
 // contain unverified nodes. However, if there are no verified nodes at all, the result
 // will contain unverified nodes.
-func (tab *Table) findnodeByID(target enode.ID, nresults int, preferLive bool) *nodesByDistance {
+func (tab *Table) findnodeByID(target enode.ID, nresults int, preferLive bool) (nodes nodesByDistance) {
+	nodes.target = target
 	tab.mutex.Lock()
 	defer tab.mutex.Unlock()
 
 	// Scan all buckets. There might be a better way to do this, but there aren't that many
 	// buckets, so this solution should be fine. The worst-case complexity of this loop
 	// is O(tab.len() * nresults).
-	nodes := &nodesByDistance{target: target}
-	liveNodes := &nodesByDistance{target: target}
-	for _, b := range &tab.buckets {
-		for _, n := range b.entries {
-			nodes.push(n.Node, nresults)
-			if preferLive && n.isValidatedLive {
-				liveNodes.push(n.Node, nresults)
+	if preferLive {
+		for _, b := range &tab.buckets {
+			for _, n := range b.entries {
+				if n.isValidatedLive {
+					nodes.push(n.Node, nresults)
+				}
 			}
 		}
+		if len(nodes.entries) > 0 {
+			return
+		}
 	}
 
-	if preferLive && len(liveNodes.entries) > 0 {
-		return liveNodes
+	for _, b := range &tab.buckets {
+		for _, n := range b.entries {
+			nodes.push(n.Node, nresults)
+		}
 	}
-	return nodes
+	return
 }
 
 // appendBucketNodes adds nodes at the given distance to the result slice.

p2p/discover/table_test.go

@@ -597,3 +597,77 @@ func newkey() *ecdsa.PrivateKey {
 	}
 	return key
 }
+
+// BenchmarkTable_findnodeByID exercises findnodeByID across table sizes, result
+// counts, and liveness ratios. The _PreferLive_NoLive cases cover the fallback
+// path where preferLive is requested but no validated-live nodes exist.
+func BenchmarkTable_findnodeByID(b *testing.B) {
+	benchmarks := []struct {
+		name       string
+		tableSize  int
+		nresults   int
+		preferLive bool
+		liveRatio  float64 // fraction of nodes marked validated-live
+	}{
+		{"SmallTable_5Results_NoPreferLive", 50, 5, false, 0.0},
+		{"SmallTable_16Results_NoPreferLive", 50, 16, false, 0.0},
+		{"SmallTable_5Results_PreferLive_AllLive", 50, 5, true, 1.0},
+		{"SmallTable_16Results_PreferLive_AllLive", 50, 16, true, 1.0},
+		{"SmallTable_5Results_PreferLive_HalfLive", 50, 5, true, 0.5},
+		{"SmallTable_16Results_PreferLive_HalfLive", 50, 16, true, 0.5},
+		{"SmallTable_5Results_PreferLive_NoLive", 50, 5, true, 0.0},
+		{"SmallTable_16Results_PreferLive_NoLive", 50, 16, true, 0.0},
+
+		{"MediumTable_5Results_NoPreferLive", 200, 5, false, 0.0},
+		{"MediumTable_16Results_NoPreferLive", 200, 16, false, 0.0},
+		{"MediumTable_5Results_PreferLive_AllLive", 200, 5, true, 1.0},
+		{"MediumTable_16Results_PreferLive_AllLive", 200, 16, true, 1.0},
+		{"MediumTable_5Results_PreferLive_HalfLive", 200, 5, true, 0.5},
+		{"MediumTable_16Results_PreferLive_HalfLive", 200, 16, true, 0.5},
+		{"MediumTable_5Results_PreferLive_NoLive", 200, 5, true, 0.0},
+		{"MediumTable_16Results_PreferLive_NoLive", 200, 16, true, 0.0},
+
+		{"FullTable_5Results_NoPreferLive", nBuckets * bucketSize, 5, false, 0.0},
+		{"FullTable_16Results_NoPreferLive", nBuckets * bucketSize, 16, false, 0.0},
+		{"FullTable_5Results_PreferLive_AllLive", nBuckets * bucketSize, 5, true, 1.0},
+		{"FullTable_16Results_PreferLive_AllLive", nBuckets * bucketSize, 16, true, 1.0},
+		{"FullTable_5Results_PreferLive_HalfLive", nBuckets * bucketSize, 5, true, 0.5},
+		{"FullTable_16Results_PreferLive_HalfLive", nBuckets * bucketSize, 16, true, 0.5},
+		{"FullTable_5Results_PreferLive_NoLive", nBuckets * bucketSize, 5, true, 0.0},
+		{"FullTable_16Results_PreferLive_NoLive", nBuckets * bucketSize, 16, true, 0.0},
+	}
+
+	for _, bm := range benchmarks {
+		b.Run(bm.name, func(b *testing.B) {
+			tab, db := newTestTable(newPingRecorder(), Config{})
+			defer db.Close()
+			defer tab.close()
+			<-tab.initDone
+
+			self := tab.self().ID()
+			nodes := make([]*enode.Node, bm.tableSize)
+			for i := range nodes {
+				// Spread across buckets by varying log-distance in the valid range.
+				d := bucketMinDistance + 1 + i%nBuckets
+				nodes[i] = nodeAtDistance(self, d, intIP(i))
+			}
+
+			liveCount := int(float64(len(nodes)) * bm.liveRatio)
+			if liveCount > 0 {
+				fillTable(tab, nodes[:liveCount], true)
+			}
+			if liveCount < len(nodes) {
+				fillTable(tab, nodes[liveCount:], false)
+			}
+
+			var target enode.ID
+			rand.Read(target[:])
+
+			b.ResetTimer()
+			b.ReportAllocs()
+			for i := 0; i < b.N; i++ {
+				_ = tab.findnodeByID(target, bm.nresults, bm.preferLive)
+			}
+		})
+	}
+}