OF-3171: Fix token bucket refill precision and clarify refill logic

Preserve fractional refill time in TokenBucketRateLimiter by carrying sub-token remainder across refill cycles, instead of discarding it when whole tokens are added.

Author: guusdk

xmppserver/src/main/java/org/jivesoftware/util/TokenBucketRateLimiter.java

@@ -15,7 +15,11 @@
  */
 package org.jivesoftware.util;
 
+import com.google.common.annotations.VisibleForTesting;
+
 import java.time.Duration;
+import java.util.Objects;
+import java.util.function.LongSupplier;
 
 /**
  * A thread-safe, synchronized token-bucket rate limiter with metrics.
@@ -29,9 +33,12 @@ public final class TokenBucketRateLimiter
 {
     private final long capacity;
     private final long refillTokensPerSecond;
+    private final LongSupplier nanoTimeSupplier;
 
     private long availableTokens;
     private long lastRefillTimeNanos;
+    // Leftover refill value in scaled units (1 token = 1_000_000_000 units).
+    private long refillRemainder;
 
     private long acceptedEvents;
     private long rejectedEvents;
@@ -46,14 +53,16 @@ public final class TokenBucketRateLimiter
     /**
      * Creates an unlimited rate limiter. Use {@link #unlimited()} to obtain an instance.
      */
-    private TokenBucketRateLimiter()
+    private TokenBucketRateLimiter(final LongSupplier nanoTimeSupplier)
     {
+        this.nanoTimeSupplier = Objects.requireNonNull(nanoTimeSupplier, "nanoTimeSupplier must not be null");
         this.unlimited = true;
         this.refillTokensPerSecond = 0;
         this.capacity = 0;
         this.availableTokens = 0;
         this.lastRefillTimeNanos = 0;
-        this.startTimeNanos = System.nanoTime();
+        this.refillRemainder = 0;
+        this.startTimeNanos = nanoTime();
     }
 
     /**
@@ -63,6 +72,21 @@ private TokenBucketRateLimiter()
      * @param maxBurst maximum number of permits that can accumulate
      */
     public TokenBucketRateLimiter(final long permitsPerSecond, final long maxBurst)
+    {
+        this(permitsPerSecond, maxBurst, System::nanoTime);
+    }
+
+    /**
+     * Creates a new rate limiter with a custom clock.
+     *
+     * Normally, the system clock is used. This constructor mainly exists for tests.
+     *
+     * @param permitsPerSecond sustained rate of permits
+     * @param maxBurst         maximum number of permits that can accumulate
+     * @param nanoTimeSupplier custom clock
+     */
+    @VisibleForTesting
+    TokenBucketRateLimiter(final long permitsPerSecond, final long maxBurst, final LongSupplier nanoTimeSupplier)
     {
         if (permitsPerSecond <= 0) {
             throw new IllegalArgumentException("permitsPerSecond must be > 0");
@@ -71,11 +95,13 @@ public TokenBucketRateLimiter(final long permitsPerSecond, final long maxBurst)
             throw new IllegalArgumentException("maxBurst must be > 0");
         }
 
+        this.nanoTimeSupplier = Objects.requireNonNull(nanoTimeSupplier, "nanoTimeSupplier must not be null");
         this.unlimited = false;
         this.refillTokensPerSecond = permitsPerSecond;
         this.capacity = maxBurst;
         this.availableTokens = maxBurst;
-        this.lastRefillTimeNanos = System.nanoTime();
+        this.lastRefillTimeNanos = nanoTime();
+        this.refillRemainder = 0;
         this.startTimeNanos = this.lastRefillTimeNanos;
     }
 
