github.com/MetalBlockchain/metalgo@v1.11.9/network/throttling/inbound_resource_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  	"testing"
     9  	"time"
    10  
    11  	"github.com/prometheus/client_golang/prometheus"
    12  	"github.com/stretchr/testify/require"
    13  	"go.uber.org/mock/gomock"
    14  
    15  	"github.com/MetalBlockchain/metalgo/ids"
    16  	"github.com/MetalBlockchain/metalgo/snow/networking/tracker"
    17  	"github.com/MetalBlockchain/metalgo/utils/math/meter"
    18  	"github.com/MetalBlockchain/metalgo/utils/resource"
    19  	"github.com/MetalBlockchain/metalgo/utils/timer/mockable"
    20  )
    21  
    22  func TestNewSystemThrottler(t *testing.T) {
    23  	ctrl := gomock.NewController(t)
    24  	require := require.New(t)
    25  	reg := prometheus.NewRegistry()
    26  	clock := mockable.Clock{}
    27  	clock.Set(time.Now())
    28  	resourceTracker, err := tracker.NewResourceTracker(reg, resource.NoUsage, meter.ContinuousFactory{}, time.Second)
    29  	require.NoError(err)
    30  	cpuTracker := resourceTracker.CPUTracker()
    31  
    32  	config := SystemThrottlerConfig{
    33  		Clock:           clock,
    34  		MaxRecheckDelay: time.Second,
    35  	}
    36  	targeter := tracker.NewMockTargeter(ctrl)
    37  	throttlerIntf, err := NewSystemThrottler("", reg, config, cpuTracker, targeter)
    38  	require.NoError(err)
    39  	require.IsType(&systemThrottler{}, throttlerIntf)
    40  	throttler := throttlerIntf.(*systemThrottler)
    41  	require.Equal(clock, config.Clock)
    42  	require.Equal(time.Second, config.MaxRecheckDelay)
    43  	require.Equal(cpuTracker, throttler.tracker)
    44  	require.Equal(targeter, throttler.targeter)
    45  }
    46  
    47  func TestSystemThrottler(t *testing.T) {
    48  	ctrl := gomock.NewController(t)
    49  	require := require.New(t)
    50  
    51  	// Setup
    52  	mockTracker := tracker.NewMockTracker(ctrl)
    53  	maxRecheckDelay := 100 * time.Millisecond
    54  	config := SystemThrottlerConfig{
    55  		MaxRecheckDelay: maxRecheckDelay,
    56  	}
    57  	vdrID, nonVdrID := ids.GenerateTestNodeID(), ids.GenerateTestNodeID()
    58  	targeter := tracker.NewMockTargeter(ctrl)
    59  	throttler, err := NewSystemThrottler("", prometheus.NewRegistry(), config, mockTracker, targeter)
    60  	require.NoError(err)
    61  
    62  	// Case: Actual usage <= target usage; should return immediately
    63  	// for both validator and non-validator
    64  	targeter.EXPECT().TargetUsage(vdrID).Return(1.0).Times(1)
    65  	mockTracker.EXPECT().Usage(vdrID, gomock.Any()).Return(0.9).Times(1)
    66  
    67  	throttler.Acquire(context.Background(), vdrID)
    68  
    69  	targeter.EXPECT().TargetUsage(nonVdrID).Return(1.0).Times(1)
    70  	mockTracker.EXPECT().Usage(nonVdrID, gomock.Any()).Return(0.9).Times(1)
    71  
    72  	throttler.Acquire(context.Background(), nonVdrID)
    73  
    74  	// Case: Actual usage > target usage; we should wait.
    75  	// In the first loop iteration inside acquire,
    76  	// say the actual usage exceeds the target.
    77  	targeter.EXPECT().TargetUsage(vdrID).Return(0.0).Times(1)
    78  	mockTracker.EXPECT().Usage(vdrID, gomock.Any()).Return(1.0).Times(1)
    79  	// Note we'll only actually wait [maxRecheckDelay]. We set [timeUntilAtDiskTarget]
    80  	// much larger to assert that the min recheck frequency is honored below.
    81  	timeUntilAtDiskTarget := 100 * maxRecheckDelay
    82  	mockTracker.EXPECT().TimeUntilUsage(vdrID, gomock.Any(), gomock.Any()).Return(timeUntilAtDiskTarget).Times(1)
    83  
    84  	// The second iteration, say the usage is OK.
    85  	targeter.EXPECT().TargetUsage(vdrID).Return(1.0).Times(1)
    86  	mockTracker.EXPECT().Usage(vdrID, gomock.Any()).Return(0.0).Times(1)
    87  
    88  	onAcquire := make(chan struct{})
    89  
    90  	// Check for validator
    91  	go func() {
    92  		throttler.Acquire(context.Background(), vdrID)
    93  		onAcquire <- struct{}{}
    94  	}()
    95  	// Make sure the min re-check frequency is honored
    96  	select {
    97  	// Use 5*maxRecheckDelay and not just maxRecheckDelay to give a buffer
    98  	// and avoid flakiness. If the min re-check freq isn't honored,
    99  	// we'll wait [timeUntilAtDiskTarget].
   100  	case <-time.After(5 * maxRecheckDelay):
   101  		require.FailNow("should have returned after about [maxRecheckDelay]")
   102  	case <-onAcquire:
   103  	}
   104  
   105  	targeter.EXPECT().TargetUsage(nonVdrID).Return(0.0).Times(1)
   106  	mockTracker.EXPECT().Usage(nonVdrID, gomock.Any()).Return(1.0).Times(1)
   107  
   108  	mockTracker.EXPECT().TimeUntilUsage(nonVdrID, gomock.Any(), gomock.Any()).Return(timeUntilAtDiskTarget).Times(1)
   109  
   110  	targeter.EXPECT().TargetUsage(nonVdrID).Return(1.0).Times(1)
   111  	mockTracker.EXPECT().Usage(nonVdrID, gomock.Any()).Return(0.0).Times(1)
   112  
   113  	// Check for non-validator
   114  	go func() {
   115  		throttler.Acquire(context.Background(), nonVdrID)
   116  		onAcquire <- struct{}{}
   117  	}()
   118  	// Make sure the min re-check frequency is honored
   119  	select {
   120  	// Use 5*maxRecheckDelay and not just maxRecheckDelay to give a buffer
   121  	// and avoid flakiness. If the min re-check freq isn't honored,
   122  	// we'll wait [timeUntilAtDiskTarget].
   123  	case <-time.After(5 * maxRecheckDelay):
   124  		require.FailNow("should have returned after about [maxRecheckDelay]")
   125  	case <-onAcquire:
   126  	}
   127  }
   128  
   129  func TestSystemThrottlerContextCancel(t *testing.T) {
   130  	require := require.New(t)
   131  	ctrl := gomock.NewController(t)
   132  
   133  	// Setup
   134  	mockTracker := tracker.NewMockTracker(ctrl)
   135  	maxRecheckDelay := 10 * time.Second
   136  	config := SystemThrottlerConfig{
   137  		MaxRecheckDelay: maxRecheckDelay,
   138  	}
   139  	vdrID := ids.GenerateTestNodeID()
   140  	targeter := tracker.NewMockTargeter(ctrl)
   141  	throttler, err := NewSystemThrottler("", prometheus.NewRegistry(), config, mockTracker, targeter)
   142  	require.NoError(err)
   143  
   144  	// Case: Actual usage > target usage; we should wait.
   145  	// Mock the tracker so that the first loop iteration inside acquire,
   146  	// it says the actual usage exceeds the target.
   147  	// There should be no second iteration because we've already returned.
   148  	targeter.EXPECT().TargetUsage(vdrID).Return(0.0).Times(1)
   149  	mockTracker.EXPECT().Usage(vdrID, gomock.Any()).Return(1.0).Times(1)
   150  	mockTracker.EXPECT().TimeUntilUsage(vdrID, gomock.Any(), gomock.Any()).Return(maxRecheckDelay).Times(1)
   151  	onAcquire := make(chan struct{})
   152  	// Pass a canceled context into Acquire so that it returns immediately.
   153  	ctx, cancel := context.WithCancel(context.Background())
   154  	cancel()
   155  	go func() {
   156  		throttler.Acquire(ctx, vdrID)
   157  		onAcquire <- struct{}{}
   158  	}()
   159  	select {
   160  	case <-onAcquire:
   161  	case <-time.After(maxRecheckDelay / 2):
   162  		// Make sure Acquire returns well before the second check (i.e. "immediately")
   163  		require.Fail("should have returned immediately")
   164  	}
   165  }