Tidy-up of BestFinalized

beacon-chain/p2p/peers/status.go

@@ -21,6 +21,7 @@ package peers
 
 import (
 	"errors"
+	"sort"
 	"sync"
 	"time"
 
@@ -319,41 +320,56 @@ func (p *Status) Decay() {
 
 // BestFinalized returns the highest finalized epoch equal to or higher than ours that is agreed upon by the majority of peers.
 // This method may not return the absolute highest finalized, but the finalized epoch in which most peers can serve blocks.
-// Ideally, all peers would be reporting the same finalized epoch.
-// Returns the best finalized root, epoch number, and list of peers that agree.
+// Ideally, all peers would be reporting the same finalized epoch but some may be behind due to their own latency, or because of
+// their finalized epoch at the time we queried them.
+// Returns the best finalized root, epoch number, and list of peers that are at or beyond that epoch.
 func (p *Status) BestFinalized(maxPeers int, ourFinalizedEpoch uint64) ([]byte, uint64, []peer.ID) {
+	connected := p.Connected()
 	finalized := make(map[[32]byte]uint64)
 	rootToEpoch := make(map[[32]byte]uint64)
-	for _, pid := range p.Connected() {
+	pidEpochs := make(map[peer.ID]uint64)
+	potentialPIDs := make([]peer.ID, 0, len(connected))
+	for _, pid := range connected {
 		peerChainState, err := p.ChainState(pid)
 		if err == nil && peerChainState != nil && peerChainState.FinalizedEpoch >= ourFinalizedEpoch {
-			r := bytesutil.ToBytes32(peerChainState.FinalizedRoot)
-			finalized[r]++
-			rootToEpoch[r] = peerChainState.FinalizedEpoch
+			root := bytesutil.ToBytes32(peerChainState.FinalizedRoot)
+			finalized[root]++
+			rootToEpoch[root] = peerChainState.FinalizedEpoch
+			pidEpochs[pid] = peerChainState.FinalizedEpoch
+			potentialPIDs = append(potentialPIDs, pid)
 		}
 	}
 
-	var mostVotedFinalizedRoot [32]byte
+	// Select the target epoch, which is the epoch most peers agree upon.
+	var targetRoot [32]byte
 	var mostVotes uint64
 	for root, count := range finalized {
 		if count > mostVotes {
 			mostVotes = count
-			mostVotedFinalizedRoot = root
+			targetRoot = root
 		}
 	}
-
-	var pids []peer.ID
-	for _, pid := range p.Connected() {
-		peerChainState, err := p.ChainState(pid)
-		if err == nil && peerChainState != nil && peerChainState.FinalizedEpoch >= rootToEpoch[mostVotedFinalizedRoot] {
-			pids = append(pids, pid)
-			if len(pids) >= maxPeers {
-				break
-			}
+	targetEpoch := rootToEpoch[targetRoot]
+
+	// Sort PIDs by finalized epoch, in decreasing order.
+	sort.Slice(potentialPIDs, func(i, j int) bool {
+		return pidEpochs[potentialPIDs[i]] > pidEpochs[potentialPIDs[j]]
+	})
+
+	// Trim potential peers to those on or after target epoch.
+	for i, pid := range potentialPIDs {
+		if pidEpochs[pid] < targetEpoch {
+			potentialPIDs = potentialPIDs[:i]
+			break
 		}
 	}
 
-	return mostVotedFinalizedRoot[:], rootToEpoch[mostVotedFinalizedRoot], pids
+	// Trim potential peers to at most maxPeers.
+	if len(potentialPIDs) > maxPeers {
+		potentialPIDs = potentialPIDs[:maxPeers]
+	}
+
+	return targetRoot[:], targetEpoch, potentialPIDs
 }
 
 // fetch is a helper function that fetches a peer status, possibly creating it.

beacon-chain/p2p/peers/status_test.go

@@ -335,6 +335,58 @@ func TestDecay(t *testing.T) {
 	}
 }
 
+func TestTrimmedOrderedPeers(t *testing.T) {
+	p := peers.NewStatus(1)
+
+	expectedTarget := uint64(2)
+	maxPeers := 3
+
+	// Peer 1
+	pid1 := addPeer(t, p, peers.PeerConnected)
+	p.SetChainState(pid1, &pb.Status{
+		FinalizedEpoch: 3,
+	})
+	// Peer 2
+	pid2 := addPeer(t, p, peers.PeerConnected)
+	p.SetChainState(pid2, &pb.Status{
+		FinalizedEpoch: 4,
+	})
+	// Peer 3
+	pid3 := addPeer(t, p, peers.PeerConnected)
+	p.SetChainState(pid3, &pb.Status{
+		FinalizedEpoch: 5,
+	})
+	// Peer 4
+	pid4 := addPeer(t, p, peers.PeerConnected)
+	p.SetChainState(pid4, &pb.Status{
+		FinalizedEpoch: 2,
+	})
+	// Peer 5
+	pid5 := addPeer(t, p, peers.PeerConnected)
+	p.SetChainState(pid5, &pb.Status{
+		FinalizedEpoch: 2,
+	})
+
+	_, target, pids := p.BestFinalized(maxPeers, 0)
+	if target != expectedTarget {
+		t.Errorf("Incorrect target epoch retrieved; wanted %v but got %v", expectedTarget, target)
+	}
+	if len(pids) != maxPeers {
+		t.Errorf("Incorrect number of peers retrieved; wanted %v but got %v", maxPeers, len(pids))
+	}
+
+	// Expect the returned list to be ordered by finalized epoch and trimmed to max peers.
+	if pids[0] != pid3 {
+		t.Errorf("Incorrect first peer; wanted %v but got %v", pid3, pids[0])
+	}
+	if pids[1] != pid2 {
+		t.Errorf("Incorrect second peer; wanted %v but got %v", pid2, pids[1])
+	}
+	if pids[2] != pid1 {
+		t.Errorf("Incorrect third peer; wanted %v but got %v", pid1, pids[2])
+	}
+}
+
 func TestBestPeer(t *testing.T) {
 	maxBadResponses := 2
 	expectedFinEpoch := uint64(4)