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  }