AIMD limit now correctly releases permits when token is used. (#9421)

Added tests for AIMD and Fixed limit to validate this. Signed-off-by: Tomas Langer <[email protected]>
helidon-io · Oct 22, 2024 · c3c4eb3 · c3c4eb3
1 parent bd5bab5
commit c3c4eb3
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 3 deletions.
diff --git a/.../concurrency/limits/src/main/java/io/helidon/common/concurrency/limits/AimdLimitImpl.java b/.../concurrency/limits/src/main/java/io/helidon/common/concurrency/limits/AimdLimitImpl.java
@@ -176,18 +176,27 @@ private AimdToken(Supplier<Long> clock, AtomicInteger concurrentRequests) {
 
         @Override
         public void dropped() {
-            updateWithSample(startTime, clock.get(), currentRequests, false);
+            try {
+                updateWithSample(startTime, clock.get(), currentRequests, false);
+            } finally {
+                AimdLimitImpl.this.semaphore.release();
+            }
         }
 
         @Override
         public void ignore() {
             concurrentRequests.decrementAndGet();
+            AimdLimitImpl.this.semaphore.release();
         }
 
         @Override
         public void success() {
-            updateWithSample(startTime, clock.get(), currentRequests, true);
-            concurrentRequests.decrementAndGet();
+            try {
+                updateWithSample(startTime, clock.get(), currentRequests, true);
+                concurrentRequests.decrementAndGet();
+            } finally {
+                AimdLimitImpl.this.semaphore.release();
+            }
         }
     }
 }
diff --git a/.../concurrency/limits/src/test/java/io/helidon/common/concurrency/limits/AimdLimitTest.java b/.../concurrency/limits/src/test/java/io/helidon/common/concurrency/limits/AimdLimitTest.java
@@ -17,6 +17,7 @@
 package io.helidon.common.concurrency.limits;
 
 import java.time.Duration;
+import java.util.Optional;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -26,6 +27,7 @@
 import org.junit.jupiter.api.Test;
 
 import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.CoreMatchers.not;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
@@ -165,4 +167,18 @@ public void testSemaphoreReleased() throws Exception {
             limit.invoke(() -> {});
         }
     }
+
+    @Test
+    public void testSemaphoreReleasedWithToken() {
+        Limit limit = AimdLimit.builder()
+                .minLimit(5)
+                .initialLimit(5)
+                .build();
+
+        for (int i = 0; i < 5000; i++) {
+            Optional<LimitAlgorithm.Token> token = limit.tryAcquire();
+            assertThat(token, not(Optional.empty()));
+            token.get().success();
+        }
+    }
 }
diff --git a/...concurrency/limits/src/test/java/io/helidon/common/concurrency/limits/FixedLimitTest.java b/...concurrency/limits/src/test/java/io/helidon/common/concurrency/limits/FixedLimitTest.java
@@ -19,6 +19,7 @@
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -28,6 +29,7 @@
 import org.junit.jupiter.api.Test;
 
 import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.CoreMatchers.not;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.hasSize;
@@ -208,4 +210,19 @@ public void testSemaphoreReleasedWithQueue() throws Exception {
             });
         }
     }
+
+    @Test
+    public void testSemaphoreReleasedWithToken() {
+        Limit limit = FixedLimit.builder()
+                .permits(5)
+                .queueLength(10)
+                .queueTimeout(Duration.ofMillis(100))
+                .build();
+
+        for (int i = 0; i < 5000; i++) {
+            Optional<LimitAlgorithm.Token> token = limit.tryAcquire();
+            assertThat(token, not(Optional.empty()));
+            token.get().success();
+        }
+    }
 }