Fix to correctly set TokenBucket capacity when requested amount greater than capacity (#2879)

Author: joviegas

.changes/next-release/bugfix-AWSSDKforJavav2-66cf80a.json

@@ -0,0 +1,6 @@
+{
+    "category": "AWS SDK for Java v2", 
+    "contributor": "", 
+    "type": "bugfix", 
+    "description": "Fix to set TokenBucket capacity correctly when requested amount greater than capacity"
+}

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/retry/RateLimitingTokenBucket.java

@@ -185,15 +185,14 @@ public OptionalDouble acquireNonBlocking(double amount, boolean fastFail) {
      * @return The unfulfilled amount.
      */
     double tryAcquireCapacity(double amount) {
+        double result;
         if (amount <= currentCapacity) {
-            currentCapacity = currentCapacity - amount;
-            amount = 0;
+            result = 0;
         } else {
-            amount = amount - currentCapacity;
-            currentCapacity = 0;
+            result = amount - currentCapacity;
         }
-
-        return amount;
+        currentCapacity = currentCapacity - amount;
+        return result;
     }
 
     private void initialize() {
@@ -347,9 +346,10 @@ synchronized void calculateTimeWindow() {
 
     /**
      * Sleep for a given amount of seconds.
+     *
      * @param seconds The amount of time to sleep in seconds.
      */
-    private static void sleep(double seconds) {
+    void sleep(double seconds) {
         long millisToSleep = (long) (seconds * 1000);
         try {
             Thread.sleep(millisToSleep);

core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/retry/RateLimitingTokenBucketTest.java

@@ -15,6 +15,7 @@
 
 package software.amazon.awssdk.core.internal.retry;
 
+import static org.assertj.core.api.Assertions.as;
 import static org.assertj.core.api.Assertions.assertThat;
 
 import java.time.Duration;
@@ -66,6 +67,8 @@ public void acquire_capacityInsufficient_sleepsForRequiredTime() {
         long elapsed = System.nanoTime() - a;
 
         assertThat(acquired).isTrue();
+        assertThat(tb.getCurrentCapacity()).isNegative();
+        assertThat(tb.getCurrentCapacity()).isEqualTo(-1.0);
         assertThat(Duration.ofNanos(elapsed).getSeconds()).isEqualTo(1);
     }
 
@@ -96,6 +99,53 @@ public void tryAcquireCapacity_capacitySufficient_returns0() {
         tb.setCurrentCapacity(5.0);
 
         assertThat(tb.tryAcquireCapacity(5.0)).isZero();
+        assertThat(tb.getCurrentCapacity()).isZero();
+    }
+
+    @Test
+    public void tryAcquireCapacity_amountGreaterThanCapacity_returnsNonZero() {
+        RateLimitingTokenBucket tb = new RateLimitingTokenBucket();
+        tb.setCurrentCapacity(5.0);
+
+        assertThat(tb.tryAcquireCapacity(8.0)).isEqualTo(3);
+        assertThat(tb.getCurrentCapacity()).isNegative();
+        assertThat(tb.getCurrentCapacity()).isEqualTo(-3);
+    }
+
+
+    @Test
+    public void acquire_amountGreaterThanNonZeroPositiveCapacity_setsNegativeCapacity() {
+        RateLimitingTokenBucket tb = Mockito.spy(new RateLimitingTokenBucket());
+
+        // stub out sleep , since we do not actually want to wait for sleep time
+        Mockito.doAnswer(invocationOnMock -> null).when(tb).sleep(2);
+
+        tb.setFillRate(1.0);
+        tb.setCurrentCapacity(1.0);
+        tb.enable();
+
+        boolean acquired = tb.acquire(3.0);
+        assertThat(acquired).isTrue();
+        assertThat(tb.getCurrentCapacity()).isNegative();
+        assertThat(tb.getCurrentCapacity()).isEqualTo(-2.0);
+    }
+
+    @Test
+    public void acquire_amountGreaterThanNegativeCapacity_setsNegativeCapacity() {
+        RateLimitingTokenBucket tb = Mockito.spy(new RateLimitingTokenBucket());
+
+        // stub out sleep , since we do not actually want to wait for sleep time
+        Mockito.doAnswer(invocationOnMock -> null).when(tb).sleep(3);
+
+        tb.setFillRate(1.0);
+        tb.setCurrentCapacity(-1.0);
+        tb.enable();
+
+        boolean acquired = tb.acquire(2.0);
+
+        assertThat(acquired).isTrue();
+        assertThat(tb.getCurrentCapacity()).isNegative();
+        assertThat(tb.getCurrentCapacity()).isEqualTo(-3.0);
     }
 
     @Test