github.com/ava-labs/avalanchego@v1.11.11/network/p2p/peer_tracker_test.go (about) 1 // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved. 2 // See the file LICENSE for licensing terms. 3 4 package p2p 5 6 import ( 7 "testing" 8 9 "github.com/prometheus/client_golang/prometheus" 10 "github.com/stretchr/testify/require" 11 12 "github.com/ava-labs/avalanchego/ids" 13 "github.com/ava-labs/avalanchego/utils/logging" 14 "github.com/ava-labs/avalanchego/version" 15 ) 16 17 func TestPeerTracker(t *testing.T) { 18 require := require.New(t) 19 p, err := NewPeerTracker( 20 logging.NoLog{}, 21 "", 22 prometheus.NewRegistry(), 23 nil, 24 nil, 25 ) 26 require.NoError(err) 27 28 // Connect some peers 29 numExtraPeers := 10 30 numPeers := desiredMinResponsivePeers + numExtraPeers 31 peerIDs := make([]ids.NodeID, numPeers) 32 peerVersion := &version.Application{ 33 Major: 1, 34 Minor: 2, 35 Patch: 3, 36 } 37 38 for i := range peerIDs { 39 peerIDs[i] = ids.GenerateTestNodeID() 40 p.Connected(peerIDs[i], peerVersion) 41 } 42 43 responsivePeers := make(map[ids.NodeID]bool) 44 45 // Expect requests to go to new peers until we have desiredMinResponsivePeers responsive peers. 46 for i := 0; i < desiredMinResponsivePeers+numExtraPeers/2; i++ { 47 peer, ok := p.SelectPeer() 48 require.True(ok) 49 require.NotZero(peer) 50 51 _, exists := responsivePeers[peer] 52 require.Falsef(exists, "expected connecting to a new peer, but got the same peer twice: peer %s iteration %d", peer, i) 53 responsivePeers[peer] = true 54 55 p.RegisterRequest(peer) // mark the peer as having a message sent to it 56 } 57 58 // Mark some peers as responsive and others as not responsive 59 i := 0 60 for peer := range responsivePeers { 61 if i < desiredMinResponsivePeers { 62 p.RegisterResponse(peer, 10) 63 } else { 64 responsivePeers[peer] = false // remember which peers were not responsive 65 p.RegisterFailure(peer) 66 } 67 i++ 68 } 69 70 // Expect requests to go to responsive or new peers, so long as they are available 71 numRequests := 50 72 for i := 0; i < numRequests; i++ { 73 peer, ok := p.SelectPeer() 74 require.True(ok) 75 require.NotZero(peer) 76 77 responsive, ok := responsivePeers[peer] 78 if ok { 79 require.Truef(responsive, "expected connecting to a responsive peer, but got a peer that was not responsive: peer %s iteration %d", peer, i) 80 p.RegisterResponse(peer, 10) 81 } else { 82 responsivePeers[peer] = false // remember that we connected to this peer 83 p.RegisterRequest(peer) // mark the peer as having a message sent to it 84 p.RegisterFailure(peer) // mark the peer as non-responsive 85 } 86 } 87 88 // Disconnect from peers that were previously responsive and ones we didn't connect to yet. 89 for _, peer := range peerIDs { 90 responsive, ok := responsivePeers[peer] 91 if ok && responsive || !ok { 92 p.Disconnected(peer) 93 } 94 } 95 96 // Requests should fall back on non-responsive peers when no other choice is left 97 peer, ok := p.SelectPeer() 98 require.True(ok) 99 require.NotZero(peer) 100 101 responsive, ok := responsivePeers[peer] 102 require.True(ok) 103 require.Falsef(responsive, "expected connecting to a non-responsive peer, but got a peer that was responsive: peer %s", peer) 104 }