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 
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
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
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()
        self._semaphore: asyncio.Semaphore | AbstractAsyncContextManager[Any] | None = (
            None
        )

    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
        """
        if self._semaphore is None:
            self._semaphore = (
                asyncio.Semaphore(self.max_concurrent)
                if self.max_concurrent
                else nullcontext()
            )

        async with self._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)

    @asynccontextmanager
    async def acquire_ctx(
        self, amount: float = 1, timeout: float | None = None
    ) -> AsyncGenerator[AsyncTokenBucket]:
        """Async context manager that acquires a specific amount of capacity.

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

        Yields:
            AsyncGenerator[AsyncTokenBucket]: The AsyncTokenBucket instance
        """
        await self.acquire(amount=amount, timeout=timeout)
        try:
            yield self
        finally:
            pass

    async def reconcile(self, actual: float, estimated: float) -> None:
        """Reconcile the actual tokens/capacity used against the estimated amount.

        Args:
            actual: The actual amount of tokens/capacity used
            estimated: The estimated amount of tokens/capacity previously acquired
        """
        diff = actual - estimated
        if diff > 0:
            await self.acquire(amount=diff)
        elif diff < 0:
            async with self._lock:
                self._fill()
                self._bucket_level = min(self.capacity, self._bucket_level - diff)

    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] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> 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 
301
302
303
304
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 
306
307
308
309
310
311
312
313
async def __aexit__(
    self,
    exc_type: type[BaseException] | None,
    exc_val: BaseException | None,
    exc_tb: TracebackType | None,
) -> 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 
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
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)

acquire_ctx(amount=1, timeout=None) async

Async context manager that acquires a specific amount of capacity.

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

Yields:

Type Description
AsyncGenerator[AsyncTokenBucket]

AsyncGenerator[AsyncTokenBucket]: The AsyncTokenBucket instance
Source code in limitor/token_bucket/core.py 
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
@asynccontextmanager
async def acquire_ctx(
    self, amount: float = 1, timeout: float | None = None
) -> AsyncGenerator[AsyncTokenBucket]:
    """Async context manager that acquires a specific amount of capacity.

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

    Yields:
        AsyncGenerator[AsyncTokenBucket]: The AsyncTokenBucket instance
    """
    await self.acquire(amount=amount, timeout=timeout)
    try:
        yield self
    finally:
        pass

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 
183
184
185
186
187
188
189
190
191
192
193
194
195
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)

reconcile(actual, estimated) async

Reconcile the actual tokens/capacity used against the estimated amount.

Parameters:

Name Type Description Default
actual float

The actual amount of tokens/capacity used

 required
estimated float

The estimated amount of tokens/capacity previously acquired

 required
Source code in limitor/token_bucket/core.py 
286
287
288
289
290
291
292
293
294
295
296
297
298
299
async def reconcile(self, actual: float, estimated: float) -> None:
    """Reconcile the actual tokens/capacity used against the estimated amount.

    Args:
        actual: The actual amount of tokens/capacity used
        estimated: The estimated amount of tokens/capacity previously acquired
    """
    diff = actual - estimated
    if diff > 0:
        await self.acquire(amount=diff)
    elif diff < 0:
        async with self._lock:
            self._fill()
            self._bucket_level = min(self.capacity, self._bucket_level - diff)

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 
 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
 96
 97
 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
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

    @contextmanager
    def acquire_ctx(self, amount: float = 1) -> Generator[SyncTokenBucket]:
        """Context manager that acquires a specific amount of capacity.

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

        Yields:
            Generator[SyncTokenBucket]: The SyncTokenBucket instance
        """
        self.acquire(amount=amount)
        try:
            yield self
        finally:
            pass

    def reconcile(self, actual: float, estimated: float) -> None:
        """Reconcile the actual tokens/capacity used against the estimated amount.

        Args:
            actual: The actual amount of tokens/capacity used
            estimated: The estimated amount of tokens/capacity previously acquired
        """
        diff = actual - estimated
        if diff > 0:
            self.acquire(amount=diff)
        elif diff < 0:
            with self._lock:
                self._fill()
                self._bucket_level = min(self.capacity, self._bucket_level - diff)

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

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> 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 
128
129
130
131
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 
133
134
135
136
137
138
139
140
def __exit__(
    self,
    exc_type: type[BaseException] | None,
    exc_val: BaseException | None,
    exc_tb: TracebackType | None,
) -> 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 
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
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

acquire_ctx(amount=1)

Context manager that acquires a specific amount of capacity.

Parameters:

Name Type Description Default
amount float

The amount of capacity to acquire, defaults to 1

 1

Yields:

Type Description
Generator[SyncTokenBucket]

Generator[SyncTokenBucket]: The SyncTokenBucket instance
Source code in limitor/token_bucket/core.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
@contextmanager
def acquire_ctx(self, amount: float = 1) -> Generator[SyncTokenBucket]:
    """Context manager that acquires a specific amount of capacity.

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

    Yields:
        Generator[SyncTokenBucket]: The SyncTokenBucket instance
    """
    self.acquire(amount=amount)
    try:
        yield self
    finally:
        pass

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 
55
56
57
58
59
60
61
62
63
64
65
66
67
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)

reconcile(actual, estimated)

Reconcile the actual tokens/capacity used against the estimated amount.

Parameters:

Name Type Description Default
actual float

The actual amount of tokens/capacity used

 required
estimated float

The estimated amount of tokens/capacity previously acquired

 required
Source code in limitor/token_bucket/core.py 
113
114
115
116
117
118
119
120
121
122
123
124
125
126
def reconcile(self, actual: float, estimated: float) -> None:
    """Reconcile the actual tokens/capacity used against the estimated amount.

    Args:
        actual: The actual amount of tokens/capacity used
        estimated: The estimated amount of tokens/capacity previously acquired
    """
    diff = actual - estimated
    if diff > 0:
        self.acquire(amount=diff)
    elif diff < 0:
        with self._lock:
            self._fill()
            self._bucket_level = min(self.capacity, self._bucket_level - diff)