diff --git a/peer/peer_tracker_test.go b/peer/peer_tracker_test.go new file mode 100644 index 0000000000..f4a510fc4e --- /dev/null +++ b/peer/peer_tracker_test.go @@ -0,0 +1,88 @@ +// (c) 2019-2022, Ava Labs, Inc. All rights reserved. +// See the file LICENSE for licensing terms. + +package peer + +import ( + "testing" + + "github.com/ava-labs/avalanchego/ids" + "github.com/stretchr/testify/require" +) + +func TestPeerTracker(t *testing.T) { + require := require.New(t) + p := NewPeerTracker() + + // Connect some peers + numExtraPeers := 10 + numPeers := desiredMinResponsivePeers + numExtraPeers + peerIDs := make([]ids.NodeID, numPeers) + + for i := range peerIDs { + peerIDs[i] = ids.GenerateTestNodeID() + p.Connected(peerIDs[i], defaultPeerVersion) + } + + responsivePeers := make(map[ids.NodeID]bool) + + // Expect requests to go to new peers until we have desiredMinResponsivePeers responsive peers. + for i := 0; i < desiredMinResponsivePeers+numExtraPeers/2; i++ { + peer, ok := p.GetAnyPeer(nil) + require.True(ok) + require.NotNil(peer) + + _, exists := responsivePeers[peer] + require.Falsef(exists, "expected connecting to a new peer, but got the same peer twice: peer %s iteration %d", peer, i) + responsivePeers[peer] = true + + p.TrackPeer(peer) // mark the peer as having a message sent to it + } + + // Mark some peers as responsive and others as not responsive + i := 0 + for peer := range responsivePeers { + if i < desiredMinResponsivePeers { + p.TrackBandwidth(peer, 10) + } else { + responsivePeers[peer] = false // remember which peers were not responsive + p.TrackBandwidth(peer, 0) + } + i++ + } + + // Expect requests to go to responsive or new peers, so long as they are available + numRequests := 50 + for i := 0; i < numRequests; i++ { + peer, ok := p.GetAnyPeer(nil) + require.True(ok) + require.NotNil(peer) + + responsive, ok := responsivePeers[peer] + if ok { + require.Truef(responsive, "expected connecting to a responsive peer, but got a peer that was not responsive: peer %s iteration %d", peer, i) + p.TrackBandwidth(peer, 10) + } else { + responsivePeers[peer] = false // remember that we connected to this peer + p.TrackPeer(peer) // mark the peer as having a message sent to it + p.TrackBandwidth(peer, 0) // mark the peer as non-responsive + } + } + + // Disconnect from peers that were previously responsive and ones we didn't connect to yet. + for _, peer := range peerIDs { + responsive, ok := responsivePeers[peer] + if ok && responsive || !ok { + p.Disconnected(peer) + } + } + + // Requests should fall back on non-responsive peers when no other choice is left + peer, ok := p.GetAnyPeer(nil) + require.True(ok) + require.NotNil(peer) + + responsive, ok := responsivePeers[peer] + require.True(ok) + require.Falsef(responsive, "expected connecting to a non-responsive peer, but got a peer that was responsive: peer %s", peer) +}