From bddeae6470a66d189a0e114002b82278d3c9d88a Mon Sep 17 00:00:00 2001
From: Yue Yuan <yyuan@squareup.com>
Date: Wed, 13 Nov 2024 22:09:30 -0500
Subject: [PATCH] This PR introduces an optional `version` in
 `RateLimitConfiguration`, which allows client code to replace the
 configuration on the fly. Bucket4j checks the version against the persisted
 version, and replace the configuration dynamically if version > persisted
 version. [doc](https://bucket4j.com/8.8.0/toc.html#configuration-replacement)

GitOrigin-RevId: 8662847831bdaba6771062b0ab49699f04bd3bb0
---
 .../com/squareup/exemplar/actions/RateLimitedAction.kt    | 1 +
 wisp/wisp-rate-limiting/api/wisp-rate-limiting.api        | 6 ++++++
 .../wisp/ratelimiting/bucket4j/Bucket4jRateLimiter.kt     | 7 ++++++-
 .../kotlin/wisp/ratelimiting/RateLimitConfiguration.kt    | 8 ++++++++
 4 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/samples/exemplar/src/main/kotlin/com/squareup/exemplar/actions/RateLimitedAction.kt b/samples/exemplar/src/main/kotlin/com/squareup/exemplar/actions/RateLimitedAction.kt
index 6df971fd98e..c04207c5002 100644
--- a/samples/exemplar/src/main/kotlin/com/squareup/exemplar/actions/RateLimitedAction.kt
+++ b/samples/exemplar/src/main/kotlin/com/squareup/exemplar/actions/RateLimitedAction.kt
@@ -41,4 +41,5 @@ object ExampleRateLimitConfiguration : RateLimitConfiguration {
   override val name = "ExpensiveRateLimitedAction"
   override val refillAmount = 10L
   override val refillPeriod: Duration = Duration.ofMinutes(1L)
+  override val version = 0L // increment the version when updating the configuration
 }
diff --git a/wisp/wisp-rate-limiting/api/wisp-rate-limiting.api b/wisp/wisp-rate-limiting/api/wisp-rate-limiting.api
index c4f20fb5ff8..08c5913d157 100644
--- a/wisp/wisp-rate-limiting/api/wisp-rate-limiting.api
+++ b/wisp/wisp-rate-limiting/api/wisp-rate-limiting.api
@@ -3,6 +3,11 @@ public abstract interface class wisp/ratelimiting/RateLimitConfiguration {
 	public abstract fun getName ()Ljava/lang/String;
 	public abstract fun getRefillAmount ()J
 	public abstract fun getRefillPeriod ()Ljava/time/Duration;
+	public abstract fun getVersion ()Ljava/lang/Long;
+}
+
+public final class wisp/ratelimiting/RateLimitConfiguration$DefaultImpls {
+	public static fun getVersion (Lwisp/ratelimiting/RateLimitConfiguration;)Ljava/lang/Long;
 }
 
 public abstract interface class wisp/ratelimiting/RateLimitPruner {
@@ -119,5 +124,6 @@ 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 getVersion ()Ljava/lang/Long;
 }
 
diff --git a/wisp/wisp-rate-limiting/bucket4j/src/main/kotlin/wisp/ratelimiting/bucket4j/Bucket4jRateLimiter.kt b/wisp/wisp-rate-limiting/bucket4j/src/main/kotlin/wisp/ratelimiting/bucket4j/Bucket4jRateLimiter.kt
index 8c6b2a2264f..90ab4a10ff6 100644
--- a/wisp/wisp-rate-limiting/bucket4j/src/main/kotlin/wisp/ratelimiting/bucket4j/Bucket4jRateLimiter.kt
+++ b/wisp/wisp-rate-limiting/bucket4j/src/main/kotlin/wisp/ratelimiting/bucket4j/Bucket4jRateLimiter.kt
@@ -4,6 +4,7 @@ import io.github.bucket4j.Bandwidth
 import io.github.bucket4j.BucketConfiguration
 import io.github.bucket4j.ConsumptionProbe
 import io.github.bucket4j.EstimationProbe
+import io.github.bucket4j.TokensInheritanceStrategy
 import io.github.bucket4j.distributed.BucketProxy
 import io.github.bucket4j.distributed.proxy.ProxyManager
 import io.micrometer.core.instrument.MeterRegistry
@@ -106,7 +107,11 @@ class Bucket4jRateLimiter @JvmOverloads constructor(
       .addLimit(configuration.toBandwidth())
       .build()
 
-    return bucketProxy.builder().build(key, bucketConfig)
+    return bucketProxy.builder().apply {
+      configuration.version?.let {
+        this.withImplicitConfigurationReplacement(it, TokensInheritanceStrategy.PROPORTIONALLY)
+      }
+    }.build(key, bucketConfig)
   }
 
   private fun RateLimitConfiguration.toBandwidth(): Bandwidth {
diff --git a/wisp/wisp-rate-limiting/src/main/kotlin/wisp/ratelimiting/RateLimitConfiguration.kt b/wisp/wisp-rate-limiting/src/main/kotlin/wisp/ratelimiting/RateLimitConfiguration.kt
index e66167f7ca5..d4c0aba5d45 100644
--- a/wisp/wisp-rate-limiting/src/main/kotlin/wisp/ratelimiting/RateLimitConfiguration.kt
+++ b/wisp/wisp-rate-limiting/src/main/kotlin/wisp/ratelimiting/RateLimitConfiguration.kt
@@ -26,4 +26,12 @@ interface RateLimitConfiguration {
    * The period of time over which [refillAmount] tokens are added back to the bucket
    */
   val refillPeriod: Duration
+
+  /**
+   * The version of the configuration. This allows implicit configuration replacement.
+   * Make sure to increase the version when changing the configuration.
+   * Desired version should start from 1
+   */
+  val version: Long?
+    get() = null // returns null to be backward compatible
 }