github.com/blend/go-sdk@v1.20220411.3/ratelimiter/leaky_bucket.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package ratelimiter
     9  
    10  import (
    11  	"time"
    12  )
    13  
    14  var (
    15  	_ RateLimiter = (*LeakyBucket)(nil)
    16  )
    17  
    18  // NewLeakyBucket returns a new token bucket rate limiter.
    19  // The rate is formed by `numActions` and `quantum`; the resulting rate is numActions/quantum.
    20  func NewLeakyBucket(numActions int, quantum time.Duration) *LeakyBucket {
    21  	return &LeakyBucket{
    22  		NumActions: numActions,
    23  		Quantum:    quantum,
    24  		Tokens:     make(map[string]*Token),
    25  		Now:        func() time.Time { return time.Now().UTC() },
    26  	}
    27  }
    28  
    29  // LeakyBucket implements the token bucket rate limiting algorithm.
    30  type LeakyBucket struct {
    31  	NumActions int
    32  	Quantum    time.Duration
    33  	Tokens     map[string]*Token
    34  	Now        func() time.Time
    35  }
    36  
    37  // Check returns true if an id has exceeded the rate limit, and false otherwise.
    38  func (lb *LeakyBucket) Check(id string) bool {
    39  	now := lb.Now()
    40  
    41  	if lb.Tokens == nil {
    42  		lb.Tokens = make(map[string]*Token)
    43  	}
    44  
    45  	token, ok := lb.Tokens[id]
    46  	if !ok {
    47  		lb.Tokens[id] = &Token{Count: 1, Last: now}
    48  		return false
    49  	}
    50  
    51  	elapsed := now.Sub(token.Last) // how long since the last call
    52  	// uint64 is used here because of how mantissa bones these calculations in float64
    53  	leakBy := uint64(lb.NumActions) * (uint64(elapsed) / uint64(lb.Quantum))
    54  
    55  	token.Count = token.Count - float64(leakBy) // remove by the rate per quantum
    56  	if token.Count < 0 {
    57  		token.Count = 0
    58  	}
    59  	token.Last = now
    60  	token.Count++
    61  
    62  	return token.Count >= float64(lb.NumActions)
    63  }
    64  
    65  // Token is an individual id's work.
    66  type Token struct {
    67  	Count float64   // the rate adjusted count; initialize at max*rate, remove rate tokens per call
    68  	Last  time.Time // last is used to calculate the elapsed, and subsequently the rate
    69  }