github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/network/p2p/utils/ratelimiter/rate_limiter.go (about)

     1  package ratelimiter
     2  
     3  import (
     4  	"time"
     5  
     6  	"github.com/libp2p/go-libp2p/core/peer"
     7  	"golang.org/x/time/rate"
     8  
     9  	"github.com/onflow/flow-go/module/component"
    10  	"github.com/onflow/flow-go/module/irrecoverable"
    11  	"github.com/onflow/flow-go/network/p2p"
    12  	"github.com/onflow/flow-go/network/p2p/utils/ratelimiter/internal"
    13  )
    14  
    15  const (
    16  	cleanUpTickInterval = 10 * time.Minute
    17  	rateLimiterTTL      = 10 * time.Minute
    18  )
    19  
    20  // RateLimiter generic rate limiter
    21  type RateLimiter struct {
    22  	component.Component
    23  	// limiterMap map that stores a rate limiter with metadata per peer.
    24  	limiterMap *internal.RateLimiterMap
    25  	// limit amount of messages allowed per second.
    26  	limit rate.Limit
    27  	// burst amount of messages allowed at one time.
    28  	burst int
    29  	// rateLimitLockoutDuration the amount of time that has to pass before a peer is allowed to connect.
    30  	rateLimitLockoutDuration time.Duration
    31  }
    32  
    33  var _ component.Component = (*RateLimiter)(nil)
    34  var _ p2p.RateLimiter = (*RateLimiter)(nil)
    35  
    36  // NewRateLimiter returns a new RateLimiter.
    37  func NewRateLimiter(limit rate.Limit, burst int, lockoutDuration time.Duration, opts ...p2p.RateLimiterOpt) *RateLimiter {
    38  	l := &RateLimiter{
    39  		limiterMap:               internal.NewLimiterMap(rateLimiterTTL, cleanUpTickInterval),
    40  		limit:                    limit,
    41  		burst:                    burst,
    42  		rateLimitLockoutDuration: lockoutDuration,
    43  	}
    44  
    45  	for _, opt := range opts {
    46  		opt(l)
    47  	}
    48  
    49  	l.Component = component.NewComponentManagerBuilder().
    50  		AddWorker(func(ctx irrecoverable.SignalerContext, ready component.ReadyFunc) {
    51  			ready()
    52  			l.limiterMap.CleanupLoop(ctx)
    53  		}).Build()
    54  
    55  	return l
    56  }
    57  
    58  // Allow checks the cached limiter for the peer and returns limiterMap.Allow().
    59  // If a limiter is not cached for a peer one is created. This func can be overridden
    60  // and the message size parameter can be used with AllowN.
    61  func (r *RateLimiter) Allow(peerID peer.ID, _ int) bool {
    62  	limiter := r.GetLimiter(peerID)
    63  	if !limiter.AllowN(time.Now(), 1) {
    64  		r.limiterMap.UpdateLastRateLimit(peerID, time.Now())
    65  		return false
    66  	}
    67  
    68  	return true
    69  }
    70  
    71  // IsRateLimited returns true is a peer is currently rate limited.
    72  func (r *RateLimiter) IsRateLimited(peerID peer.ID) bool {
    73  	metadata, ok := r.limiterMap.Get(peerID)
    74  	if !ok {
    75  		return false
    76  	}
    77  	return time.Since(metadata.LastRateLimit()) < r.rateLimitLockoutDuration
    78  }
    79  
    80  // GetLimiter returns limiter for the peerID, if a limiter does not exist one is created and stored.
    81  func (r *RateLimiter) GetLimiter(peerID peer.ID) *rate.Limiter {
    82  	if metadata, ok := r.limiterMap.Get(peerID); ok {
    83  		return metadata.Limiter()
    84  	}
    85  
    86  	limiter := rate.NewLimiter(r.limit, r.burst)
    87  	r.limiterMap.Store(peerID, limiter)
    88  
    89  	return limiter
    90  }
    91  
    92  // UpdateLastRateLimit updates the last time a peer was rate limited in the limiter map.
    93  func (r *RateLimiter) UpdateLastRateLimit(peerID peer.ID, lastRateLimit time.Time) {
    94  	r.limiterMap.UpdateLastRateLimit(peerID, lastRateLimit)
    95  }