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

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