Refactor SmoothWarmupLimiter and TokenBucketLimiter to improve wait time calculations and enhance logic for permit handling

Author: chenzhiguo

joylive-core/joylive-governance-api/src/main/java/com/jd/live/agent/governance/invoke/ratelimit/tokenbucket/SmoothWarmupLimiter.java

@@ -16,6 +16,7 @@
 import com.jd.live.agent.governance.policy.service.limit.RateLimitPolicy;
 import com.jd.live.agent.governance.policy.service.limit.SlidingWindow;
 
+import static java.lang.Math.min;
 
 /**
  * SmoothWarmupLimiter
@@ -82,9 +83,27 @@ protected double getMaxStoredPermits() {
         return thresholdPermits + 2.0 * warmupMicros / (permitIntervalMicros + coldIntervalMicros);
     }
 
+    @Override
+    protected long waitForStorePermits(double storedPermits, double targetPermits) {
+        double availablePermitsAboveThreshold = storedPermits - thresholdPermits;
+        long micros = 0;
+        // measuring the integral on the right part of the function (the climbing line)
+        if (availablePermitsAboveThreshold > 0.0) {
+            double permitsAboveThresholdToTake = min(availablePermitsAboveThreshold, targetPermits);
+            double length =
+                    permitsToTime(availablePermitsAboveThreshold)
+                            + permitsToTime(availablePermitsAboveThreshold - permitsAboveThresholdToTake);
+            micros = (long) (permitsAboveThresholdToTake * length / 2.0);
+            targetPermits -= permitsAboveThresholdToTake;
+        }
+        // measuring the integral on the left part of the function (the horizontal line)
+        micros += (long) (permitIntervalMicros * targetPermits);
+        return micros;
+    }
+
     @Override
     protected double coolDownIntervalMicros() {
-        return storedPermits > thresholdPermits ? coldIntervalMicros : permitIntervalMicros;
+        return warmupMicros / maxStoredPermits;
     }
 
     private double permitsToTime(double permits) {

joylive-core/joylive-governance-api/src/main/java/com/jd/live/agent/governance/invoke/ratelimit/tokenbucket/TokenBucketLimiter.java

@@ -154,16 +154,26 @@ protected boolean isOver() {
     protected long waitForRequiredPermits(long permits, long startTimeMicros, long timeoutMicros) {
         long nowMicros = stopwatch.readMicros();
         refresh(nowMicros);
-        double available = min(permits, storedPermits);
-        double lack = permits - available;
-        long waitTimeMicros = waitForStorePermits(storedPermits, available) + (long) (lack * permitIntervalMicros);
-        long momentAvailable = saturatedAdd(this.nextPermitMicros, waitTimeMicros);
 
-        long microsToSleep = max(0L, momentAvailable - nowMicros);
+        // 1. First, calculate the time needed to wait for the current request itself.
+        //    If the current time has already exceeded nextPermitMicros, there is no need to wait (wait = 0).
+        //    If the current time has not yet reached nextPermitMicros (which is a hole left by the previous request), then a wait is required.
+        long microsToSleep = max(0L, this.nextPermitMicros - nowMicros);
+
+        // 2. Check if there is a timeout
+        //    Only when the "debt of the ancestors" is too heavy, and the time I need to wait exceeds my Timeout, do I give up.
         if (microsToSleep > timeoutMicros) {
             return TIMEOUT;
         }
-        this.nextPermitMicros = momentAvailable;
+
+        // 3. Calculate the additional time cost of consuming the token for this request (this is the pit I'm leaving for posterity)
+        double available = min(permits, storedPermits);
+        double lack = permits - available;
+        long waitTimeMicros = waitForStorePermits(storedPermits, available) + (long) (lack * permitIntervalMicros);
+
+        // 4.Time to update next token available = current next time + cost incurred this time
+        this.nextPermitMicros = saturatedAdd(this.nextPermitMicros, waitTimeMicros);
+
         storedPermits -= available;
         return microsToSleep;
     }