self-limiters 0.1.2

Distributed async rate limiters, using Redis

distributed async client rate limiters

This package implements an async distributed semaphore, and the token bucket algorithm. Both implementations are FIFO, and require redis.

Between them, the two implementations make it possible to regulate traffic with respect to:

• Concurrency-based limits (5 active requests at the same time), or

• Time-based limits (5 requests every 10 seconds)

This was written for rate-limiting, but the semaphore and token bucket implementations can be used for anything.

Installation

pip install self-limiters

Usage

Both implementations are written as async context managers.

Semaphore

The Semaphore context manager is used as follows:

from self_limiters import Semaphore


# allows 10 concurrent requests
concurrency_limited_queue = Semaphore(  # <-- init
    name="foo",
    capacity=10,
    max_sleep=15,
    redis_url="redis://localhost:6379"
)

async def main():
    async with concurrency_limited_queue:  # <-- usage
        client.get(...)

A MaxSleepExceededError is raised after max_sleep seconds, if max_sleep is specified.

Token bucket

The TokenBucket context manager is used like this:

from self_limiters import TokenBucket

# allows 10 requests per minute
time_limited_queue = TokenBucket(  # <-- init
    name="unique-resource-name",
    capacity=10,
    refill_frequency=60,
    refill_amount=10,  #
    max_sleep=600,
    redis_url="redis://localhost:6379"
)

async def main():
    async with time_limited_queue:  # <-- usage
        client.get(...)

If max_sleep is set and the sleep time exceeds max-sleep a MaxSleepExceededError is raised.

Decorators

The package doesn't ship any decorators, but if you want to apply these to a function, you can roll your own, like this:

from self_limiters import Semaphore


def limit(name, capacity):
    # Initialize semaphore or token bucket
    limiter = Semaphore(name=name, capacity=capacity, redis_url="redis://127.0.0.1:6389")

    def middle(f):
        async def inner(*args, **kwargs):
            # Use context manager within the decorator
            async with limiter:
                return await f(*args, **kwargs)
        return inner
    return middle


# The decorator will then ensure we sleep until ready
@limit(name="my-user-api", capacity=50)
def retrieve_user(id: UUID) -> User:
    ...

Implementation breakdown

The package was written for performance, and it seems that the most performant way to implement these algorithms is by leveraging Lua scripts. I initially wrote this in pure Rust, but Lua scripts present a couple of really great benefits:

• Lua scripts are executed on a redis instance, and lets us implement close to the entire implementation logic in a single script. This means our client can make 1 request to redis to run the script instead of 1 request per redis call needed. The time saved by reducing the number of requests is huge.

• The initial rust implementation (pre-lua scripts) had to use the redlock algorithm to ensure fairness. With Lua scripts (since redis instance are single threaded), our implementations are FIFO out of the box.

  This makes our implementation a lot faster, since we no longer need locks, and it simultaneously makes the code much, much simpler.

With Lua scripts, this is how our flows ended up being:

The semaphore implementation

1. Run create_semaphore.lua to create a list, which will be the foundation of our semaphore. This is skipped if it has already been created.
2. Run BLPOP to non-blockingly wait until the semaphore has capacity. When it does, we pop from the list.
3. Then run another release_semaphore.lua to "release" the semaphore by adding back the capacity we popped.

So in total we make 3 calls to redis (we would have made 6 without the scripts), which are all non-blocking.

The token bucket implementation

Here, things are even simpler. The steps are:

1. Run schedule.lua to retrieve a wake-up time.
2. Sleep until then.

We make 1 call instead of 3, and both of these are also non-blocking.

In other words, the limiters' impact on an application event-loop should be completely negligible.

Benchmarks

When testing on Github actions we get sub-millisecond runtimes:

• processing 100 nodes with the semaphore implementation takes ~0.6ms per instance
• processing 100 nodes with the token bucket implementation takes ~0.03ms per instance

The majority of this should also still be waiting for i/o.

Implementation details

This section breaks down how both implementations are written, in detail. It assumes you know how a Python (async) context manager works. If this is new to you, take a look here first!

The semaphore implementation

The semaphore implementation is useful when you need to limit a process to n concurrent actions. For example if you have 10 web servers, and you're interacting with an API that will only tolerate 5 concurrent requests before locking you out.

The flow can be broken down as follows:

• The create_semaphore Lua script calls SETNX on the key of the queue plus a postfix (if the name specified in the class instantiation is "my-queue", then the queue name will be __self-limiters:my-queue and setnx will be called for __self-limiters:my-queue-exists). If the returned value is 1 it means the queue we will use for our semaphore does not exist yet and needs to be created.

  (It might strike you as weird to maintain a separate value, just to indicate whether a list exists, when we could just check the list itself. It would be nice if we could use EXISTS on the list directly, but unfortunately a list is considered not to exist when all elements are popped (i.e., when a semaphore is fully acquired), so I don't see another way of doing this. Contributions are very welcome if you do!

• Then if the queue needs to be created we call RPUSH with the number of arguments equal to the capacity value used when initializing the semaphore instance. For a semaphore with a capacity of 5, we call RPUSH 1 1 1 1 1, where the values are completely arbitrary.

• Once the list/queue has been created, we call BLPOP to block until it's our turn. BLPOP is FIFO by default. We also make sure to specify the max_sleep based on the initialized semaphore instance setting. If nothing was passed we allow sleeping forever.

• On __aexit__ we call another script to RPUSH a 1 value back into the queue (i.e., release the semaphore) and set an expiry on the queue and the string value we called SETNX on.

  The expires are a half measure for dealing with dropped capacity. If a node holding the semaphore dies, the capacity might never be returned. If, however, there is no one using the semaphore for the duration of the expiry value, all values will be cleared, and the semaphore will be recreated at full capacity next time it's used. The expiry is 30 seconds at the time of writing, but could be made configurable.

The token bucket implementation

The token bucket implementation is useful when you need to limit a process by time. For example, to 1 request per minute, or 500 every 10 seconds.

The implementation is forward-looking. It works out the time there would have been capacity in the bucket for a given client and returns that time. From there we can asynchronously sleep until it's time to perform our rate limited action.

The flow can be broken down as follows:

• Call the schedule Lua script which first GETs the state of the bucket.

  The bucket state contains the last slot scheduled and the number of tokens left for that slot. With a capacity of 1, having a tokens_left_for_slot variable makes no sense, but if there's capacity of 2 or more, it is possible that we will need to schedule multiple clients to the same slot.

  The script then works out whether to decrement the tokens_left_for_slot value, or to increment the slot value wrt. the frequency variable.

  Finally, we store the bucket state again using SETEX. This allows us to store the state and set expiry at the same time. The default expiry is 30 at the time of writing, but could be made configurable.

  One thing to note, is that this would not work if it wasn't for the fact that redis is single threaded, so Lua scripts on Redis are FIFO. Without this we would need locks and a lot more logic.

• Then we just sleep. Simple!

Contributing

Please do!

Feedback on the implementation, issues, and PRs are welcome. See CONTRIBUTING.md for more details.