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

     1  package ratelimit
     2  
     3  import (
     4  	"time"
     5  
     6  	"golang.org/x/time/rate"
     7  
     8  	"github.com/libp2p/go-libp2p/core/peer"
     9  
    10  	"github.com/koko1123/flow-go-1/network/p2p"
    11  	"github.com/koko1123/flow-go-1/network/p2p/unicast/ratelimit/internal/limiter_map"
    12  )
    13  
    14  // BandWidthRateLimiter unicast rate limiter that limits the bandwidth that can be sent
    15  // by a peer per some configured interval.
    16  type BandWidthRateLimiter struct {
    17  	limiters                 *limiter_map.RateLimiterMap
    18  	limit                    rate.Limit
    19  	burst                    int
    20  	now                      p2p.GetTimeNow
    21  	rateLimitLockoutDuration time.Duration // the amount of time that has to pass before a peer is allowed to connect
    22  }
    23  
    24  // NewBandWidthRateLimiter returns a new BandWidthRateLimiter. The cleanup loop will be started in a
    25  // separate goroutine and should be stopped by calling Close.
    26  func NewBandWidthRateLimiter(limit rate.Limit, burst int, lockout time.Duration, opts ...p2p.RateLimiterOpt) *BandWidthRateLimiter {
    27  	l := &BandWidthRateLimiter{
    28  		limiters:                 limiter_map.NewLimiterMap(rateLimiterTTL, cleanUpTickInterval),
    29  		limit:                    limit,
    30  		burst:                    burst,
    31  		now:                      time.Now,
    32  		rateLimitLockoutDuration: lockout * time.Second,
    33  	}
    34  
    35  	for _, opt := range opts {
    36  		opt(l)
    37  	}
    38  
    39  	return l
    40  }
    41  
    42  // Allow checks the cached limiter for the peer and returns limiter.AllowN(msg.Size())
    43  // which will check if a peer is able to send a message of msg.Size().
    44  // If a limiter is not cached one is created.
    45  func (b *BandWidthRateLimiter) Allow(peerID peer.ID, msgSize int) bool {
    46  	limiter := b.getLimiter(peerID)
    47  	if !limiter.AllowN(b.now(), msgSize) {
    48  		b.limiters.UpdateLastRateLimit(peerID, b.now())
    49  		return false
    50  	}
    51  
    52  	return true
    53  }
    54  
    55  // IsRateLimited returns true is a peer is currently rate limited.
    56  func (b *BandWidthRateLimiter) IsRateLimited(peerID peer.ID) bool {
    57  	metadata, ok := b.limiters.Get(peerID)
    58  	if !ok {
    59  		return false
    60  	}
    61  	return time.Since(metadata.LastRateLimit()) < b.rateLimitLockoutDuration
    62  }
    63  
    64  // SetTimeNowFunc overrides the default time.Now func with the GetTimeNow func provided.
    65  func (b *BandWidthRateLimiter) SetTimeNowFunc(now p2p.GetTimeNow) {
    66  	b.now = now
    67  }
    68  
    69  // Start starts cleanup loop for underlying caches.
    70  func (b *BandWidthRateLimiter) Start() {
    71  	go b.limiters.CleanupLoop()
    72  }
    73  
    74  // Stop sends cleanup signal to underlying rate limiters and rate limited peers maps. After the rate limiter
    75  // is stopped it can not be reused.
    76  func (b *BandWidthRateLimiter) Stop() {
    77  	b.limiters.Close()
    78  }
    79  
    80  // getLimiter returns limiter for the peerID, if a limiter does not exist one is created and stored.
    81  func (b *BandWidthRateLimiter) getLimiter(peerID peer.ID) *rate.Limiter {
    82  	if metadata, ok := b.limiters.Get(peerID); ok {
    83  		return metadata.Limiter()
    84  	}
    85  
    86  	limiter := rate.NewLimiter(b.limit, b.burst)
    87  	b.limiters.Store(peerID, limiter)
    88  
    89  	return limiter
    90  }