github.com/larrabee/ratelimit@v1.0.6-0.20191102113931-712217ec4fdc/real_bucket.go (about)

     1  // Copyright 2014 Canonical Ltd.
     2  // Licensed under the LGPLv3 with static-linking exception.
     3  // See LICENCE file for details.
     4  
     5  // Package ratelimit provides an efficient token bucket implementation
     6  // that can be used to limit the rate of arbitrary things.
     7  // See http://en.wikipedia.org/wiki/Token_bucket.
     8  package ratelimit
     9  
    10  import (
    11  	"math"
    12  	"sync"
    13  	"time"
    14  	_ "unsafe"
    15  )
    16  
    17  // The algorithm that this implementation uses does computational work
    18  // only when tokens are removed from the bucket, and that work completes
    19  // in short, bounded-constant time (RealBucket.Wait benchmarks at 175ns on
    20  // my laptop).
    21  //
    22  // Time is measured in equal measured ticks, a given interval
    23  // (fillInterval) apart. On each tick a number of tokens (quantum) are
    24  // added to the bucket.
    25  //
    26  // When any of the methods are called the bucket updates the number of
    27  // tokens that are in the bucket, and it records the current tick
    28  // number too. Note that it doesn't record the current time - by
    29  // keeping things in units of whole ticks, it's easy to dish out tokens
    30  // at exactly the right intervals as measured from the start time.
    31  //
    32  // This allows us to calculate the number of tokens that will be
    33  // available at some time in the future with a few simple arithmetic
    34  // operations.
    35  //
    36  // The main reason for being able to transfer multiple tokens on each tick
    37  // is so that we can represent rates greater than 1e9 (the resolution of the Go
    38  // time package) tokens per second, but it's also useful because
    39  // it means we can easily represent situations like "a person gets
    40  // five tokens an hour, replenished on the hour".
    41  
    42  // RealBucket represents a token bucket that fills at a predetermined rate.
    43  // Methods on RealBucket may be called concurrently.
    44  type RealBucket struct {
    45  	// startTime holds the moment when the bucket was
    46  	// first created and ticks began.
    47  	startTime int64
    48  
    49  	// capacity holds the overall capacity of the bucket.
    50  	capacity int64
    51  
    52  	// quantum holds how many tokens are added on
    53  	// each tick.
    54  	quantum int64
    55  
    56  	// fillInterval holds the interval between each tick.
    57  	fillInterval time.Duration
    58  
    59  	// mu guards the fields below it.
    60  	mu sync.Mutex
    61  
    62  	// availableTokens holds the number of available
    63  	// tokens as of the associated latestTick.
    64  	// It will be negative when there are consumers
    65  	// waiting for tokens.
    66  	availableTokens int64
    67  
    68  	// latestTick holds the latest tick for which
    69  	// we know the number of tokens in the bucket.
    70  	latestTick int64
    71  }
    72  
    73  // NewBucket returns a new token bucket that fills at the
    74  // rate of one token every fillInterval, up to the given
    75  // maximum capacity. Both arguments must be
    76  // positive. The bucket is initially full.
    77  func NewBucket(fillInterval time.Duration, capacity int64) (*RealBucket, error) {
    78  	return NewBucketWithQuantum(fillInterval, capacity, 1)
    79  }
    80  
    81  // rateMargin specifes the allowed variance of actual
    82  // rate from specified rate. 1% seems reasonable.
    83  const rateMargin = 0.01
    84  
    85  // NewBucketWithRate returns a token bucket that fills the bucket
    86  // at the rate of rate tokens per second up to the given
    87  // maximum capacity. Because of limited clock resolution,
    88  // at high rates, the actual rate may be up to 1% different from the
    89  // specified rate.
    90  func NewBucketWithRate(rate float64, capacity int64) (*RealBucket, error) {
    91  	// Use the same bucket each time through the loop
    92  	// to save allocations.
    93  	tb, err := NewBucketWithQuantum(1, capacity, 1)
    94  	if err != nil {
    95  		return nil, err
    96  	}
    97  	for quantum := int64(1); quantum < 1<<50; quantum = nextQuantum(quantum) {
    98  		fillInterval := time.Duration(1e9 * float64(quantum) / rate)
    99  		if fillInterval <= 0 {
   100  			continue
   101  		}
   102  		tb.fillInterval = fillInterval
   103  		tb.quantum = quantum
   104  		if diff := math.Abs(tb.Rate() - rate); diff/rate <= rateMargin {
   105  			return tb, nil
   106  		}
   107  	}
   108  	return nil, &QuantumError{rate}
   109  }
   110  
   111  // NewBucketWithQuantum is similar to NewBucket, but allows
   112  // the specification of the quantum size - quantum tokens
   113  // are added every fillInterval.
   114  func NewBucketWithQuantum(fillInterval time.Duration, capacity, quantum int64) (*RealBucket, error) {
   115  	if fillInterval <= 0 {
   116  		return nil, &ValueError{FieldFillInterval, fillInterval.Nanoseconds()}
   117  	}
   118  	if capacity <= 0 {
   119  		return nil, &ValueError{FieldCapacity, capacity}
   120  	}
   121  	if quantum <= 0 {
   122  		return nil, &ValueError{FieldQuantum, quantum}
   123  	}
   124  	return &RealBucket{
   125  		startTime:       nanotime(),
   126  		latestTick:      0,
   127  		fillInterval:    fillInterval,
   128  		capacity:        capacity,
   129  		quantum:         quantum,
   130  		availableTokens: capacity,
   131  	}, nil
   132  }
   133  
   134  // nextQuantum returns the next quantum to try after q.
   135  // We grow the quantum exponentially, but slowly, so we
   136  // get a good fit in the lower numbers.
   137  func nextQuantum(q int64) int64 {
   138  	q1 := q * 11 / 10
   139  	if q1 == q {
   140  		q1++
   141  	}
   142  	return q1
   143  }
   144  
   145  // Wait takes count tokens from the bucket, waiting until they are
   146  // available.
   147  func (tb *RealBucket) Wait(count int64) {
   148  	if d := tb.Take(count); d > 0 {
   149  		time.Sleep(d)
   150  	}
   151  }
   152  
   153  // WaitMaxDuration is like Wait except that it will
   154  // only take tokens from the bucket if it needs to wait
   155  // for no greater than maxWait. It reports whether
   156  // any tokens have been removed from the bucket
   157  // If no tokens have been removed, it returns immediately.
   158  func (tb *RealBucket) WaitMaxDuration(count int64, maxWait time.Duration) bool {
   159  	d, ok := tb.TakeMaxDuration(count, maxWait)
   160  	if d > 0 {
   161  		time.Sleep(d)
   162  	}
   163  	return ok
   164  }
   165  
   166  const infinityDuration time.Duration = 0x7fffffffffffffff
   167  
   168  // Take takes count tokens from the bucket without blocking. It returns
   169  // the time that the caller should wait until the tokens are actually
   170  // available.
   171  //
   172  // Note that if the request is irrevocable - there is no way to return
   173  // tokens to the bucket once this method commits us to taking them.
   174  func (tb *RealBucket) Take(count int64) time.Duration {
   175  	tb.mu.Lock()
   176  	d, _ := tb.take(nanotime(), count, infinityDuration)
   177  	tb.mu.Unlock()
   178  	return d
   179  }
   180  
   181  // TakeMaxDuration is like Take, except that
   182  // it will only take tokens from the bucket if the wait
   183  // time for the tokens is no greater than maxWait.
   184  //
   185  // If it would take longer than maxWait for the tokens
   186  // to become available, it does nothing and reports false,
   187  // otherwise it returns the time that the caller should
   188  // wait until the tokens are actually available, and reports
   189  // true.
   190  func (tb *RealBucket) TakeMaxDuration(count int64, maxWait time.Duration) (time.Duration, bool) {
   191  	tb.mu.Lock()
   192  	t, ok := tb.take(nanotime(), count, maxWait)
   193  	tb.mu.Unlock()
   194  	return t, ok
   195  }
   196  
   197  // TakeAvailable takes up to count immediately available tokens from the
   198  // bucket. It returns the number of tokens removed, or zero if there are
   199  // no available tokens. It does not block.
   200  func (tb *RealBucket) TakeAvailable(count int64) int64 {
   201  	tb.mu.Lock()
   202  	avail := tb.takeAvailable(nanotime(), count)
   203  	tb.mu.Unlock()
   204  	return avail
   205  }
   206  
   207  // takeAvailable is the internal version of TakeAvailable - it takes the
   208  // current time as an argument to enable easy testing.
   209  func (tb *RealBucket) takeAvailable(now int64, count int64) int64 {
   210  	if count <= 0 {
   211  		return 0
   212  	}
   213  	tb.adjustavailableTokens(tb.currentTick(now))
   214  	if tb.availableTokens <= 0 {
   215  		return 0
   216  	}
   217  	if count > tb.availableTokens {
   218  		count = tb.availableTokens
   219  	}
   220  	tb.availableTokens -= count
   221  	return count
   222  }
   223  
   224  // Available returns the number of available tokens. It will be negative
   225  // when there are consumers waiting for tokens. Note that if this
   226  // returns greater than zero, it does not guarantee that calls that take
   227  // tokens from the buffer will succeed, as the number of available
   228  // tokens could have changed in the meantime. This method is intended
   229  // primarily for metrics reporting and debugging.
   230  func (tb *RealBucket) Available() int64 {
   231  	return tb.available(nanotime())
   232  }
   233  
   234  // available is the internal version of available - it takes the current time as
   235  // an argument to enable easy testing.
   236  func (tb *RealBucket) available(now int64) int64 {
   237  	tb.mu.Lock()
   238  	tb.adjustavailableTokens(tb.currentTick(now))
   239  	t := tb.availableTokens
   240  	tb.mu.Unlock()
   241  	return t
   242  }
   243  
   244  // Capacity returns the capacity that the bucket was created with.
   245  func (tb *RealBucket) Capacity() int64 {
   246  	return tb.capacity
   247  }
   248  
   249  // Rate returns the fill rate of the bucket, in tokens per second.
   250  func (tb *RealBucket) Rate() float64 {
   251  	return 1e9 * float64(tb.quantum) / float64(tb.fillInterval)
   252  }
   253  
   254  // take is the internal version of Take - it takes the current time as
   255  // an argument to enable easy testing.
   256  func (tb *RealBucket) take(now int64, count int64, maxWait time.Duration) (time.Duration, bool) {
   257  	if count <= 0 {
   258  		return 0, true
   259  	}
   260  
   261  	tick := tb.currentTick(now)
   262  	tb.adjustavailableTokens(tick)
   263  	avail := tb.availableTokens - count
   264  	if avail >= 0 {
   265  		tb.availableTokens = avail
   266  		return 0, true
   267  	}
   268  	// Round up the missing tokens to the nearest multiple
   269  	// of quantum - the tokens won't be available until
   270  	// that tick.
   271  
   272  	// endTick holds the tick when all the requested tokens will
   273  	// become available.
   274  	endTick := tick + (-avail+tb.quantum-1)/tb.quantum
   275  	endTime := tb.startTime + (endTick * int64(tb.fillInterval))
   276  	waitTime := time.Duration(endTime - now)
   277  	if waitTime > maxWait {
   278  		return 0, false
   279  	}
   280  	tb.availableTokens = avail
   281  	return waitTime, true
   282  }
   283  
   284  // currentTick returns the current time tick, measured
   285  // from tb.startTime.
   286  func (tb *RealBucket) currentTick(now int64) int64 {
   287  	return (now - tb.startTime) / int64(tb.fillInterval)
   288  }
   289  
   290  // adjustavailableTokens adjusts the current number of tokens
   291  // available in the bucket at the given time, which must
   292  // be in the future (positive) with respect to tb.latestTick.
   293  func (tb *RealBucket) adjustavailableTokens(tick int64) {
   294  	lastTick := tb.latestTick
   295  	tb.latestTick = tick
   296  	if tb.availableTokens >= tb.capacity {
   297  		return
   298  	}
   299  	tb.availableTokens += (tick - lastTick) * tb.quantum
   300  	if tb.availableTokens > tb.capacity {
   301  		tb.availableTokens = tb.capacity
   302  	}
   303  	return
   304  }
   305  
   306  //go:linkname nanotime runtime.nanotime
   307  func nanotime() int64