#90 Exposing TokenBucket as a standalone component #100
Conversation
kevkevy3000
commented
Nov 23, 2023
kevkevy3000
commented
- Created new class buckets.py that contains the standalone component for TokenBucket
- Changed TokenBucketLimiter and Retry API/component to pull rate limiting functionality from TokenBucket class
…to use the newly exposed TokenBucket component
There was a problem hiding this comment.
@kevkevy3000 Great start, Kev 👍
Please give comments a look to see if they make sense and don't forget about tests 👀
It will probably make sense to have tests specific to TokenBucket (even though we are partially testing it via the existing ratelimiter tests) to ensure that tokens() and empty() properties are showing actual info as time goes on.
hyx/ratelimit/buckets.py
Outdated
| """ | ||
| Token Bucket Logic | ||
| Replenish tokens as time passes on. If tokens are available, executions can be allowed. | ||
| Otherwise, it's going to be rejected with RateLimitExceeded |
There was a problem hiding this comment.
TokenBucket is decoupled from the rate limiting, so it probably doesn't keep using RateLimitExceeded exception here. I would create a new exception like EmptyBucket and throw it here.
hyx/ratelimit/buckets.py
Outdated
| def empty(self) -> bool: | ||
| return self._tokens <= 0 | ||
|
|
||
| async def acquire(self) -> None: |
There was a problem hiding this comment.
What about calling it just take()?
…n more token bucket tests
There was a problem hiding this comment.
@kevkevy3000 this is great Kevin 👍 Just a few more really minor things and we are merging it
hyx/ratelimit/buckets.py
Outdated
| self._tokens = tokens_to_add - 1 | ||
| return | ||
|
|
||
| def __replenish(self) -> None: |
There was a problem hiding this comment.
Could we just use one underscore there? 🙏
| def __replenish(self) -> None: | |
| def _replenish(self) -> None: |
hyx/ratelimit/buckets.py
Outdated
| next_replenish = self._next_replenish_at | ||
| until_next_replenish = next_replenish - now | ||
|
|
||
| if until_next_replenish <= 0: |
There was a problem hiding this comment.
It's better to check for until_next_replenish > 0 first and return right away if so. That will make the code more "flat" which is easier to read 👀 :
if until_next_replenish > 0:
return
# the code to replenish the bucket
hyx/ratelimit/managers.py
Outdated
|
|
||
| self._tokens = self._bucket_size | ||
| self._next_replenish_at = self._loop.time() + self._token_per_secs | ||
| self._token_bucket = TokenBucket(max_executions, per_time_secs, bucket_size) | ||
|
|
||
| @property | ||
| def tokens(self) -> float: |
There was a problem hiding this comment.
Maybe in this case, it's better just to expose the bucket (and it will provide the remaining information around token state).
There was a problem hiding this comment.
Should we just have one property for just token_bucket and remove the existing two properties?
There was a problem hiding this comment.
Should we just have one property for just token_bucket and remove the existing two properties?
Yes, that! I would call the property just bucket
tests/test_ratelimiter/test_api.py
Outdated
| @@ -45,7 +45,7 @@ async def test__ratelimiter__limit_exceeded() -> None: | |||
| async def calc() -> float: | |||
| return 42 | |||
|
|
|||
| with pytest.raises(RateLimitExceeded): | |||
| with pytest.raises(EmptyBucket): | |||
There was a problem hiding this comment.
Could we please keep the RateLimitExceeded exception? It's more tied to the ratelimiter "domain" while the bucket error feels like exposing detail of the underlying implementation to ratelimiter consumers.
There was a problem hiding this comment.
Yeah, so the tokenbucket class throws an empty bucket exception, and the token bucket rate limiter class essentially pulls that logic from the tokenbucket class. Should we just make the acquire function in tokenbucketlimiter catch this exception and throw the ratelimitexceeded exception?
There was a problem hiding this comment.
Should we just make the acquire function in tokenbucketlimiter catch this exception and throw the ratelimitexceeded exception?
@kevkevy3000 yes, that's actually what I meant. Might seem redundant, but it brings some flexibility to the rate limiter component around when to fail with rate limit error.
hyx/ratelimit/buckets.py
Outdated
| return self._tokens <= 0 | ||
|
|
||
| async def take(self) -> None: | ||
| print(self.tokens) |
There was a problem hiding this comment.
Ops, we have left some debug statements 👀
There was a problem hiding this comment.
Can we please name it test_buckets.py to keep it consistent with other test file names?
| await bucket.take() | ||
|
|
||
|
|
||
| async def test__ratelimiter__fully_replenish_after_time_period() -> None: |
There was a problem hiding this comment.
Thank you for adding this case we were discussing in the comments 👍
| await bucket.take() | ||
|
|
||
|
|
||
| async def test__ratelimiter__fully_replenish_after_time_period() -> None: |
There was a problem hiding this comment.
| async def test__ratelimiter__fully_replenish_after_time_period() -> None: | |
| async def test__token_bucket__fully_replenish_after_time_period() -> None: |
hyx/ratelimit/buckets.py
Outdated
| """ | ||
| Token Bucket Logic | ||
| Replenish tokens as time passes on. If tokens are available, executions can be allowed. | ||
| Otherwise, it's going to be rejected with EmptyBucket |
There was a problem hiding this comment.
| Otherwise, it's going to be rejected with EmptyBucket | |
| Otherwise, it's going to be rejected with a EmptyBucket error |
