rate-limiter.py

# Thread-safe (or async-safe) rate limiter for Python
#
# Very simple but effective: no buckets, tokens, queues, or sleep
# Allows for both smooth limiting and windowed bursts
# Intended for a single client (not distributed) with no long-term memory

import threading
from functools import wraps
from threading import Thread, Timer
from typing import Callable, ParamSpec, TypeVar

P = ParamSpec("P")
T = TypeVar("T")


def limit_rate(
    frequency: float, burst: int = 1
) -> Callable[[Callable[P, T]], Callable[P, T]]:
    period = burst / frequency

    def decorator(func: Callable[P, T]) -> Callable[P, T]:
        # Note: if you need a fair lock (maintaining FIFO), use your own
        semaphore = threading.BoundedSemaphore(burst)

        @wraps(func)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
            # Wait for a permit (blocking)
            semaphore.acquire()
            timer = Timer(period, semaphore.release)
            timer.daemon = True  # Allow program to exit without waiting for timer
            timer.start()
            # Proceed with request, timer will release permit in the background
            return func(*args, **kwargs)

        return wrapper

    return decorator


@limit_rate(1 / 2)
def hello() -> None:
    print("hello")


if __name__ == "__main__":
    print("Running sync")
    # prints hello every two seconds
    threads = [Thread(target=hello) for _ in range(5)]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join()


# Equivalent async version, equally simple
# If desired, the two versions could be combined by using `inspect.iscoroutinefunction`

import asyncio
from typing import Awaitable


def limit_rate_async(
    frequency: float, burst: int = 1
) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
    period = burst / frequency

    def decorator(
        func: Callable[P, Awaitable[T]],
    ) -> Callable[P, Awaitable[T]]:
        # Note: if you need a fair lock (maintaining FIFO), use your own
        # or if you always want burst = 1, asyncio.Lock is fair
        semaphore = asyncio.BoundedSemaphore(burst)

        @wraps(func)
        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
            # Wait for a permit (blocking)
            await semaphore.acquire()
            loop = asyncio.get_running_loop()
            loop.call_later(period, semaphore.release)
            # Proceed with request, permit will be released asynchronously
            return await func(*args, **kwargs)

        return wrapper

    return decorator


@limit_rate_async(3, 2)
async def hello_async() -> None:
    print("hello async")


async def main() -> None:
    # prints hello in groups of two, finishes after about 3 seconds
    await asyncio.gather(*(hello_async() for _ in range(9)))


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