Skip to content

core

Token Bucket Rate Limiter Implementation

AsyncTokenBucket

Asynchronous Leaky Bucket Rate Limiter

Parameters:

Name Type Description Default
bucket_config BucketConfig | None

Configuration for the token bucket with the max capacity and time period in seconds

 None
max_concurrent int | None

Maximum number of concurrent requests allowed to acquire capacity

 None
Note

This implementation is synchronous and supports bursts up to the capacity within the specified time period

Source code in limitor/token_bucket/core.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
class AsyncTokenBucket:
    """Asynchronous Leaky Bucket Rate Limiter

    Args:
        bucket_config: Configuration for the token bucket with the max capacity and time period in seconds
        max_concurrent: Maximum number of concurrent requests allowed to acquire capacity

    Note:
        This implementation is synchronous and supports bursts up to the capacity within the specified time period
    """

    def __init__(self, bucket_config: BucketConfig | None = None, max_concurrent: int | None = None):
        config = bucket_config or BucketConfig()
        self.capacity = config.capacity
        self.seconds = config.seconds

        self.fill_rate = self.capacity / self.seconds

        self._bucket_level = self.capacity
        self._last_fill = time.monotonic()

        self.max_concurrent = max_concurrent
        self._lock = asyncio.Lock()

    def _fill(self) -> None:
        """Fill the bucket based on the elapsed time since the last fill"""
        now = time.monotonic()
        elapsed = now - self._last_fill
        self._bucket_level = min(self.capacity, self._bucket_level + elapsed * self.fill_rate)
        self._last_fill = now

    def capacity_info(self, amount: float = 1) -> Capacity:
        """Get the current capacity information of the token bucket

        Args:
            amount: The amount of capacity to check for, defaults to 1

        Returns:
            A named tuple indicating if the bucket has enough capacity and how much more is needed
        """
        self._fill()
        # we need at least `amount` tokens to proceed
        needed = amount - self._bucket_level
        return Capacity(has_capacity=needed <= 0, needed_capacity=needed)

    async def _acquire_logic(self, amount: float = 1) -> None:
        """Core logic for acquiring capacity from the token bucket.

        Args:
            amount: The amount of capacity to check for, defaults to 1

        Notes:
            Adding a lock here ensures that the acquire logic is atomic, but it also means that the
                requests are going to be done in the order they were received  i.e. not out-of-order like
                most async programs.
            The benefit is that with multiple concurrent requests, we can ensure that the bucket level
                is updated correctly and that we don't have multiple requests trying to update the bucket level
                at the same time, which could lead to an inconsistent state i.e. a race condition.
        """
        async with self._lock:  # ensures atomicity given we can have multiple concurrent requests
            capacity_info = self.capacity_info(amount=amount)
            while not capacity_info.has_capacity:
                needed = capacity_info.needed_capacity
                # amount we need to wait to leak (either part or the entire capacity)
                # needed is guaranteed to be positive here, so we can use it directly
                wait_time = needed / self.fill_rate
                if wait_time > 0:
                    await asyncio.sleep(wait_time)

                capacity_info = self.capacity_info(amount=amount)

            self._bucket_level -= amount

    async def _semaphore_acquire(self, amount: float = 1) -> None:
        """Acquire capacity using a semaphore to limit concurrency.

        Args:
            amount: The amount of capacity to acquire, defaults to 1
        """
        semaphore = asyncio.Semaphore(self.max_concurrent) if self.max_concurrent else nullcontext()
        async with semaphore:
            await self._acquire_logic(amount)

    async def acquire(self, amount: float = 1, timeout: float | None = None) -> None:
        """Acquire capacity from the token bucket, waiting asynchronously until allowed.

        Supports timeouts and cancellations.

        Args:
            amount: The amount of capacity to acquire, defaults to 1
            timeout: Optional timeout in seconds for the acquire operation

        Raises:
            TimeoutError: If the acquire operation times out after the specified timeout period
        """
        validate_amount(self, amount=amount)

        if timeout is not None:
            try:
                await asyncio.wait_for(self._semaphore_acquire(amount), timeout=timeout)
            except TimeoutError as error:
                raise TimeoutError(f"Acquire timed out after {timeout} seconds for amount={amount}") from error
        else:
            await self._semaphore_acquire(amount)

    async def __aenter__(self) -> AsyncTokenBucket:
        """Enter the context manager, acquiring resources if necessary"""
        await self.acquire()
        return self

    async def __aexit__(self, exc_type: type[BaseException], exc_val: BaseException, exc_tb: TracebackType) -> None:
        """Exit the context manager, releasing any resources if necessary"""
        return None

