github.com/koko1123/flow-go-1@v0.29.6/network/p2p/unicast/ratelimit/bandwidth_rate_limiter_test.go (about)

     1  package ratelimit
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	"github.com/koko1123/flow-go-1/model/flow"
     8  	libp2pmessage "github.com/koko1123/flow-go-1/model/libp2p/message"
     9  	"github.com/koko1123/flow-go-1/network"
    10  
    11  	"github.com/stretchr/testify/require"
    12  	"golang.org/x/time/rate"
    13  
    14  	"github.com/koko1123/flow-go-1/network/channels"
    15  	"github.com/koko1123/flow-go-1/utils/unittest"
    16  )
    17  
    18  // TestBandWidthRateLimiter_Allow ensures rate limiter allows messages as expected.
    19  func TestBandWidthRateLimiter_Allow(t *testing.T) {
    20  	//limiter limit will be set up to 1000 bytes/sec
    21  	limit := rate.Limit(1000)
    22  
    23  	//burst per interval
    24  	burst := 1000
    25  
    26  	// setup bandwidth rate limiter
    27  	bandwidthRateLimiter := NewBandWidthRateLimiter(limit, burst, 1)
    28  
    29  	id, _ := unittest.IdentityWithNetworkingKeyFixture()
    30  	peerID, err := unittest.PeerIDFromFlowID(id)
    31  	require.NoError(t, err)
    32  
    33  	// create message with about 400bytes (300 random bytes + 100bytes message info)
    34  	b := make([]byte, 300)
    35  	for i := range b {
    36  		b[i] = byte('X')
    37  	}
    38  
    39  	msg, err := network.NewOutgoingScope(
    40  		flow.IdentifierList{unittest.IdentifierFixture()},
    41  		channels.TestNetworkChannel,
    42  		&libp2pmessage.TestMessage{
    43  			Text: string(b),
    44  		},
    45  		unittest.NetworkCodec().Encode,
    46  		network.ProtocolTypeUnicast)
    47  	require.NoError(t, err)
    48  
    49  	allowed := bandwidthRateLimiter.Allow(peerID, msg.Size())
    50  	require.True(t, allowed)
    51  	allowed = bandwidthRateLimiter.Allow(peerID, msg.Size())
    52  	require.True(t, allowed)
    53  	allowed = bandwidthRateLimiter.Allow(peerID, msg.Size())
    54  	require.False(t, allowed)
    55  
    56  	// wait for 1 second, the rate limiter should allow 3 messages again
    57  	time.Sleep(1 * time.Second)
    58  	allowed = bandwidthRateLimiter.Allow(peerID, msg.Size())
    59  	require.True(t, allowed)
    60  	allowed = bandwidthRateLimiter.Allow(peerID, msg.Size())
    61  	require.True(t, allowed)
    62  	allowed = bandwidthRateLimiter.Allow(peerID, msg.Size())
    63  	require.False(t, allowed)
    64  }
    65  
    66  // TestBandWidthRateLimiter_IsRateLimited ensures IsRateLimited returns true for peers that are rate limited.
    67  func TestBandWidthRateLimiter_IsRateLimited(t *testing.T) {
    68  	//limiter limit will be set up to 1000 bytes/sec
    69  	limit := rate.Limit(1000)
    70  
    71  	//burst per interval
    72  	burst := 1000
    73  
    74  	// setup bandwidth rate limiter
    75  	bandwidthRateLimiter := NewBandWidthRateLimiter(limit, burst, 1)
    76  
    77  	// for the duration of a simulated second we will send 3 messages. Each message is about
    78  	// 400 bytes, the 3rd message will put our limiter over the 1000 byte limit at 1200 bytes. Thus
    79  	// the 3rd message should be rate limited.
    80  	id, _ := unittest.IdentityWithNetworkingKeyFixture()
    81  	peerID, err := unittest.PeerIDFromFlowID(id)
    82  	require.NoError(t, err)
    83  
    84  	// create message with about 400bytes (300 random bytes + 100bytes message info)
    85  	b := make([]byte, 500000)
    86  	for i := range b {
    87  		b[i] = byte('X')
    88  	}
    89  
    90  	require.False(t, bandwidthRateLimiter.IsRateLimited(peerID))
    91  
    92  	msg, err := network.NewOutgoingScope(
    93  		flow.IdentifierList{unittest.IdentifierFixture()},
    94  		channels.TestNetworkChannel,
    95  		&libp2pmessage.TestMessage{
    96  			Text: string(b),
    97  		},
    98  		unittest.NetworkCodec().Encode,
    99  		network.ProtocolTypeUnicast)
   100  	require.NoError(t, err)
   101  
   102  	allowed := bandwidthRateLimiter.Allow(peerID, msg.Size())
   103  	require.False(t, allowed)
   104  	require.True(t, bandwidthRateLimiter.IsRateLimited(peerID))
   105  
   106  	// wait for 1 second, the rate limiter should reset
   107  	time.Sleep(1 * time.Second)
   108  	require.False(t, bandwidthRateLimiter.IsRateLimited(peerID))
   109  }