token-throttler

Token throttler

Token throttler is an extendable rate-limiting library somewhat based on a token bucket algorithm.

Table of contents

1. Installation
2. Features
3. Usage
   1. Manual usage example
   2. Decorator usage example
   3. FastAPI usage example
4. Storage
5. Configuration
   1. Global configuration usage
   2. Instance configuration usage

1. Installation

Token throttler is available on PyPI:

$ python -m pip install token-throttler

Token throttler officially supports Python >= 3.8.

NOTE: Depending on the storage engine you pick, you can install token throttler with extras:

$ python -m pip install token-throttler[redis]

2. Features

• Blocking (TokenThrottler) and non-blocking (TokenThrottlerAsync)
• Global throttler(s) configuration
• Configurable token throttler cost and identifier
• Multiple buckets per throttler per identifier
• Buckets can be added/removed manually or by a dict configuration
• Manual usage or usage via decorator
• Decorator usage supports async code too
• Custom decorator can be written
• Extendable storage engine (eg. Redis)

3. Usage

Token throttler supports both manual usage and via decorator.

Decorator usage supports both async and sync.

1) Manual usage example:

from token_throttler import TokenBucket, TokenThrottler
from token_throttler.storage import RuntimeStorage

throttler: TokenThrottler = TokenThrottler(1, RuntimeStorage())
throttler.add_bucket("hello_world", TokenBucket(10, 10))
throttler.add_bucket("hello_world", TokenBucket(30, 20))


def hello_world() -> None:
    print("Hello World")


for i in range(10):
    throttler.consume("hello_world")
    hello_world()

if throttler.consume("hello_world"):
    hello_world()
else:
    print(
        f"bucket_one ran out of tokens, retry in: {throttler.wait_time('hello_world')}"
    )

2) Decorator usage example:

from token_throttler import TokenBucket, TokenThrottler, TokenThrottlerException
from token_throttler.storage import RuntimeStorage

throttler: TokenThrottler = TokenThrottler(1, RuntimeStorage())
throttler.add_bucket("hello_world", TokenBucket(10, 10))
throttler.add_bucket("hello_world", TokenBucket(30, 20))


@throttler.enable("hello_world")
def hello_world() -> None:
    print("Hello World")


for i in range(10):
    hello_world()

try:
    hello_world()
except TokenThrottlerException:
    print(
        f"bucket_one ran out of tokens, retry in: {throttler.wait_time('hello_world')}"
    )

3) FastAPI usage example:

from fastapi import Depends, FastAPI, Request
from pydantic import BaseModel
from token_throttler import TokenThrottler, TokenBucket
from token_throttler.storage import RuntimeStorage
from token_throttler.ext.fastapi import FastAPIThrottler

app: FastAPI = FastAPI()
ban_hammer: TokenThrottler = TokenThrottler(cost=1, storage=RuntimeStorage())


class User(BaseModel):
    id: int
    name: str


u1: User = User(id=1, name="Test")
users: list[User] = [u1]


def create_buckets() -> None:
    for user in users:
        ban_hammer.add_bucket(
            str(user.id), TokenBucket(replenish_time=10, max_tokens=10)
        )


# For example purposes only - feel free to load your users with the bucket(s) anywhere else
app.add_event_handler("startup", create_buckets)


# Add your auth logic here
def get_auth_user() -> User:
    return u1


# Since this is a FastAPI dependency, you can have the `Request` object here too if needed (e.g. storing the JWT token in
# a request lifecycle).
# You can also use sub-dependencies like shown below - `auth_user`
async def get_user_id(
        request: Request, auth_user: User = Depends(get_auth_user)
) -> str:
    return str(auth_user.id)


# You can configure the dependency per route, on a router level or on global application level.
# Pass your unique user identifier (e.g. JWT subject, apikey or user's ID) as a callback parameter 
# of the `enable` method.
@app.get(
    "/throttle", dependencies=[Depends(FastAPIThrottler(ban_hammer).enable(get_user_id))]
)
async def read_users():
    return {"detail": "This is throttled URL"}

For other examples see examples directory.

4. Storage

TokenThrottler supports RuntimeStorage and RedisStorage. TokenThrottlerAsync supports RedisStorageAsync

If you want your own storage engine, feel free to extend the token_throttler.storage.BucketStorage or token_throttler.storage.BucketStorageAsync classes.

For storage examples see examples directory.

5. Configuration

Token throttler supports global and per instance configurations.

Configuration params:

• IDENTIFIER_FAIL_SAFE (default False) - if invalid identifier is given as a param for the consume method and IDENTIFIER_FAIL_SAFE is set to True, no KeyError exception will be raised and consume will act like a limitless bucket is being consumed.
• ENABLE_THREAD_LOCK (default False) - if set to True, throttler will acquire a thread lock upon calling consume method and release the lock once the consume is finished. This avoids various race conditions at a slight performance cost.

Global configuration usage

Global configuration means that all the throttler instances will use the same configuration from the ThrottlerConfigGlobal class singleton (!) instance - default_config.

from token_throttler import TokenBucket, TokenThrottler, default_config
from token_throttler.storage import RuntimeStorage

default_config.set({
   "ENABLE_THREAD_LOCK": False,
   "IDENTIFIER_FAIL_SAFE": True,
})
throttler: TokenThrottler = TokenThrottler(1, RuntimeStorage())
throttler.add_bucket("hello_world", TokenBucket(10, 10))
throttler.add_bucket("hello_world", TokenBucket(30, 20))
...

Whenever you change the options inside default_config object, all of the TokenThrottler instances, that are not using the instance specific configuration, will get updated with the new settings.

NOTE: If any modifications are attempted on the default_config object during runtime, a RuntimeWarning will be emitted. The default_config object is designed to be a singleton instance of the ThrottlerConfigGlobal class, responsible for storing and managing configuration settings. To ensure the proper functioning of the throttling mechanism and to maintain the integrity of the configuration, it is recommended to avoid making runtime changes to the default_config object. Any such modifications may result in unexpected behavior or inconsistencies in the throttling behavior. Instead, consider utilizing instance-specific configuration.

Instance configuration usage

Instance configuration is to be used when TokenThrottler instance(s) needs a different configuration that the global one. If no instance configuration is passed, TokenThrottler object will use the global configuration.

from token_throttler import ThrottlerConfig, TokenBucket, TokenThrottler
from token_throttler.storage import RuntimeStorage


throttler: TokenThrottler = TokenThrottler(
    1,
    RuntimeStorage(),
    ThrottlerConfig(
        {
            "ENABLE_THREAD_LOCK": True,
            "IDENTIFIER_FAIL_SAFE": True,
        }
    ),
)
throttler.add_bucket("hello_world", TokenBucket(10, 10))
throttler.add_bucket("hello_world", TokenBucket(30, 20))
...