__aenter__() async

Enter the context manager, acquiring resources if necessary

Source code in limitor/token_bucket/core.py 
203
204
205
206
async def __aenter__(self) -> AsyncTokenBucket:
    """Enter the context manager, acquiring resources if necessary"""
    await self.acquire()
    return self

__aexit__(exc_type, exc_val, exc_tb) async

Exit the context manager, releasing any resources if necessary

Source code in limitor/token_bucket/core.py 
208
209
210
async def __aexit__(self, exc_type: type[BaseException], exc_val: BaseException, exc_tb: TracebackType) -> None:
    """Exit the context manager, releasing any resources if necessary"""
    return None

acquire(amount=1, timeout=None) async

Acquire capacity from the token bucket, waiting asynchronously until allowed.

Supports timeouts and cancellations.

Parameters:

Name Type Description Default
amount float

The amount of capacity to acquire, defaults to 1

 1
timeout float | None

Optional timeout in seconds for the acquire operation

 None

Raises:

Type Description
TimeoutError

If the acquire operation times out after the specified timeout period
Source code in limitor/token_bucket/core.py 
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
async def acquire(self, amount: float = 1, timeout: float | None = None) -> None:
    """Acquire capacity from the token bucket, waiting asynchronously until allowed.

    Supports timeouts and cancellations.

    Args:
        amount: The amount of capacity to acquire, defaults to 1
        timeout: Optional timeout in seconds for the acquire operation

    Raises:
        TimeoutError: If the acquire operation times out after the specified timeout period
    """
    validate_amount(self, amount=amount)

    if timeout is not None:
        try:
            await asyncio.wait_for(self._semaphore_acquire(amount), timeout=timeout)
        except TimeoutError as error:
            raise TimeoutError(f"Acquire timed out after {timeout} seconds for amount={amount}") from error
    else:
        await self._semaphore_acquire(amount)

capacity_info(amount=1)

Get the current capacity information of the token bucket

Parameters:

Name Type Description Default
amount float

The amount of capacity to check for, defaults to 1

 1

Returns:

Type Description
Capacity

A named tuple indicating if the bucket has enough capacity and how much more is needed
Source code in limitor/token_bucket/core.py 
129
130
131
132
133
134
135
136
137
138
139
140
141
def capacity_info(self, amount: float = 1) -> Capacity:
    """Get the current capacity information of the token bucket

    Args:
        amount: The amount of capacity to check for, defaults to 1

    Returns:
        A named tuple indicating if the bucket has enough capacity and how much more is needed
    """
    self._fill()
    # we need at least `amount` tokens to proceed
    needed = amount - self._bucket_level
    return Capacity(has_capacity=needed <= 0, needed_capacity=needed)

SyncTokenBucket

Token Bucket Rate Limiter

Parameters:

Name Type Description Default
bucket_config BucketConfig | None

Configuration for the token bucket with the max capacity and time period in seconds

 None
Note

This implementation is synchronous and supports bursts up to the capacity within the specified time period

