github.com/ava-labs/avalanchego@v1.11.11/network/throttling/inbound_msg_buffer_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 14 "github.com/ava-labs/avalanchego/ids" 15 ) 16 17 // Test inboundMsgBufferThrottler 18 func TestMsgBufferThrottler(t *testing.T) { 19 require := require.New(t) 20 throttler, err := newInboundMsgBufferThrottler(prometheus.NewRegistry(), 3) 21 require.NoError(err) 22 23 nodeID1, nodeID2 := ids.GenerateTestNodeID(), ids.GenerateTestNodeID() 24 // Acquire shouldn't block for first 3 25 throttler.Acquire(context.Background(), nodeID1) 26 throttler.Acquire(context.Background(), nodeID1) 27 throttler.Acquire(context.Background(), nodeID1) 28 require.Len(throttler.nodeToNumProcessingMsgs, 1) 29 require.Equal(uint64(3), throttler.nodeToNumProcessingMsgs[nodeID1]) 30 31 // Acquire shouldn't block for other node 32 throttler.Acquire(context.Background(), nodeID2) 33 throttler.Acquire(context.Background(), nodeID2) 34 throttler.Acquire(context.Background(), nodeID2) 35 require.Len(throttler.nodeToNumProcessingMsgs, 2) 36 require.Equal(uint64(3), throttler.nodeToNumProcessingMsgs[nodeID1]) 37 require.Equal(uint64(3), throttler.nodeToNumProcessingMsgs[nodeID2]) 38 39 // Acquire should block for 4th acquire 40 done := make(chan struct{}) 41 go func() { 42 throttler.Acquire(context.Background(), nodeID1) 43 done <- struct{}{} 44 }() 45 select { 46 case <-done: 47 require.FailNow("should block on acquiring") 48 case <-time.After(50 * time.Millisecond): 49 } 50 51 throttler.release(nodeID1) 52 // fourth acquire should be unblocked 53 <-done 54 require.Len(throttler.nodeToNumProcessingMsgs, 2) 55 require.Equal(uint64(3), throttler.nodeToNumProcessingMsgs[nodeID2]) 56 57 // Releasing from other node should have no effect 58 throttler.release(nodeID2) 59 throttler.release(nodeID2) 60 throttler.release(nodeID2) 61 62 // Release remaining 3 acquires 63 throttler.release(nodeID1) 64 throttler.release(nodeID1) 65 throttler.release(nodeID1) 66 require.Empty(throttler.nodeToNumProcessingMsgs) 67 } 68 69 // Test inboundMsgBufferThrottler when an acquire is cancelled 70 func TestMsgBufferThrottlerContextCancelled(t *testing.T) { 71 require := require.New(t) 72 throttler, err := newInboundMsgBufferThrottler(prometheus.NewRegistry(), 3) 73 require.NoError(err) 74 75 vdr1Context, vdr1ContextCancelFunc := context.WithCancel(context.Background()) 76 nodeID1 := ids.GenerateTestNodeID() 77 // Acquire shouldn't block for first 3 78 throttler.Acquire(vdr1Context, nodeID1) 79 throttler.Acquire(vdr1Context, nodeID1) 80 throttler.Acquire(vdr1Context, nodeID1) 81 require.Len(throttler.nodeToNumProcessingMsgs, 1) 82 require.Equal(uint64(3), throttler.nodeToNumProcessingMsgs[nodeID1]) 83 84 // Acquire should block for 4th acquire 85 done := make(chan struct{}) 86 go func() { 87 throttler.Acquire(vdr1Context, nodeID1) 88 done <- struct{}{} 89 }() 90 select { 91 case <-done: 92 require.FailNow("should block on acquiring") 93 case <-time.After(50 * time.Millisecond): 94 } 95 96 // Acquire should block for 5th acquire 97 done2 := make(chan struct{}) 98 go func() { 99 throttler.Acquire(vdr1Context, nodeID1) 100 done2 <- struct{}{} 101 }() 102 select { 103 case <-done2: 104 require.FailNow("should block on acquiring") 105 case <-time.After(50 * time.Millisecond): 106 } 107 108 // Unblock fifth acquire 109 vdr1ContextCancelFunc() 110 select { 111 case <-done2: 112 case <-time.After(50 * time.Millisecond): 113 require.FailNow("cancelling context should unblock Acquire") 114 } 115 select { 116 case <-done: 117 case <-time.After(50 * time.Millisecond): 118 require.FailNow("should be blocked") 119 } 120 }