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

     1  package ratelimit
     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/network/p2p"
    10  	"github.com/onflow/flow-go/network/p2p/utils/ratelimiter"
    11  )
    12  
    13  // BandWidthRateLimiter unicast rate limiter that limits the bandwidth that can be sent
    14  // by a peer per some configured interval.
    15  type BandWidthRateLimiter struct {
    16  	*ratelimiter.RateLimiter
    17  }
    18  
    19  // NewBandWidthRateLimiter returns a new BandWidthRateLimiter. The cleanup loop will be started in a
    20  // separate goroutine and should be stopped by calling Close.
    21  func NewBandWidthRateLimiter(limit rate.Limit, burst int, lockout time.Duration, opts ...p2p.RateLimiterOpt) *BandWidthRateLimiter {
    22  	l := &BandWidthRateLimiter{
    23  		RateLimiter: ratelimiter.NewRateLimiter(limit, burst, lockout, opts...),
    24  	}
    25  
    26  	return l
    27  }
    28  
    29  // Allow checks the cached limiter for the peer and returns limiter.AllowN(msg.Size())
    30  // which will check if a peer is able to send a message of msg.Size().
    31  // If a limiter is not cached one is created.
    32  func (b *BandWidthRateLimiter) Allow(peerID peer.ID, msgSize int) bool {
    33  	limiter := b.GetLimiter(peerID)
    34  	if !limiter.AllowN(time.Now(), msgSize) {
    35  		b.UpdateLastRateLimit(peerID, time.Now())
    36  		return false
    37  	}
    38  
    39  	return true
    40  }