Skip to content

Commit

Permalink
This PR allows users to configure the refill strategy for bucket4j. See
Browse files Browse the repository at this point in the history
[bucket4j documentation](https://bucket4j.com/7.5.0/toc.html#refill).
Greedy is the default refill strategy used by bucket4j, which will
refill tokens back to the bucket4 continuously in proportion to the time
elapsed.
On the other hand, Interval refill strategy tops off the bucket at the
end of the interval, no matter when the last token was used.

GitOrigin-RevId: 3213b0bdf9bedadf6754dc0741910907e9bfca8f
  • Loading branch information
yyuan-squareup authored and svc-squareup-copybara committed Dec 11, 2024
1 parent d392e63 commit b1be9c9
Show file tree
Hide file tree
Showing 6 changed files with 129 additions and 5 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ import redis.clients.jedis.ConnectionPoolConfig
import wisp.deployment.TESTING
import wisp.ratelimiting.RateLimiter
import wisp.ratelimiting.testing.TestRateLimitConfig
import wisp.ratelimiting.testing.TestRateLimitConfigRefillGreedily

@MiskTest(startService = true)
class RedisRateLimiterTest {
Expand Down Expand Up @@ -100,6 +101,70 @@ class RedisRateLimiterTest {
assertThat(counter).isEqualTo(10)
}

@Test
fun `test bucket refilled at the end of the interval after consuming all tokens`() {
val increment = TestRateLimitConfig.refillPeriod.dividedBy(5)
repeat(5) {
val result = rateLimiter.consumeToken(KEY, TestRateLimitConfig)
assertThat(result.didConsume).isTrue()
assertThat(result.remaining).isEqualTo(TestRateLimitConfig.capacity - 1 - it)
fakeClock.add(increment)
}

assertThat(rateLimiter.availableTokens(KEY, TestRateLimitConfig)).isEqualTo(5L)
}

@Test
fun `test bucket refilled at the end of the interval after consuming some tokens`() {
val increment = TestRateLimitConfig.refillPeriod.dividedBy(5)
repeat(3) {
val result = rateLimiter.consumeToken(KEY, TestRateLimitConfig)
assertThat(result.didConsume).isTrue()
assertThat(result.remaining).isEqualTo(TestRateLimitConfig.capacity - 1 - it)
fakeClock.add(increment)
}

assertThat(rateLimiter.availableTokens(KEY, TestRateLimitConfig)).isEqualTo(2L)
fakeClock.add(increment)
assertThat(rateLimiter.availableTokens(KEY, TestRateLimitConfig)).isEqualTo(2L)
fakeClock.add(increment) // the clock now has past the end of the interval
assertThat(rateLimiter.availableTokens(KEY, TestRateLimitConfig)).isEqualTo(5L)
}

@Test
fun `test bucket refilled continuously after each increment`() {
repeat(5) {
val result = rateLimiter.consumeToken(KEY, TestRateLimitConfigRefillGreedily)
assertThat(result.didConsume).isTrue()
assertThat(result.remaining).isEqualTo(TestRateLimitConfigRefillGreedily.capacity - 1 - it)
}
assertThat(rateLimiter.availableTokens(KEY, TestRateLimitConfigRefillGreedily)).isEqualTo(0L)
assertThat(rateLimiter.consumeToken(KEY, TestRateLimitConfigRefillGreedily).didConsume).isFalse()

val increment = TestRateLimitConfigRefillGreedily.refillPeriod.dividedBy(5)
repeat(5) {
// One token is added back after each increment
fakeClock.add(increment)
assertThat(rateLimiter.availableTokens(KEY, TestRateLimitConfigRefillGreedily)).isEqualTo(it + 1L)
}
}

@Test
fun `test bucket refilled continuously`() {
val increment = TestRateLimitConfigRefillGreedily.refillPeriod.dividedBy(5)
repeat(5) {
val result = rateLimiter.consumeToken(KEY, TestRateLimitConfigRefillGreedily)
assertThat(result.didConsume).isTrue()
assertThat(result.remaining).isEqualTo(TestRateLimitConfigRefillGreedily.capacity - 1)

assertThat(rateLimiter.availableTokens(KEY, TestRateLimitConfigRefillGreedily)).isEqualTo(4L)
fakeClock.add(increment)
assertThat(rateLimiter.availableTokens(KEY, TestRateLimitConfigRefillGreedily)).isEqualTo(5L)
}

assertThat(rateLimiter.availableTokens(KEY, TestRateLimitConfigRefillGreedily)).isEqualTo(5L)
}

companion object {
private const val KEY = "test_key"
}
Expand Down
21 changes: 21 additions & 0 deletions wisp/wisp-rate-limiting/api/wisp-rate-limiting.api
Original file line number Diff line number Diff line change
@@ -1,12 +1,22 @@
public final class wisp/ratelimiting/RateLimitBucketRefillStrategy : java/lang/Enum {
public static final field GREEDY Lwisp/ratelimiting/RateLimitBucketRefillStrategy;
public static final field INTERVAL Lwisp/ratelimiting/RateLimitBucketRefillStrategy;
public static fun getEntries ()Lkotlin/enums/EnumEntries;
public static fun valueOf (Ljava/lang/String;)Lwisp/ratelimiting/RateLimitBucketRefillStrategy;
public static fun values ()[Lwisp/ratelimiting/RateLimitBucketRefillStrategy;
}

public abstract interface class wisp/ratelimiting/RateLimitConfiguration {
public abstract fun getCapacity ()J
public abstract fun getName ()Ljava/lang/String;
public abstract fun getRefillAmount ()J
public abstract fun getRefillPeriod ()Ljava/time/Duration;
public abstract fun getRefillStrategy ()Lwisp/ratelimiting/RateLimitBucketRefillStrategy;
public abstract fun getVersion ()Ljava/lang/Long;
}

public final class wisp/ratelimiting/RateLimitConfiguration$DefaultImpls {
public static fun getRefillStrategy (Lwisp/ratelimiting/RateLimitConfiguration;)Lwisp/ratelimiting/RateLimitBucketRefillStrategy;
public static fun getVersion (Lwisp/ratelimiting/RateLimitConfiguration;)Ljava/lang/Long;
}

Expand Down Expand Up @@ -124,6 +134,17 @@ public final class wisp/ratelimiting/testing/TestRateLimitConfig : wisp/ratelimi
public fun getName ()Ljava/lang/String;
public fun getRefillAmount ()J
public fun getRefillPeriod ()Ljava/time/Duration;
public fun getRefillStrategy ()Lwisp/ratelimiting/RateLimitBucketRefillStrategy;
public fun getVersion ()Ljava/lang/Long;
}

public final class wisp/ratelimiting/testing/TestRateLimitConfigRefillGreedily : wisp/ratelimiting/RateLimitConfiguration {
public static final field INSTANCE Lwisp/ratelimiting/testing/TestRateLimitConfigRefillGreedily;
public fun getCapacity ()J
public fun getName ()Ljava/lang/String;
public fun getRefillAmount ()J
public fun getRefillPeriod ()Ljava/time/Duration;
public fun getRefillStrategy ()Lwisp/ratelimiting/RateLimitBucketRefillStrategy;
public fun getVersion ()Ljava/lang/Long;
}

Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import io.github.bucket4j.distributed.BucketProxy
import io.github.bucket4j.distributed.proxy.ProxyManager
import io.micrometer.core.instrument.MeterRegistry
import io.micrometer.core.instrument.Metrics
import wisp.ratelimiting.RateLimitBucketRefillStrategy
import wisp.ratelimiting.RateLimitConfiguration
import wisp.ratelimiting.RateLimiter
import wisp.ratelimiting.RateLimiterMetrics
Expand Down Expand Up @@ -115,10 +116,19 @@ class Bucket4jRateLimiter @JvmOverloads constructor(
}

private fun RateLimitConfiguration.toBandwidth(): Bandwidth {
return Bandwidth.builder()
.capacity(capacity)
.refillIntervally(refillAmount, refillPeriod)
.initialTokens(capacity)
.build()
return if (refillStrategy == RateLimitBucketRefillStrategy.GREEDY) {
Bandwidth.builder()
.capacity(capacity)
.refillGreedy(refillAmount, refillPeriod)
.initialTokens(capacity)
.build()
}
else {
Bandwidth.builder()
.capacity(capacity)
.refillIntervally(refillAmount, refillPeriod)
.initialTokens(capacity)
.build()
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
package wisp.ratelimiting

enum class RateLimitBucketRefillStrategy {
/*
* The bucket will be filled continuously at the specified rate
*/
GREEDY,
/*
* The bucket will be topped off at the end of the interval,
* no matter when the last token was consumed.
*/
INTERVAL
}
Original file line number Diff line number Diff line change
Expand Up @@ -34,4 +34,7 @@ interface RateLimitConfiguration {
*/
val version: Long?
get() = null // returns null to be backward compatible

val refillStrategy: RateLimitBucketRefillStrategy
get() = RateLimitBucketRefillStrategy.INTERVAL
}
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package wisp.ratelimiting.testing

import wisp.ratelimiting.RateLimitBucketRefillStrategy
import wisp.ratelimiting.RateLimitConfiguration
import java.time.Duration

Expand All @@ -12,3 +13,14 @@ object TestRateLimitConfig : RateLimitConfiguration {
override val refillAmount = BUCKET_CAPACITY
override val refillPeriod: Duration = REFILL_DURATION
}

object TestRateLimitConfigRefillGreedily : RateLimitConfiguration {
private const val BUCKET_CAPACITY = 5L
private val REFILL_DURATION: Duration = Duration.ofSeconds(30L)

override val capacity = BUCKET_CAPACITY
override val name = "test_configuration_refill_greedily"
override val refillAmount = BUCKET_CAPACITY
override val refillPeriod: Duration = REFILL_DURATION
override val refillStrategy: RateLimitBucketRefillStrategy = RateLimitBucketRefillStrategy.GREEDY
}

0 comments on commit b1be9c9

Please sign in to comment.