dimas-b commented on code in PR #3433:
URL: https://github.com/apache/polaris/pull/3433#discussion_r2699390225
##########
runtime/service/src/test/java/org/apache/polaris/service/ratelimiter/TokenBucketTest.java:
##########
@@ -93,6 +93,46 @@ void testConcurrent() throws InterruptedException {
Assertions.assertEquals(maxTokens, numAcquired.get());
}
+ @Test
+ void testFractionalTokenAccrual() {
+ MutableClock clock = MutableClock.of(Instant.now(), ZoneOffset.UTC);
+ TokenBucket tokenBucket = new TokenBucket(1, 1, clock);
+
+ assertCanAcquire(tokenBucket, 1);
+ assertCannotAcquire(tokenBucket);
+
+ for (int i = 0; i < 9; i++) {
+ clock.add(Duration.ofMillis(100));
+ assertCannotAcquire(tokenBucket);
+ }
+
+ // After 900ms total, we should have 0.9 tokens - not enough yet
+ // Add 100ms more to reach exactly 1 token
+ clock.add(Duration.ofMillis(100));
+ Assertions.assertTrue(tokenBucket.tryAcquire());
Review Comment:
nit: why not `assertCanAcquire(tokenBucket, 1)` as above?
##########
runtime/service/src/main/java/org/apache/polaris/service/ratelimiter/TokenBucket.java:
##########
@@ -25,20 +25,29 @@
* amount of tokens. Each successful "tryAcquire" costs 1 token.
*/
public class TokenBucket {
- private final double tokensPerMilli;
- private final long maxTokens;
+ /** Conversion factor: 1 token = 1000 milli-tokens. */
+ private static final long MILLI_TOKENS_PER_TOKEN = 1000L;
+
+ private static final long MAX_TOKENS_LIMIT = Long.MAX_VALUE /
MILLI_TOKENS_PER_TOKEN;
+
+ private final long tokensPerSecond;
+ private final long maxMilliTokens;
private final InstantSource instantSource;
- private long tokens;
- private long lastTokenGenerationMillis;
+ private long milliTokens;
+ private long lastUpdateMillis;
public TokenBucket(long tokensPerSecond, long maxTokens, InstantSource
instantSource) {
- this.tokensPerMilli = tokensPerSecond / 1000D;
- this.maxTokens = maxTokens;
+ if (maxTokens > MAX_TOKENS_LIMIT) {
+ throw new IllegalArgumentException(
+ "maxTokens must be <= " + MAX_TOKENS_LIMIT + " to avoid overflow");
+ }
+ this.tokensPerSecond = tokensPerSecond;
+ this.maxMilliTokens = maxTokens * MILLI_TOKENS_PER_TOKEN;
this.instantSource = instantSource;
- tokens = maxTokens;
- lastTokenGenerationMillis = instantSource.millis();
+ milliTokens = maxMilliTokens;
Review Comment:
nit: why not `this.***` as with the above assignments?
##########
runtime/service/src/main/java/org/apache/polaris/service/ratelimiter/TokenBucket.java:
##########
@@ -47,15 +56,22 @@ public TokenBucket(long tokensPerSecond, long maxTokens,
InstantSource instantSo
* @return whether a token was successfully acquired and spent
*/
public synchronized boolean tryAcquire() {
- // Grant tokens for the time that has passed since our last tryAcquire()
long t = instantSource.millis();
- long millisPassed = Math.subtractExact(t, lastTokenGenerationMillis);
- lastTokenGenerationMillis = t;
- tokens = Math.min(maxTokens, tokens + ((long) (millisPassed *
tokensPerMilli)));
+ long millisPassed = Math.subtractExact(t, lastUpdateMillis);
+ lastUpdateMillis = t;
Review Comment:
IMHO, it would be more robust (and easier to understand) to advance
`lastUpdateMillis` only when `millisPassed` is large enough to grant at least
one token. We'll move `lastUpdateMillis` only to the point that exactly matches
the boundary that results in that exact grant. This way we can count exact
tokens (not milli tokens), which increases the operational range of token
values (without overflow).
##########
runtime/service/src/main/java/org/apache/polaris/service/ratelimiter/TokenBucket.java:
##########
@@ -47,15 +56,22 @@ public TokenBucket(long tokensPerSecond, long maxTokens,
InstantSource instantSo
* @return whether a token was successfully acquired and spent
*/
public synchronized boolean tryAcquire() {
- // Grant tokens for the time that has passed since our last tryAcquire()
long t = instantSource.millis();
- long millisPassed = Math.subtractExact(t, lastTokenGenerationMillis);
- lastTokenGenerationMillis = t;
- tokens = Math.min(maxTokens, tokens + ((long) (millisPassed *
tokensPerMilli)));
+ long millisPassed = Math.subtractExact(t, lastUpdateMillis);
+ lastUpdateMillis = t;
+
+ long grant = millisPassed * tokensPerSecond;
+ if (tokensPerSecond != 0 && grant / tokensPerSecond != millisPassed) {
+ // Overflow occurred. If this much time passed, the bucket would be full
anyway.
+ milliTokens = maxMilliTokens;
Review Comment:
This case is _not_ tested by `TokenBucketTest`. Is it necessary?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
