PostgreSQL Implementation of Token Bucket Algorithm.

create_function.sql

CREATE OR REPLACE FUNCTION get_token_from_rate_limited_bucket(rate_id TEXT)
RETURNS BIGINT AS $body$
DECLARE
  new_tk BIGINT := 0;
begin
  WITH calculated_rate_limit AS (
    SELECT
      id AS bucket_id,
      FLOOR(LEAST(tk + refill_per_sec * extract(epoch FROM NOW() - ts), window_seconds * refill_per_sec))::BIGINT as amount
    FROM rate_limiting_token_buckets
    WHERE id = rate_id
  )
  UPDATE rate_limiting_token_buckets
  SET tk = CASE WHEN calculated_rate_limit.amount >= 1
                THEN calculated_rate_limit.amount - 1
                ELSE rate_limiting_token_buckets.tk
          END,
      ts = CASE WHEN calculated_rate_limit.amount >= 1
                THEN NOW()
                ELSE rate_limiting_token_buckets.ts
          end
  FROM calculated_rate_limit
  WHERE rate_limiting_token_buckets.id = calculated_rate_limit.bucket_id
        AND calculated_rate_limit.amount >= 1
  RETURNING calculated_rate_limit.amount INTO new_tk;
  RETURN new_tk;
END;
$body$ LANGUAGE plpgsql;

create_table.sql

CREATE TABLE rate_limiting_token_buckets (
  id TEXT PRIMARY KEY,
  tk BIGINT NOT NULL,
  ts TIMESTAMPTZ NOT NULL,
  refill_per_sec NUMERIC(11, 7) DEFAULT 166, -- 9999.9999999 to -9999.9999999
  window_seconds INT DEFAULT 60
);

PostgreSQL Implementation of Token Bucket Algorithm.md

Token Bucket

This is an opinionated implementation of the Token Bucket algorithm for PostgreSQL.

Why?

I had several tasks, running on several intances, trying to query ChatGPT API at the same time.
I needed a solution to rate limit the number of queries made to OPENAI API, this information had to be available accross several services.
I only had one database (pg) available.

Recommendation

If you do not need to specify refill_per_sec as a decimal number, convert it to INT for improved performances.

Usage

Create a bucket

Example A:

• up to 9000 queries per minute
• duration of window 60 seconds
• allows bursts of 9000 queries

-- id: 'BucketA' -- unique, also used as name
-- tk: 10 -- initial number of tokens available
-- ts: CURRENT_TIMESTAMP -- because it needs an initial value
-- refill_per_sec: 150 -- 150*60= 9000 requests/minute
-- window_seconds: 60 --  window duration in seconds
INSERT INTO rate_limiting_token_buckets (id, tk, ts, refill_per_sec, window_seconds)
VALUES ('BucketA', 10, CURRENT_TIMESTAMP, 150, 60)

Example B:

• up to 200 queries per day
• duration of window 86400 seconds (1 day)
• allows bursts of 200 queries

-- id: 'BucketB' -- unique, also used as name
-- tk: 10 -- initial number of tokens available
-- ts: CURRENT_TIMESTAMP -- because it needs an initial value
-- refill_per_sec: 0.0023149 -- 200/86400= 0.0023149 requests/second
-- window_seconds: 86400 -- window duration in seconds
INSERT INTO rate_limiting_token_buckets (id, tk, ts, refill_per_sec, window_seconds)
VALUES ('BucketB', 10, CURRENT_TIMESTAMP, 0.0023149, 86400)

Get a token from the bucket

SELECT get_token_from_rate_limited_bucket('BucketA');

get_token_from_rate_limited_bucket return values:

• BIGINT OR NULL:
  • When the bucket isn't empty, it returns the number of tokens left including the one you've taken. Therefor:
    • A return value of 1 means you've taken the last token.
  • When the bucket is empty (no more tokens available for the window period), it returns NULL.
    • After that, the bucket will refill at the rate set in refill_per_sec.
  • When the requested bucket doesn't exist, it returns NULL.

FAQ

Can get_token_from_rate_limited_bucket return 0 instead of NULL when there aren't any token available?

Yes, remove this condition before creating the function AND calculated_rate_limit.amount >= 1.

But beware, removing this condition will trigger an UPDATE on the bucket even when it is empty. Therefor:

• The update is recorded in the transaction log, regardless of whether the data was altered.
• The row is locked as part of the update operation.
• Any associated triggers will be executed.

This overhead is typically minimal but can become noticeable in high-volume transactions.