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 }