Andrew Chen

Andrew Chen

SDE @ StripeMS CS '19 @ UW-MadisonEnthusiastic about creating tools to make life better.

Token Bucket

Back

Credit to: Link

from time import time

class TokenBucket(object):
  """An implementation of the token bucket algorithm.

  >>> bucket = TokenBucket(80, 0.5)
  >>> print bucket.consume(10)
  True
  >>> print bucket.consume(90)
  False
  """
  def __init__(self, tokens, fill_rate):
      """tokens is the total tokens in the bucket. fill_rate is the
      rate in tokens/second that the bucket will be refilled."""
      self.capacity = float(tokens)
      self._tokens = float(tokens)
      self.fill_rate = float(fill_rate)
      self.timestamp = time()

  def consume(self, tokens):
      """Consume tokens from the bucket. Returns True if there were
      sufficient tokens otherwise False."""
      if tokens <= self.tokens:
          self._tokens -= tokens
      else:
          return False
      return True

  def get_tokens(self):
      if self._tokens < self.capacity:
          now = time()
          delta = self.fill_rate * (now - self.timestamp)
          self._tokens = min(self.capacity, self._tokens + delta)
          self.timestamp = now
      return self._tokens
  tokens = property(get_tokens)

if __name__ == '__main__':
  from time import sleep
  bucket = TokenBucket(80, 1)
  print "tokens =", bucket.tokens
  print "consume(10) =", bucket.consume(10)
  print "consume(10) =", bucket.consume(10)
  sleep(1)
  print "tokens =", bucket.tokens
  sleep(1)
  print "tokens =", bucket.tokens
  print "consume(90) =", bucket.consume(90)
  print "tokens =", bucket.tokens