From 1ff74d28da37adcb6969b00b69c977a7506e5b2b Mon Sep 17 00:00:00 2001
From: Walter Oggioni <oggioni.walter@gmail.com>
Date: Sat, 25 Jan 2025 01:18:05 +0800
Subject: [PATCH] fixed Bucket wait time estimation

---
 gradle.properties                             |  4 +--
 .../woggioni/jwo/internal/LocalBucket.java    | 35 ++++++++++---------
 2 files changed, 21 insertions(+), 18 deletions(-)

diff --git a/gradle.properties b/gradle.properties
index dc48623..b60471b 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -3,6 +3,6 @@ org.gradle.parallel=true
 org.gradle.caching=true
 
 gitea.maven.url = https://gitea.woggioni.net/api/packages/woggioni/maven
-jwo.version = 2025.01.23
-lys.version = 2025.01.17
+jwo.version = 2025.01.26
+lys.version = 2025.01.24
 guice.version = 5.0.1
diff --git a/src/main/java/net/woggioni/jwo/internal/LocalBucket.java b/src/main/java/net/woggioni/jwo/internal/LocalBucket.java
index 865e76c..7591590 100644
--- a/src/main/java/net/woggioni/jwo/internal/LocalBucket.java
+++ b/src/main/java/net/woggioni/jwo/internal/LocalBucket.java
@@ -35,12 +35,8 @@ public class LocalBucket implements Bucket {
         this.lastFill = new AtomicLong(currentTimestamp);
     }
 
-    long getTokenPrivate(long nTokens, long now) {
-        if(nTokens > maxCapacity) throw new IllegalArgumentException("The requested number of tokens exceeds the bucket max capacity");
+    private long getTokenPrivate(long nTokens, long now, long previousFillTime, long currentFillTime) {
         final LongBinaryOperator tickCalculator = (lf, currentTimestamp) -> (currentTimestamp - lf) / fillPeriod;
-        final LongBinaryOperator timestampCalculator = (lf, currentTimestamp) -> lf + tickCalculator.applyAsLong(lf, currentTimestamp) * fillPeriod;
-        final long previousFillTime = lastFill.getAndAccumulate(now, timestampCalculator);
-        final long currentFillTime = timestampCalculator.applyAsLong(previousFillTime, now);
         if (currentFillTime != previousFillTime) {
             final long ticks = tickCalculator.applyAsLong(previousFillTime, now);
             final LongUnaryOperator filledAmountCalculator = currentTokens -> Math.min(currentTokens + ticks * fillAmount, maxCapacity);
@@ -67,33 +63,40 @@ public class LocalBucket implements Bucket {
         }
     }
 
+
     @Override
     public boolean removeTokens(long nTokens) {
         return removeTokens(nTokens, System.nanoTime());
     }
 
     @Override
-    public boolean removeTokens(long nTokens, long currentTimestamp) {
-        return getTokenPrivate(nTokens, currentTimestamp) >= nTokens;
+    public boolean removeTokens(long nTokens, long now) {
+        if(nTokens > maxCapacity) throw new IllegalArgumentException("The requested number of tokens exceeds the bucket max capacity");
+        final LongBinaryOperator tickCalculator = (lf, currentTimestamp) -> (currentTimestamp - lf) / fillPeriod;
+        final LongBinaryOperator timestampCalculator = (lf, currentTimestamp) -> lf + tickCalculator.applyAsLong(lf, currentTimestamp) * fillPeriod;
+        final long previousFillTime = lastFill.getAndAccumulate(now, timestampCalculator);
+        final long currentFillTime = timestampCalculator.applyAsLong(previousFillTime, now);
+        long result = getTokenPrivate(nTokens, now, previousFillTime, currentFillTime);
+        return result >= nTokens;
     }
 
     @Override
     public long removeTokensWithEstimate(long nTokens) {
-        final long previousTokenAmount = getTokenPrivate(nTokens, System.nanoTime());
-        if(previousTokenAmount >= nTokens) {
-            return -1;
-        } else {
-            return ceilDiv((nTokens - previousTokenAmount) * fillPeriod, fillAmount);
-        }
+        return removeTokensWithEstimate(nTokens, System.nanoTime());
     }
 
     @Override
-    public long removeTokensWithEstimate(long nTokens, long currentTimestamp) {
-        final long previousTokenAmount = getTokenPrivate(nTokens, currentTimestamp);
+    public long removeTokensWithEstimate(long nTokens, long now) {
+        if(nTokens > maxCapacity) throw new IllegalArgumentException("The requested number of tokens exceeds the bucket max capacity");
+        final LongBinaryOperator tickCalculator = (lf, currentTimestamp) -> (currentTimestamp - lf) / fillPeriod;
+        final LongBinaryOperator timestampCalculator = (lf, currentTimestamp) -> lf + tickCalculator.applyAsLong(lf, currentTimestamp) * fillPeriod;
+        final long previousFillTime = lastFill.getAndAccumulate(now, timestampCalculator);
+        final long currentFillTime = timestampCalculator.applyAsLong(previousFillTime, now);
+        long previousTokenAmount = getTokenPrivate(nTokens, now, previousFillTime, currentFillTime);
         if(previousTokenAmount >= nTokens) {
             return -1;
         } else {
-            return ceilDiv((nTokens - previousTokenAmount) * fillPeriod, fillAmount);
+            return Math.max(ceilDiv((nTokens - previousTokenAmount) * fillPeriod, fillAmount), currentFillTime + fillPeriod - now);
         }
     }
 }