github.com/ethersphere/bee/v2@v2.2.0/pkg/ratelimit/ratelimit.go (about)

     1  // Copyright 2021 The Swarm Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package ratelimit provides a mechanism to rate limit requests based on a string key,
     6  // refill rate and burst amount. Under the hood, it's a token bucket of size burst amount,
     7  // that refills at the refill rate.
     8  package ratelimit
     9  
    10  import (
    11  	"context"
    12  	"sync"
    13  	"time"
    14  
    15  	"golang.org/x/time/rate"
    16  )
    17  
    18  type Limiter struct {
    19  	mtx     sync.Mutex
    20  	limiter map[string]*rate.Limiter
    21  	rate    rate.Limit
    22  	burst   int
    23  }
    24  
    25  // New returns a new Limiter object with refresh rate and burst amount
    26  func New(r time.Duration, burst int) *Limiter {
    27  	return &Limiter{
    28  		limiter: make(map[string]*rate.Limiter),
    29  		rate:    rate.Every(r),
    30  		burst:   burst,
    31  	}
    32  }
    33  
    34  // Allow checks if the limiter that belongs to 'key' has not exceeded the limit.
    35  func (l *Limiter) Allow(key string, count int) bool {
    36  	return l.getLimiter(key).AllowN(time.Now(), count)
    37  }
    38  
    39  // Wait blocks until the limiter permits n events to happen. Returns the time duration
    40  // the limiter waited for to allow the number of events to occur.
    41  func (l *Limiter) Wait(ctx context.Context, key string, count int) (time.Duration, error) {
    42  	limiter := l.getLimiter(key)
    43  
    44  	n := time.Now()
    45  
    46  	if limiter.AllowN(n, count) {
    47  		return 0, nil
    48  	}
    49  
    50  	err := limiter.WaitN(ctx, count)
    51  
    52  	return time.Since(n), err
    53  }
    54  
    55  // Clear deletes the limiter that belongs to 'key'
    56  func (l *Limiter) getLimiter(key string) *rate.Limiter {
    57  	l.mtx.Lock()
    58  	defer l.mtx.Unlock()
    59  
    60  	limiter, ok := l.limiter[key]
    61  	if !ok {
    62  		limiter = rate.NewLimiter(l.rate, l.burst)
    63  		l.limiter[key] = limiter
    64  	}
    65  
    66  	return limiter
    67  }
    68  
    69  // Clear deletes the limiter that belongs to 'key'
    70  func (l *Limiter) Clear(key string) {
    71  	l.mtx.Lock()
    72  	defer l.mtx.Unlock()
    73  
    74  	delete(l.limiter, key)
    75  }