Skip to content

Implement a token bucket used for adaptive retries#1

Merged
ubaskota merged 5 commits intoadaptive_retriesfrom
token_bucket
Jan 8, 2026
Merged

Implement a token bucket used for adaptive retries#1
ubaskota merged 5 commits intoadaptive_retriesfrom
token_bucket

Conversation

@ubaskota
Copy link
Owner

Issue #, if available:

Description of changes:
Implements an async token bucket that will be used for client side rate limiting in Adaptive retries. The TokenBucket class provides configurable rate limiting with automatic token refill over time.

Testing:

  • All tests pass with proper mocking to avoid time dependencies
  • Concurrent test verifies no deadlocks
  • Type checking and linting pass without errors

Known Limitation:

  • Potential resource starvation for large requests will be addressed while implementation a retry strategy for Adaptive retries.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

Copy link

@nateprewitt nateprewitt left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good start! I have some questions and feedback. Let me know if you need anything clarified or if I'm missing something.

Comment on lines +354 to +355
async with token_bucket._lock: # type: ignore
token_bucket._refill() # type: ignore

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we talked about this previously, but we don't want to be calling private methods if we can avoid it. Is there another way we can test this that better reflects the public interfaces?

Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, we did talk about this earlier. I updated the other tests to use a public method, but for this we are specifically testing whether the refill exceeds the max capacity when the rate is increased, so I decided to leave it as it is.

Copy link

@jonathan343 jonathan343 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

very minor nit. Looks good to me otherwise

Comment on lines +488 to +494
"""
TokenBucket provides a collection of arbitrary tokens while managing issuance
and refilling over time. This is controlled by a fill rate that can be variably
adjusted. When tokens aren't available, the bucket will enforce a delay before
attempting to reacquire tokens until one is available or the defined timeout is
reached.
"""

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit - See the multi-line docstring guidance from PEP 257 below:

Multi-line docstrings consist of a summary line just like a one-line docstring, followed by a blank line, followed by a more elaborate description. The summary line may be used by automatic indexing tools; it is important that it fits on one line and is separated from the rest of the docstring by a blank line. The summary line may be on the same line as the opening quotes or on the next line. The entire docstring is indented the same as the quotes at its first line (see example below).

I've provided my suggestion below:

Suggested change
"""
TokenBucket provides a collection of arbitrary tokens while managing issuance
and refilling over time. This is controlled by a fill rate that can be variably
adjusted. When tokens aren't available, the bucket will enforce a delay before
attempting to reacquire tokens until one is available or the defined timeout is
reached.
"""
"""Token bucket for rate limiting with configurable fill rate.
Manages token issuance and automatic refilling over time. When tokens aren't
available, enforces a delay until tokens are refilled or timeout is reached.
"""

@ubaskota ubaskota merged commit 97d6286 into adaptive_retries Jan 8, 2026
4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants