Skip to content

Commit b1be9c9

Browse files
yyuan-squareupsvc-squareup-copybara
yyuan-squareup
authored and
svc-squareup-copybara
committed
This PR allows users to configure the refill strategy for bucket4j. See
[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
1 parent d392e63 commit b1be9c9

File tree

6 files changed

+129
-5
lines changed

6 files changed

+129
-5
lines changed

‎misk-rate-limiting-bucket4j-redis/src/test/kotlin/misk/ratelimiting/bucket4j/redis/RedisRateLimiterTest.kt‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import redis.clients.jedis.ConnectionPoolConfig
1919
import wisp.deployment.TESTING
2020
import wisp.ratelimiting.RateLimiter
2121
import wisp.ratelimiting.testing.TestRateLimitConfig
22+
import wisp.ratelimiting.testing.TestRateLimitConfigRefillGreedily
2223

2324
@MiskTest(startService = true)
2425
class RedisRateLimiterTest {
@@ -100,6 +101,70 @@ class RedisRateLimiterTest {
100101
assertThat(counter).isEqualTo(10)
101102
}
102103

104+
@Test
105+
fun `test bucket refilled at the end of the interval after consuming all tokens`() {
106+
val increment = TestRateLimitConfig.refillPeriod.dividedBy(5)
107+
repeat(5) {
108+
val result = rateLimiter.consumeToken(KEY, TestRateLimitConfig)
109+
assertThat(result.didConsume).isTrue()
110+
assertThat(result.remaining).isEqualTo(TestRateLimitConfig.capacity - 1 - it)
111+
fakeClock.add(increment)
112+
}
113+
114+
assertThat(rateLimiter.availableTokens(KEY, TestRateLimitConfig)).isEqualTo(5L)
115+
}
116+
117+
@Test
118+
fun `test bucket refilled at the end of the interval after consuming some tokens`() {
119+
val increment = TestRateLimitConfig.refillPeriod.dividedBy(5)
120+
repeat(3) {
121+
val result = rateLimiter.consumeToken(KEY, TestRateLimitConfig)
122+
assertThat(result.didConsume).isTrue()
123+
assertThat(result.remaining).isEqualTo(TestRateLimitConfig.capacity - 1 - it)
124+
fakeClock.add(increment)
125+
}
126+
127+
assertThat(rateLimiter.availableTokens(KEY, TestRateLimitConfig)).isEqualTo(2L)
128+
fakeClock.add(increment)
129+
assertThat(rateLimiter.availableTokens(KEY, TestRateLimitConfig)).isEqualTo(2L)
130+
fakeClock.add(increment) // the clock now has past the end of the interval
131+
assertThat(rateLimiter.availableTokens(KEY, TestRateLimitConfig)).isEqualTo(5L)
132+
}
133+
134+
@Test
135+
fun `test bucket refilled continuously after each increment`() {
136+
repeat(5) {
137+
val result = rateLimiter.consumeToken(KEY, TestRateLimitConfigRefillGreedily)
138+
assertThat(result.didConsume).isTrue()
139+
assertThat(result.remaining).isEqualTo(TestRateLimitConfigRefillGreedily.capacity - 1 - it)
140+
}
141+
assertThat(rateLimiter.availableTokens(KEY, TestRateLimitConfigRefillGreedily)).isEqualTo(0L)
142+
assertThat(rateLimiter.consumeToken(KEY, TestRateLimitConfigRefillGreedily).didConsume).isFalse()
143+
144+
val increment = TestRateLimitConfigRefillGreedily.refillPeriod.dividedBy(5)
145+
repeat(5) {
146+
// One token is added back after each increment
147+
fakeClock.add(increment)
148+
assertThat(rateLimiter.availableTokens(KEY, TestRateLimitConfigRefillGreedily)).isEqualTo(it + 1L)
149+
}
150+
}
151+
152+
@Test
153+
fun `test bucket refilled continuously`() {
154+
val increment = TestRateLimitConfigRefillGreedily.refillPeriod.dividedBy(5)
155+
repeat(5) {
156+
val result = rateLimiter.consumeToken(KEY, TestRateLimitConfigRefillGreedily)
157+
assertThat(result.didConsume).isTrue()
158+
assertThat(result.remaining).isEqualTo(TestRateLimitConfigRefillGreedily.capacity - 1)
159+
160+
assertThat(rateLimiter.availableTokens(KEY, TestRateLimitConfigRefillGreedily)).isEqualTo(4L)
161+
fakeClock.add(increment)
162+
assertThat(rateLimiter.availableTokens(KEY, TestRateLimitConfigRefillGreedily)).isEqualTo(5L)
163+
}
164+
165+
assertThat(rateLimiter.availableTokens(KEY, TestRateLimitConfigRefillGreedily)).isEqualTo(5L)
166+
}
167+
103168
companion object {
104169
private const val KEY = "test_key"
105170
}

‎wisp/wisp-rate-limiting/api/wisp-rate-limiting.api‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,22 @@
1+
public final class wisp/ratelimiting/RateLimitBucketRefillStrategy : java/lang/Enum {
2+
public static final field GREEDY Lwisp/ratelimiting/RateLimitBucketRefillStrategy;
3+
public static final field INTERVAL Lwisp/ratelimiting/RateLimitBucketRefillStrategy;
4+
public static fun getEntries ()Lkotlin/enums/EnumEntries;
5+
public static fun valueOf (Ljava/lang/String;)Lwisp/ratelimiting/RateLimitBucketRefillStrategy;
6+
public static fun values ()[Lwisp/ratelimiting/RateLimitBucketRefillStrategy;
7+
}
8+
19
public abstract interface class wisp/ratelimiting/RateLimitConfiguration {
210
public abstract fun getCapacity ()J
311
public abstract fun getName ()Ljava/lang/String;
412
public abstract fun getRefillAmount ()J
513
public abstract fun getRefillPeriod ()Ljava/time/Duration;
14+
public abstract fun getRefillStrategy ()Lwisp/ratelimiting/RateLimitBucketRefillStrategy;
615
public abstract fun getVersion ()Ljava/lang/Long;
716
}
817

918
public final class wisp/ratelimiting/RateLimitConfiguration$DefaultImpls {
19+
public static fun getRefillStrategy (Lwisp/ratelimiting/RateLimitConfiguration;)Lwisp/ratelimiting/RateLimitBucketRefillStrategy;
1020
public static fun getVersion (Lwisp/ratelimiting/RateLimitConfiguration;)Ljava/lang/Long;
1121
}
1222

@@ -124,6 +134,17 @@ public final class wisp/ratelimiting/testing/TestRateLimitConfig : wisp/ratelimi
124134
public fun getName ()Ljava/lang/String;
125135
public fun getRefillAmount ()J
126136
public fun getRefillPeriod ()Ljava/time/Duration;
137+
public fun getRefillStrategy ()Lwisp/ratelimiting/RateLimitBucketRefillStrategy;
138+
public fun getVersion ()Ljava/lang/Long;
139+
}
140+
141+
public final class wisp/ratelimiting/testing/TestRateLimitConfigRefillGreedily : wisp/ratelimiting/RateLimitConfiguration {
142+
public static final field INSTANCE Lwisp/ratelimiting/testing/TestRateLimitConfigRefillGreedily;
143+
public fun getCapacity ()J
144+
public fun getName ()Ljava/lang/String;
145+
public fun getRefillAmount ()J
146+
public fun getRefillPeriod ()Ljava/time/Duration;
147+
public fun getRefillStrategy ()Lwisp/ratelimiting/RateLimitBucketRefillStrategy;
127148
public fun getVersion ()Ljava/lang/Long;
128149
}
129150

‎wisp/wisp-rate-limiting/bucket4j/src/main/kotlin/wisp/ratelimiting/bucket4j/Bucket4jRateLimiter.kt‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import io.github.bucket4j.distributed.BucketProxy
99
import io.github.bucket4j.distributed.proxy.ProxyManager
1010
import io.micrometer.core.instrument.MeterRegistry
1111
import io.micrometer.core.instrument.Metrics
12+
import wisp.ratelimiting.RateLimitBucketRefillStrategy
1213
import wisp.ratelimiting.RateLimitConfiguration
1314
import wisp.ratelimiting.RateLimiter
1415
import wisp.ratelimiting.RateLimiterMetrics
@@ -115,10 +116,19 @@ class Bucket4jRateLimiter @JvmOverloads constructor(
115116
}
116117

117118
private fun RateLimitConfiguration.toBandwidth(): Bandwidth {
118-
return Bandwidth.builder()
119-
.capacity(capacity)
120-
.refillIntervally(refillAmount, refillPeriod)
121-
.initialTokens(capacity)
122-
.build()
119+
return if (refillStrategy == RateLimitBucketRefillStrategy.GREEDY) {
120+
Bandwidth.builder()
121+
.capacity(capacity)
122+
.refillGreedy(refillAmount, refillPeriod)
123+
.initialTokens(capacity)
124+
.build()
125+
}
126+
else {
127+
Bandwidth.builder()
128+
.capacity(capacity)
129+
.refillIntervally(refillAmount, refillPeriod)
130+
.initialTokens(capacity)
131+
.build()
132+
}
123133
}
124134
}

‎wisp/wisp-rate-limiting/src/main/kotlin/wisp/ratelimiting/RateLimitBucketRefillStrategy.kt‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
package wisp.ratelimiting
2+
3+
enum class RateLimitBucketRefillStrategy {
4+
/*
5+
* The bucket will be filled continuously at the specified rate
6+
*/
7+
GREEDY,
8+
/*
9+
* The bucket will be topped off at the end of the interval,
10+
* no matter when the last token was consumed.
11+
*/
12+
INTERVAL
13+
}

‎wisp/wisp-rate-limiting/src/main/kotlin/wisp/ratelimiting/RateLimitConfiguration.kt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,4 +34,7 @@ interface RateLimitConfiguration {
3434
*/
3535
val version: Long?
3636
get() = null // returns null to be backward compatible
37+
38+
val refillStrategy: RateLimitBucketRefillStrategy
39+
get() = RateLimitBucketRefillStrategy.INTERVAL
3740
}

‎wisp/wisp-rate-limiting/src/testFixtures/kotlin/wisp/ratelimiting/testing/TestRateLimitConfig.kt‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package wisp.ratelimiting.testing
22

3+
import wisp.ratelimiting.RateLimitBucketRefillStrategy
34
import wisp.ratelimiting.RateLimitConfiguration
45
import java.time.Duration
56

@@ -12,3 +13,14 @@ object TestRateLimitConfig : RateLimitConfiguration {
1213
override val refillAmount = BUCKET_CAPACITY
1314
override val refillPeriod: Duration = REFILL_DURATION
1415
}
16+
17+
object TestRateLimitConfigRefillGreedily : RateLimitConfiguration {
18+
private const val BUCKET_CAPACITY = 5L
19+
private val REFILL_DURATION: Duration = Duration.ofSeconds(30L)
20+
21+
override val capacity = BUCKET_CAPACITY
22+
override val name = "test_configuration_refill_greedily"
23+
override val refillAmount = BUCKET_CAPACITY
24+
override val refillPeriod: Duration = REFILL_DURATION
25+
override val refillStrategy: RateLimitBucketRefillStrategy = RateLimitBucketRefillStrategy.GREEDY
26+
}

0 commit comments

Comments
 (0)