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

     1  package p2p
     2  
     3  import (
     4  	"time"
     5  
     6  	"github.com/libp2p/go-libp2p/core/peer"
     7  )
     8  
     9  // RateLimiter unicast rate limiter interface
    10  type RateLimiter interface {
    11  	// Allow returns true if a message with the give size should be allowed to be processed.
    12  	Allow(peerID peer.ID, msgSize int) bool
    13  
    14  	// IsRateLimited returns true if a peer is rate limited.
    15  	IsRateLimited(peerID peer.ID) bool
    16  
    17  	// SetTimeNowFunc allows users to override the underlying time module used.
    18  	SetTimeNowFunc(now GetTimeNow)
    19  
    20  	// Stop sends cleanup signal to underlying rate limiters and rate limited peers maps. After the rate limiter
    21  	// is stopped it can not be reused.
    22  	Stop()
    23  
    24  	// Start starts cleanup loop for underlying rate limiters and rate limited peers maps.
    25  	Start()
    26  }
    27  
    28  // GetTimeNow callback used to get the current time. This allows us to improve testing by manipulating the current time
    29  // as opposed to using time.Now directly.
    30  type GetTimeNow func() time.Time
    31  
    32  type RateLimiterOpt func(limiter RateLimiter)
    33  
    34  func WithGetTimeNowFunc(now GetTimeNow) RateLimiterOpt {
    35  	return func(limiter RateLimiter) {
    36  		limiter.SetTimeNowFunc(now)
    37  	}
    38  }