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

     1  package ratelimit
     2  
     3  import (
     4  	"time"
     5  
     6  	"github.com/libp2p/go-libp2p/core/peer"
     7  	"github.com/rs/zerolog"
     8  	"golang.org/x/time/rate"
     9  
    10  	"github.com/onflow/flow-go/network/p2p"
    11  	"github.com/onflow/flow-go/network/p2p/unicast/ratelimit"
    12  	"github.com/onflow/flow-go/network/p2p/utils/ratelimiter"
    13  )
    14  
    15  // ControlMessageRateLimiter rate limiter that rate limits the amount of
    16  type ControlMessageRateLimiter struct {
    17  	*ratelimiter.RateLimiter
    18  }
    19  
    20  var _ p2p.BasicRateLimiter = (*ControlMessageRateLimiter)(nil)
    21  
    22  // NewControlMessageRateLimiter returns a new ControlMessageRateLimiter. The cleanup loop will be started in a
    23  // separate goroutine and should be stopped by calling Close.
    24  func NewControlMessageRateLimiter(logger zerolog.Logger, limit rate.Limit, burst int) p2p.BasicRateLimiter {
    25  	if limit == 0 {
    26  		logger.Warn().Msg("control message rate limit set to 0 using noop rate limiter")
    27  		// setup noop rate limiter if rate limiting is disabled
    28  		return ratelimit.NewNoopRateLimiter()
    29  	}
    30  
    31  	// NOTE: we use a lockout duration of 0 because we only need to expose the basic functionality of the
    32  	// rate limiter and not the lockout feature.
    33  	lockoutDuration := time.Duration(0)
    34  	return &ControlMessageRateLimiter{
    35  		RateLimiter: ratelimiter.NewRateLimiter(limit, burst, lockoutDuration),
    36  	}
    37  }
    38  
    39  // Allow checks the cached limiter for the peer and returns limiter.Allow().
    40  // If a limiter is not cached for a peer one is created.
    41  func (c *ControlMessageRateLimiter) Allow(peerID peer.ID, n int) bool {
    42  	limiter := c.GetLimiter(peerID)
    43  	if !limiter.AllowN(time.Now(), n) {
    44  		c.UpdateLastRateLimit(peerID, time.Now())
    45  		return false
    46  	}
    47  
    48  	return true
    49  }