github.com/koko1123/flow-go-1@v0.29.6/network/p2p/unicast/ratelimit/message_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  // MessageRateLimiter unicast rate limiter that limits the amount of streams that can
    15  // be created per some configured interval. A new stream is created each time a libP2P
    16  // node sends a direct message.
    17  type MessageRateLimiter struct {
    18  	limiters                 *limiter_map.RateLimiterMap
    19  	limit                    rate.Limit
    20  	burst                    int
    21  	now                      p2p.GetTimeNow
    22  	rateLimitLockoutDuration time.Duration // the amount of time that has to pass before a peer is allowed to connect
    23  }
    24  
    25  // NewMessageRateLimiter returns a new MessageRateLimiter. The cleanup loop will be started in a
    26  // separate goroutine and should be stopped by calling Close.
    27  func NewMessageRateLimiter(limit rate.Limit, burst int, lockoutDuration time.Duration, opts ...p2p.RateLimiterOpt) *MessageRateLimiter {
    28  	l := &MessageRateLimiter{
    29  		limiters:                 limiter_map.NewLimiterMap(rateLimiterTTL, cleanUpTickInterval),
    30  		limit:                    limit,
    31  		burst:                    burst,
    32  		now:                      time.Now,
    33  		rateLimitLockoutDuration: lockoutDuration * time.Second,
    34  	}
    35  
    36  	for _, opt := range opts {
    37  		opt(l)
    38  	}
    39  
    40  	return l
    41  }
    42  
    43  // Allow checks the cached limiter for the peer and returns limiter.Allow().
    44  // If a limiter is not cached for a peer one is created.
    45  func (s *MessageRateLimiter) Allow(peerID peer.ID, _ int) bool {
    46  	limiter := s.getLimiter(peerID)
    47  	if !limiter.AllowN(s.now(), 1) {
    48  		s.limiters.UpdateLastRateLimit(peerID, s.now())
    49  		return false
    50  	}
    51  
    52  	return true
    53  }
    54  
    55  // IsRateLimited returns true is a peer is currently rate limited.
    56  func (s *MessageRateLimiter) IsRateLimited(peerID peer.ID) bool {
    57  	metadata, ok := s.limiters.Get(peerID)
    58  	if !ok {
    59  		return false
    60  	}
    61  	return time.Since(metadata.LastRateLimit()) < s.rateLimitLockoutDuration
    62  }
    63  
    64  // Start starts cleanup loop for underlying caches.
    65  func (s *MessageRateLimiter) Start() {
    66  	go s.limiters.CleanupLoop()
    67  }
    68  
    69  // Stop sends cleanup signal to underlying rate limiters and rate limited peers maps. After the rate limiter
    70  // is closed it can not be reused.
    71  func (s *MessageRateLimiter) Stop() {
    72  	s.limiters.Close()
    73  }
    74  
    75  // SetTimeNowFunc overrides the default time.Now func with the GetTimeNow func provided.
    76  func (s *MessageRateLimiter) SetTimeNowFunc(now p2p.GetTimeNow) {
    77  	s.now = now
    78  }
    79  
    80  // getLimiter returns limiter for the peerID, if a limiter does not exist one is created and stored.
    81  func (s *MessageRateLimiter) getLimiter(peerID peer.ID) *rate.Limiter {
    82  	if metadata, ok := s.limiters.Get(peerID); ok {
    83  		return metadata.Limiter()
    84  	}
    85  
    86  	limiter := rate.NewLimiter(s.limit, s.burst)
    87  	s.limiters.Store(peerID, limiter)
    88  
    89  	return limiter
    90  }