github.com/MetalBlockchain/metalgo@v1.11.9/node/beacon_manager_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 node
     5  
     6  import (
     7  	"sync"
     8  	"testing"
     9  
    10  	"github.com/stretchr/testify/require"
    11  	"go.uber.org/mock/gomock"
    12  
    13  	"github.com/MetalBlockchain/metalgo/ids"
    14  	"github.com/MetalBlockchain/metalgo/snow/networking/router"
    15  	"github.com/MetalBlockchain/metalgo/snow/validators"
    16  	"github.com/MetalBlockchain/metalgo/utils/constants"
    17  	"github.com/MetalBlockchain/metalgo/version"
    18  )
    19  
    20  const numValidators = 5_000
    21  
    22  // Tests that reconnects that mutate the beacon manager's current total stake
    23  // weight is consistent. Test is not deterministic.
    24  func TestBeaconManager_DataRace(t *testing.T) {
    25  	require := require.New(t)
    26  
    27  	validatorIDs := make([]ids.NodeID, 0, numValidators)
    28  	validatorSet := validators.NewManager()
    29  	for i := 0; i < numValidators; i++ {
    30  		nodeID := ids.GenerateTestNodeID()
    31  
    32  		require.NoError(validatorSet.AddStaker(constants.PrimaryNetworkID, nodeID, nil, ids.Empty, 1))
    33  		validatorIDs = append(validatorIDs, nodeID)
    34  	}
    35  
    36  	wg := &sync.WaitGroup{}
    37  
    38  	ctrl := gomock.NewController(t)
    39  	mockRouter := router.NewMockRouter(ctrl)
    40  
    41  	b := beaconManager{
    42  		Router:                  mockRouter,
    43  		beacons:                 validatorSet,
    44  		requiredConns:           numValidators,
    45  		onSufficientlyConnected: make(chan struct{}),
    46  	}
    47  
    48  	// connect numValidators validators, each with a weight of 1
    49  	wg.Add(2 * numValidators)
    50  	mockRouter.EXPECT().
    51  		Connected(gomock.Any(), gomock.Any(), gomock.Any()).
    52  		Times(2 * numValidators).
    53  		Do(func(ids.NodeID, *version.Application, ids.ID) {
    54  			wg.Done()
    55  		})
    56  
    57  	for _, nodeID := range validatorIDs {
    58  		nodeID := nodeID
    59  		go func() {
    60  			b.Connected(nodeID, version.CurrentApp, constants.PrimaryNetworkID)
    61  			b.Connected(nodeID, version.CurrentApp, ids.GenerateTestID())
    62  		}()
    63  	}
    64  	wg.Wait()
    65  
    66  	// we should have a weight of numValidators now
    67  	require.Equal(int64(numValidators), b.numConns)
    68  
    69  	// disconnect numValidators validators
    70  	wg.Add(numValidators)
    71  	mockRouter.EXPECT().
    72  		Disconnected(gomock.Any()).
    73  		Times(numValidators).
    74  		Do(func(ids.NodeID) {
    75  			wg.Done()
    76  		})
    77  
    78  	for _, nodeID := range validatorIDs {
    79  		go b.Disconnected(nodeID)
    80  	}
    81  	wg.Wait()
    82  
    83  	// we should a weight of zero now
    84  	require.Zero(b.numConns)
    85  }