Token Bucket - Rate Limiting Algorithm

Jul 24, 2023

On this page

The idea
TockenBucket
Using it

Rate limiting is the difference between a service that degrades under load and one that falls over. The token bucket is the algorithm most APIs reach for, because it does something the naive approaches don't: it allows short bursts without giving up control of the average rate.

The idea

Picture a bucket that holds tokens.

The bucket has a capacity — the most tokens it can ever hold.
Tokens are added back at a fixed refill rate.
Every request removes one token. No token, no request.

That's the whole algorithm. Two properties fall out of it:

The long-run rate is the refill rate. Over time you can't consume tokens faster than they're added.
Bursts up to the capacity are free. A client that has been idle finds a full bucket and can spend it all at once.

That second property is the point. A fixed-window counter would reject the eleventh request in a second even if the client had been silent for an hour. The token bucket lets that traffic through, then throttles it back down.

TockenBucket

TokenBucket.js

import chalk from "chalk"
 
export default class TokenBucket {
  constructor(capacity, refillAmount, refillTime) {
    this.capacity = capacity
    this.refillTime = refillTime // amount of time between refills (in sec)
    this.refillAmount = refillAmount // number of tokens to add per refill cycle
    this.db = {}
  }
 
  refillBucket(key) {
    if (this.db[key] === undefined) return null
 
    const { tokens, ts } = this.db[key]
    const currentTime = Date.now()
    const elapsedTime = Math.floor(
      (currentTime - ts) / (this.refillTime * 1000), // convert to seconds
    )
 
    const newTokens = elapsedTime * this.refillAmount
 
    this.db[key] = {
      tokens: Math.min(this.capacity, tokens + newTokens),
      ts: currentTime,
    }
 
    return this.db[key]
  }
 
  createBucket(key) {
    if (this.db[key] === undefined) {
      this.db[key] = {
        tokens: this.capacity,
        ts: Date.now(),
      }
    }
 
    return this.db[key]
  }
 
  handleRequest(key) {
    let bucket = this.createBucket(key)
    const currentTime = Date.now()
 
    // check if the time elapsed since the (convert to seconds)
    const elapsedTime = Math.floor((currentTime - bucket.ts) / 1000)
 
    if (elapsedTime >= this.refillTime) {
      bucket = this.refillBucket(key)
    } else {
      if (bucket?.tokens <= 0) {
        console.log(
          chalk.red(
            `Request[REJECTED] for ${key} (tokens - ${
              bucket.tokens
            }) -- ${new Date().toLocaleTimeString()}\n`,
          ),
        )
        return false
      }
    }
 
    if (!bucket) {
      chalk.red(
        `Request[REJECTED] for ${key} -- ${new Date().toLocaleTimeString()} -- BUCKET NOT FOUND\n`,
      )
      return false
    }
 
    console.log(
      chalk.green(
        `Request[ACCEPTED] for ${key} (tokens - ${
          bucket.tokens
        }) -- ${new Date().toLocaleTimeString()}\n`,
      ),
    )
    bucket.tokens -= 1
    return true
  }
}

Using it

TokenBucket takes a capacity, a refill amount and a refill interval, and exposes a single handleRequest(userId) — one bucket per client, keyed by id.

index.js

import TokenBucket from "./TokenBucket.js"
 
// capacity, refillAmount, refillTime (sec)
// token bucket with 2 token every 1 min
// const bucket = new TokenBucket(4, 2, 60);
 
// token bucket with capacity 4 and ading 4 tokens every 5 sec
// const bucket = new TokenBucket(4, 4, 5);
 
// token bucket with capacity 4 and ading 4 tokens every 2 sec
const bucket = new TokenBucket(4, 4, 2)
 
bucket.handleRequest("user1")
bucket.handleRequest("user1")
bucket.handleRequest("user1")
bucket.handleRequest("user1")
bucket.handleRequest("user1")
 
setTimeout(() => {
  bucket.handleRequest("user1")
  bucket.handleRequest("user1")
  bucket.handleRequest("user1")
  bucket.handleRequest("user1")
  bucket.handleRequest("user1")
  bucket.handleRequest("user1")
 
  setTimeout(() => {
    bucket.handleRequest("user1")
  }, 3000)
}, 3000)