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  }