Source code in limitor/token_bucket/core.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
class SyncTokenBucket:
    """Token Bucket Rate Limiter

    Args:
        bucket_config: Configuration for the token bucket with the max capacity and time period in seconds

    Note:
        This implementation is synchronous and supports bursts up to the capacity within the specified time period
    """

    def __init__(self, bucket_config: BucketConfig | None = None):
        # import config and set attributes
        config = bucket_config or BucketConfig()
        self.capacity = config.capacity
        self.seconds = config.seconds

        self.fill_rate = self.capacity / self.seconds  # units per second

        self._bucket_level = self.capacity  # current volume of tokens in the bucket
        self._last_fill = time.monotonic()  # last refill time

        # thread-safe
        self._lock = threading.Lock()

    def _fill(self) -> None:
        """Fill the bucket based on the elapsed time since the last fill"""
        now = time.monotonic()
        elapsed = now - self._last_fill
        self._bucket_level = min(self.capacity, self._bucket_level + elapsed * self.fill_rate)
        self._last_fill = now

    def capacity_info(self, amount: float = 1) -> Capacity:
        """Get the current capacity information of the token bucket

        Args:
            amount: The amount of capacity to check for, defaults to 1

        Returns:
            A named tuple indicating if the bucket has enough capacity and how much more is needed
        """
        self._fill()
        # we need at least `amount` tokens to proceed
        needed = amount - self._bucket_level
        return Capacity(has_capacity=needed <= 0, needed_capacity=needed)

    def acquire(self, amount: float = 1) -> None:
        """Acquire capacity from the token bucket, blocking until enough capacity is available.

        This method will block and sleep until the requested amount can be acquired
        without exceeding the bucket's capacity, simulating rate limiting.

        Args:
            amount: The amount of capacity to acquire, defaults to 1

        Notes:
            The while loop is just to make sure nothing funny happens while waiting
        """
        validate_amount(self, amount=amount)

        with self._lock:
            capacity_info = self.capacity_info(amount=amount)
            while not capacity_info.has_capacity:
                needed = capacity_info.needed_capacity
                # amount we need to wait to leak
                # needed is guaranteed to be positive here, so we can use it directly
                wait_time = needed / self.fill_rate
                if wait_time > 0:
                    time.sleep(wait_time)

                capacity_info = self.capacity_info(amount=amount)

            self._bucket_level -= amount

    def __enter__(self) -> SyncTokenBucket:
        """Enter the context manager, acquiring resources if necessary"""
        self.acquire()
        return self

    def __exit__(self, exc_type: type[BaseException], exc_val: BaseException, exc_tb: TracebackType) -> None:
        """Exit the context manager, releasing any resources if necessary"""
        return None

__enter__()

Enter the context manager, acquiring resources if necessary

Source code in limitor/token_bucket/core.py 
88
89
90
91
def __enter__(self) -> SyncTokenBucket:
    """Enter the context manager, acquiring resources if necessary"""
    self.acquire()
    return self

__exit__(exc_type, exc_val, exc_tb)

Exit the context manager, releasing any resources if necessary

Source code in limitor/token_bucket/core.py 
93
94
95
def __exit__(self, exc_type: type[BaseException], exc_val: BaseException, exc_tb: TracebackType) -> None:
    """Exit the context manager, releasing any resources if necessary"""
    return None

acquire(amount=1)

Acquire capacity from the token bucket, blocking until enough capacity is available.

This method will block and sleep until the requested amount can be acquired without exceeding the bucket's capacity, simulating rate limiting.

Parameters:

Name Type Description Default
amount float

The amount of capacity to acquire, defaults to 1

 1
Notes

The while loop is just to make sure nothing funny happens while waiting

Source code in limitor/token_bucket/core.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def acquire(self, amount: float = 1) -> None:
    """Acquire capacity from the token bucket, blocking until enough capacity is available.

    This method will block and sleep until the requested amount can be acquired
    without exceeding the bucket's capacity, simulating rate limiting.

    Args:
        amount: The amount of capacity to acquire, defaults to 1

    Notes:
        The while loop is just to make sure nothing funny happens while waiting
    """
    validate_amount(self, amount=amount)

    with self._lock:
        capacity_info = self.capacity_info(amount=amount)
        while not capacity_info.has_capacity:
            needed = capacity_info.needed_capacity
            # amount we need to wait to leak
            # needed is guaranteed to be positive here, so we can use it directly
            wait_time = needed / self.fill_rate
            if wait_time > 0:
                time.sleep(wait_time)

            capacity_info = self.capacity_info(amount=amount)

        self._bucket_level -= amount

capacity_info(amount=1)

Get the current capacity information of the token bucket

Parameters:

Name Type Description Default
amount float

The amount of capacity to check for, defaults to 1

 1

Returns:

Type Description
Capacity

A named tuple indicating if the bucket has enough capacity and how much more is needed
Source code in limitor/token_bucket/core.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
def capacity_info(self, amount: float = 1) -> Capacity:
    """Get the current capacity information of the token bucket

    Args:
        amount: The amount of capacity to check for, defaults to 1

    Returns:
        A named tuple indicating if the bucket has enough capacity and how much more is needed
    """
    self._fill()
    # we need at least `amount` tokens to proceed
    needed = amount - self._bucket_level
    return Capacity(has_capacity=needed <= 0, needed_capacity=needed)