github.com/ava-labs/avalanchego@v1.11.11/network/throttling/bandwidth_throttler_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 throttling
     5  
     6  import (
     7  	"context"
     8  	"sync"
     9  	"testing"
    10  
    11  	"github.com/prometheus/client_golang/prometheus"
    12  	"github.com/stretchr/testify/require"
    13  
    14  	"github.com/ava-labs/avalanchego/ids"
    15  	"github.com/ava-labs/avalanchego/utils/logging"
    16  )
    17  
    18  func TestBandwidthThrottler(t *testing.T) {
    19  	require := require.New(t)
    20  	// Assert initial state
    21  	config := BandwidthThrottlerConfig{
    22  		RefillRate:   8,
    23  		MaxBurstSize: 10,
    24  	}
    25  	throttlerIntf, err := newBandwidthThrottler(logging.NoLog{}, prometheus.NewRegistry(), config)
    26  	require.NoError(err)
    27  	require.IsType(&bandwidthThrottlerImpl{}, throttlerIntf)
    28  	throttler := throttlerIntf.(*bandwidthThrottlerImpl)
    29  	require.NotNil(throttler.log)
    30  	require.NotNil(throttler.limiters)
    31  	require.Equal(config.RefillRate, throttler.RefillRate)
    32  	require.Equal(config.MaxBurstSize, throttler.MaxBurstSize)
    33  	require.Empty(throttler.limiters)
    34  
    35  	// Add a node
    36  	nodeID1 := ids.GenerateTestNodeID()
    37  	throttler.AddNode(nodeID1)
    38  	require.Len(throttler.limiters, 1)
    39  
    40  	// Remove the node
    41  	throttler.RemoveNode(nodeID1)
    42  	require.Empty(throttler.limiters)
    43  
    44  	// Add the node back
    45  	throttler.AddNode(nodeID1)
    46  	require.Len(throttler.limiters, 1)
    47  
    48  	// Should be able to acquire 8
    49  	throttler.Acquire(context.Background(), 8, nodeID1)
    50  
    51  	// Make several goroutines that acquire bytes.
    52  	wg := sync.WaitGroup{}
    53  	wg.Add(int(config.MaxBurstSize) + 5)
    54  	for i := uint64(0); i < config.MaxBurstSize+5; i++ {
    55  		go func() {
    56  			throttler.Acquire(context.Background(), 1, nodeID1)
    57  			wg.Done()
    58  		}()
    59  	}
    60  	wg.Wait()
    61  }