github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/network/p2p/unicast/ratelimit/rate_limiters.go (about)

     1  package ratelimit
     2  
     3  import (
     4  	"github.com/libp2p/go-libp2p/core/peer"
     5  
     6  	"github.com/onflow/flow-go/network/channels"
     7  	"github.com/onflow/flow-go/network/p2p"
     8  )
     9  
    10  var (
    11  	ReasonMessageCount RateLimitReason = "messagecount"
    12  	ReasonBandwidth    RateLimitReason = "bandwidth"
    13  )
    14  
    15  type RateLimitReason string
    16  
    17  func (r RateLimitReason) String() string {
    18  	return string(r)
    19  }
    20  
    21  type OnRateLimitedPeerFunc func(pid peer.ID, role, msgType string, topic channels.Topic, reason RateLimitReason) // the callback called each time a peer is rate limited
    22  
    23  type RateLimitersOption func(*RateLimiters)
    24  
    25  func WithDisabledRateLimiting(disabled bool) RateLimitersOption {
    26  	return func(r *RateLimiters) {
    27  		r.disabled = disabled
    28  	}
    29  }
    30  
    31  func WithMessageRateLimiter(messageLimiter p2p.RateLimiter) RateLimitersOption {
    32  	return func(r *RateLimiters) {
    33  		r.MessageRateLimiter = messageLimiter
    34  	}
    35  }
    36  
    37  func WithBandwidthRateLimiter(bandwidthLimiter p2p.RateLimiter) RateLimitersOption {
    38  	return func(r *RateLimiters) {
    39  		r.BandWidthRateLimiter = bandwidthLimiter
    40  	}
    41  }
    42  
    43  func WithNotifier(notifier p2p.RateLimiterConsumer) RateLimitersOption {
    44  	return func(r *RateLimiters) {
    45  		r.notifier = notifier
    46  	}
    47  }
    48  
    49  // RateLimiters used to manage stream and bandwidth rate limiters
    50  type RateLimiters struct {
    51  	MessageRateLimiter   p2p.RateLimiter
    52  	BandWidthRateLimiter p2p.RateLimiter
    53  	notifier             p2p.RateLimiterConsumer
    54  	disabled             bool // flag allows rate limiter to collect metrics without rate limiting if set to false
    55  }
    56  
    57  // NewRateLimiters returns *RateLimiters
    58  func NewRateLimiters(opts ...RateLimitersOption) *RateLimiters {
    59  	r := NoopRateLimiters()
    60  
    61  	for _, opt := range opts {
    62  		opt(r)
    63  	}
    64  
    65  	return r
    66  }
    67  
    68  // Limiters returns list of all underlying rate limiters.
    69  func (r *RateLimiters) Limiters() []p2p.RateLimiter {
    70  	return []p2p.RateLimiter{r.MessageRateLimiter, r.BandWidthRateLimiter}
    71  }
    72  
    73  // MessageAllowed will return result from MessageRateLimiter.Allow. It will invoke the OnRateLimitedPeerFunc
    74  // callback each time a peer is not allowed.
    75  func (r *RateLimiters) MessageAllowed(peerID peer.ID) bool {
    76  	if r.MessageRateLimiter == nil {
    77  		return true
    78  	}
    79  
    80  	if !r.MessageRateLimiter.Allow(peerID, 0) { // 0 is not used for message rate limiter. It is only used for bandwidth rate limiter.
    81  		r.notifier.OnRateLimitedPeer(peerID, "", "", "", ReasonMessageCount.String())
    82  
    83  		// avoid rate limiting during dry run
    84  		return r.disabled
    85  	}
    86  
    87  	return true
    88  }
    89  
    90  // BandwidthAllowed will return result from BandWidthRateLimiter.Allow. It will invoke the OnRateLimitedPeerFunc
    91  // callback each time a peer is not allowed.
    92  func (r *RateLimiters) BandwidthAllowed(peerID peer.ID, originRole string, msgSize int, msgType string, msgTopic channels.Topic) bool {
    93  	if r.BandWidthRateLimiter == nil {
    94  		return true
    95  	}
    96  
    97  	if !r.BandWidthRateLimiter.Allow(peerID, msgSize) {
    98  		r.notifier.OnRateLimitedPeer(peerID, originRole, msgType, msgTopic.String(), ReasonBandwidth.String())
    99  
   100  		// avoid rate limiting during dry runs if disabled set to false
   101  		return r.disabled
   102  	}
   103  
   104  	return true
   105  }