3
\$\begingroup\$

I've written a class that implements something akin to the token bucket algorithm so that I may rate limit aysnc HTTP requests made from my application.

The code works but I'm still not sure if it feels 'right', maybe just because I am new to async programming in Python.

Specific concerns:

  • It feels weird/bad how the state (self.last, self.tokens) is managed (or not managed) meaning checks in the _refill method.
  • Maybe I shouldn't be using a context manager (__aenter__ and __aexit__) for this.
  • I'm not sure if I should be doing any cleaning-up in __aexit__
  • Implementation of the algorithm may not be 100% faithful as I am waiting until the bucket is empty before refilling as a batch . I think this may be OK because only one token can be consumed per request.
"""
Typical usage example:
    async def rateLimitedTask(bucket):
        await bucket.removeToken() #wait for token
        #do the task which has now been rate limited
        ...

    async with TokenBucket(2, 1/5) as bucket:
        #make arbitrary async calls to rateLimitedTask()
        ...
"""

import asyncio
import time
from collections import deque

class TokenBucket():
    """Context manager which provides a token bucket.

    Attributes:
        tokens (collections.deque): The tokens (which manifest as timestamps)
                                    in reverse chronological order.
        last (float):               The most recent token consumed.
        rate (float):               See __init__ arg `tokensPerS`.
        capacity (int):             See __init__ arg `capacity`.
        sleepDuration (float):      See __init__ arg `refillSleepS`.

    """

    def __init__(self, capacity=1, tokensPerS=1.0, refillSleepS=0.1):
        """Initialises the TokenBucket context manager.

        Args:
            capacity (int, optional):       The maxiumum tokens the bucket can
                                            hold. Larger is burstier.
                                            Defaults to 1.
            tokensPerS (float, optional):   The average replenishment rate in
                                            tokens per second.
                                            Defaults to 1.
            refillSleepS (float, optional): If not enough time has passed
                                            to refill; how many seconds to
                                            sleep before trying again.
                                            Defaults to 0.1.
        """
        self.capacity = capacity
        self.rate = tokensPerS
        self.last = None
        self.tokens = deque([])
        self.sleepDuration = refillSleepS

    def hasTokens(self):
        return len(self.tokens) > 0

    def isEmpty(self):
        return len(self.tokens) < 1

    async def _refill(self):
        """Refill the empty bucket back up to capacity.

        Is called only when the bucket is empty.
        The tokens themselves are timestamps of when the token was created.
        Uses monotonic time to safeguard against system clock changes.
        """
        if self.last: #prevents unnecessary run on virgin execution
            while self.capacity / (time.monotonic() - self.last) > self.rate:
                await asyncio.sleep(self.sleepDuration)
                if self.hasTokens():
                    return #another call filled the bucket already

        for _ in range(self.capacity):
            self.tokens.appendleft(time.monotonic())

    async def removeToken(self):
        """Removes a token from the bucket and returns void.

        If the bucket is empty it will await its refill before returning.
        """
        if self.isEmpty():
            await self._refill()

        self.last = self.tokens.pop()

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc, tb):
        pass

Any help would be much appreciated.

\$\endgroup\$
2
\$\begingroup\$

Parent-less class

class TokenBucket():

can drop the parens.

lower_snake_case

These variables:

  • tokensPerS
  • refillSleepS
  • sleepDuration

etc., and all of your methods, should be in lower_snake_case instead of camelCase, by convention.

Type hints

PEP484 allows you to turn this

def hasTokens(self):

into this

def hasTokens(self) -> bool:

and likewise for your other method parameters and returns. Among other things, it will allow you to drop your types from comments like this:

"""
    tokens (collections.deque): The tokens (which manifest as timestamps)
                                in reverse chronological order.
    last (float):               The most recent token consumed.
    rate (float):               See __init__ arg `tokensPerS`.
    capacity (int):             See __init__ arg `capacity`.
    sleepDuration (float):      See __init__ arg `refillSleepS`.
"""

because they'll already be in the signature.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.