@@ -86,7 +112,19 @@ public TokenBucketRateLimiter(final long permitsPerSecond, final long maxBurst)
      */
     public static TokenBucketRateLimiter unlimited()
     {
-        return new TokenBucketRateLimiter();
+        return new TokenBucketRateLimiter(System::nanoTime);
+    }
+
+    /**
+     * Returns the current time in nanoseconds from this instance clock.
+     *
+     * Usually this is the system clock, but tests can provide a custom clock.
+     *
+     * @return The time in nanoseconds.
+     */
+    private long nanoTime()
+    {
+        return nanoTimeSupplier.getAsLong();
     }
 
     /**
@@ -112,30 +150,54 @@ public synchronized boolean tryAcquire()
     }
 
     /**
-     * Refills tokens based on elapsed time.
+     * Adds tokens based on elapsed time.
      */
     private void refillIfNeeded()
     {
-        final long now = System.nanoTime();
+        final long now = nanoTime();
         final long elapsed = now - lastRefillTimeNanos;
+        // With very small intervals, no time may have passed yet.
         if (elapsed <= 0) {
             return;
         }
 
-        // If the multiplication would overflow, elapsed time is so large that the bucket would be completely refilled
-        // regardless, so cap directly at capacity.
-        final long tokensToAdd;
-        if (elapsed > Long.MAX_VALUE / refillTokensPerSecond) {
-            tokensToAdd = capacity;
-        } else {
-            tokensToAdd = Math.min(capacity, (elapsed * refillTokensPerSecond) / 1_000_000_000L);
+        if (availableTokens >= capacity) {
+            // When already full, do not store extra time as hidden credit.
+            lastRefillTimeNanos = now;
+            refillRemainder = 0;
+            return;
         }
 
-        if (tokensToAdd > 0) {
-            availableTokens = tokensToAdd >= capacity - availableTokens ? capacity : availableTokens + tokensToAdd;
-            // Only advance the timestamp when tokens are actually added, so that sub-token elapsed time is preserved
-            // and contributes to the next refill rather than being discarded.
+        final long remainingCapacity = capacity - availableTokens;
+
+        // Refill is calculated in scaled integer units: (elapsed * rate) + previous leftover.
+        // If this overflows, elapsed time is so large that the bucket must be full.
+        if (elapsed > (Long.MAX_VALUE - refillRemainder) / refillTokensPerSecond) {
+            availableTokens = capacity;
+            lastRefillTimeNanos = now;
+            refillRemainder = 0;
+            return;
+        }
+
+        // Convert elapsed time to scaled units (1_000_000_000 units = 1 token).
+        final long refillUnits = elapsed * refillTokensPerSecond + refillRemainder;
+        final long tokensToGenerate = refillUnits / 1_000_000_000L;
+        // Keep leftover units until they become at least one full token.
+        if (tokensToGenerate <= 0) {
+            return;
+        }
+
+        final long tokensToAdd = Math.min(remainingCapacity, tokensToGenerate);
+        availableTokens += tokensToAdd;
+
+        if (availableTokens >= capacity) {
+            // Once capacity is reached, drop any extra accrued value.
+            lastRefillTimeNanos = now;
+            refillRemainder = 0;
+        } else {
+            // Keep the leftover fraction so refill speed stays accurate over time.
             lastRefillTimeNanos = now;
+            refillRemainder = refillUnits % 1_000_000_000L;
         }
     }
 
@@ -220,7 +282,7 @@ public synchronized double getAcceptanceRatio() {
      * @return uptime duration
      */
     public Duration getUptime() {
-        return Duration.ofNanos(System.nanoTime() - startTimeNanos);
+        return Duration.ofNanos(nanoTime() - startTimeNanos);
     }
 
     /**

xmppserver/src/test/java/org/jivesoftware/util/TokenBucketRateLimiterTest.java

@@ -18,6 +18,8 @@
 import org.junit.jupiter.api.Test;
 
 import java.time.Duration;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.LongSupplier;
 
 import static org.junit.jupiter.api.Assertions.*;
 
@@ -154,6 +156,48 @@ void testGetAvailableTokensReflectsRefill() throws InterruptedException
         assertEquals(1, limiter.getAvailableTokens(), "getAvailableTokens should reflect tokens added by refill");
     }
 
+    /**
+     * Verifies that sub-token elapsed time is preserved and contributes to a later refill.
+     */
+    @Test
+    void testFractionalElapsedTimeIsPreservedAcrossRefills()
+    {
+        // Setup test fixture.
+        final FakeNanoClock clock = new FakeNanoClock();
+        final TokenBucketRateLimiter limiter = new TokenBucketRateLimiter(1, 2, clock);
+
+        // Execute system under test.
+        assertTrue(limiter.tryAcquire(), "First acquire should consume one of the initial burst tokens");
+        assertTrue(limiter.tryAcquire(), "Second acquire should consume the second initial burst token");
+        assertFalse(limiter.tryAcquire(), "Bucket should now be exhausted");
+
+        clock.advanceNanos(1_500_000_000L);
+        assertTrue(limiter.tryAcquire(), "1.5 seconds should refill exactly one token");
+        assertFalse(limiter.tryAcquire(), "Only one token should have been refilled so far");
+
+        clock.advanceNanos(500_000_000L);
+        assertTrue(limiter.tryAcquire(), "Remaining 0.5 second should combine with prior remainder to refill one token");
+    }
+
+    /**
+     * Verifies that elapsed time while the bucket is full does not become hidden credit.
+     */
+    @Test
+    void testElapsedTimeIsNotBankedWhileBucketIsFull()
+    {
+        // Setup test fixture.
+        final FakeNanoClock clock = new FakeNanoClock();
+        final TokenBucketRateLimiter limiter = new TokenBucketRateLimiter(1, 2, clock);
+
+        // Execute system under test.
+        clock.advanceNanos(5_000_000_000L);
+        assertEquals(2, limiter.getAvailableTokens(), "Bucket should remain capped at burst capacity");
+
+        assertTrue(limiter.tryAcquire(), "One token should be immediately available");
+        assertTrue(limiter.tryAcquire(), "Second token should be immediately available");
+        assertFalse(limiter.tryAcquire(), "No hidden refill credit should exist after draining a full bucket");
+    }
+
     /**
      * Verifies that the limiter correctly counts accepted and rejected events.
      */
@@ -360,4 +404,23 @@ void testUnlimitedLimiterInstancesAreIndependent()
         assertEquals(3, limiter1.getAcceptedEvents(), "Limiter1 should count only its own accepted events");
         assertEquals(0, limiter2.getAcceptedEvents(), "Limiter2 should not be affected by events on limiter1");
     }
+
+    /**
+     * A fake {@link System#nanoTime()} implementation that allows for manual increments of time.
+     */
+    private static final class FakeNanoClock implements LongSupplier
+    {
+        private final AtomicLong nowNanos = new AtomicLong();
+
+        @Override
+        public long getAsLong()
+        {
+            return nowNanos.get();
+        }
+
+        void advanceNanos(final long nanos)
+        {
+            nowNanos.addAndGet(nanos);
+        }
+    }
 }