How to Implement a Token Bucket Rate Limiter with asyncio in Python

import asyncio
import time


class TokenBucket:
    def __init__(self, rate_per_second, capacity):
        self.rate = rate_per_second
        self.capacity = capacity
        self.tokens = capacity
        self.last_refill = time.monotonic()
        self.lock = asyncio.Lock()

    async def acquire(self):
        async with self.lock:
            now = time.monotonic()
            elapsed = now - self.last_refill
            self.tokens = min(self.capacity, self.tokens + elapsed * self.rate)
            self.last_refill = now

            if self.tokens < 1:
                wait_time = (1 - self.tokens) / self.rate
                await asyncio.sleep(wait_time)
                self.tokens = 0
                self.last_refill = time.monotonic()
            else:
                self.tokens -= 1


async def worker(name, bucket, delay):
    for i in range(3):
        await bucket.acquire()
        print(f"Worker {name}: task {i} at {time.time():.2f}")
        await asyncio.sleep(delay)


async def main():
    bucket = TokenBucket(rate_per_second=2, capacity=2)
    tasks = [
        asyncio.create_task(worker("A", bucket, 0.3)),
        asyncio.create_task(worker("B", bucket, 0.3)),
    ]
    await asyncio.gather(*tasks)


if __name__ == "__main__":
    asyncio.run(main())