Prune idle per-host connection semaphores

client/src/main/java/org/asynchttpclient/netty/channel/CombinedConnectionSemaphore.java

@@ -16,6 +16,7 @@
 package org.asynchttpclient.netty.channel;
 
 import java.io.IOException;
+import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -33,14 +34,21 @@ public class CombinedConnectionSemaphore extends PerHostConnectionSemaphore {
     public void acquireChannelLock(Object partitionKey) throws IOException {
         long remainingTime = acquireTimeout > 0 ? acquireGlobalTimed(partitionKey) : acquireGlobal(partitionKey);
 
+        Semaphore freeConnections = reserveFreeConnectionsForHost(partitionKey);
+        boolean acquired = false;
         try {
-            if (remainingTime < 0 || !getFreeConnectionsForHost(partitionKey).tryAcquire(remainingTime, TimeUnit.MILLISECONDS)) {
+            acquired = remainingTime >= 0 && freeConnections.tryAcquire(remainingTime, TimeUnit.MILLISECONDS);
+            if (!acquired) {
                 releaseGlobal(partitionKey);
                 throw tooManyConnectionsPerHost;
             }
         } catch (InterruptedException e) {
             releaseGlobal(partitionKey);
             throw new RuntimeException(e);
+        } finally {
+            if (!acquired) {
+                releaseHostReservation(partitionKey, freeConnections);
+            }
         }
     }
 
@@ -53,9 +61,18 @@ public void acquireChannelLock(Object partitionKey, boolean nonBlocking) throws
         // nonBlocking (the caller is on the event loop): take the global permit without waiting, then the
         // per-host permit without waiting, releasing the global one if the per-host permit is unavailable.
         globalMaxConnectionSemaphore.acquireChannelLock(partitionKey, true);
-        if (!getFreeConnectionsForHost(partitionKey).tryAcquire()) {
-            releaseGlobal(partitionKey);
-            throw tooManyConnectionsPerHost;
+        Semaphore freeConnections = reserveFreeConnectionsForHost(partitionKey);
+        boolean acquired = false;
+        try {
+            acquired = freeConnections.tryAcquire();
+            if (!acquired) {
+                releaseGlobal(partitionKey);
+                throw tooManyConnectionsPerHost;
+            }
+        } finally {
+            if (!acquired) {
+                releaseHostReservation(partitionKey, freeConnections);
+            }
         }
     }
 

client/src/main/java/org/asynchttpclient/netty/channel/PerHostConnectionSemaphore.java

@@ -29,6 +29,18 @@
  */
 public class PerHostConnectionSemaphore implements ConnectionSemaphore {
 
+    private static final class TrackedSemaphore extends Semaphore {
+
+        private static final long serialVersionUID = 1L;
+
+        // Guarded by freeChannelsPerHost.compute/computeIfPresent for this entry's key.
+        private int references;
+
+        private TrackedSemaphore(int permits) {
+            super(permits);
+        }
+    }
+
     protected final ConcurrentHashMap<Object, Semaphore> freeChannelsPerHost = new ConcurrentHashMap<>();
     protected final int maxConnectionsPerHost;
     protected final IOException tooManyConnectionsPerHost;
@@ -48,29 +60,68 @@ public void acquireChannelLock(Object partitionKey) throws IOException {
 
     @Override
     public void acquireChannelLock(Object partitionKey, boolean nonBlocking) throws IOException {
+        Semaphore freeConnections = reserveFreeConnectionsForHost(partitionKey);
+        boolean acquired = false;
         try {
-            Semaphore freeConnections = getFreeConnectionsForHost(partitionKey);
             // nonBlocking (the caller is on the event loop): try once and fail fast rather than parking
             // the loop for up to acquireTimeout.
-            boolean acquired = nonBlocking
+            acquired = nonBlocking
                     ? freeConnections.tryAcquire()
                     : freeConnections.tryAcquire(acquireTimeout, TimeUnit.MILLISECONDS);
             if (!acquired) {
                 throw tooManyConnectionsPerHost;
             }
         } catch (InterruptedException e) {
             throw new RuntimeException(e);
+        } finally {
+            if (!acquired) {
+                releaseHostReservation(partitionKey, freeConnections);
+            }
         }
     }
 
     @Override
     public void releaseChannelLock(Object partitionKey) {
-        getFreeConnectionsForHost(partitionKey).release();
+        if (maxConnectionsPerHost <= 0) {
+            return;
+        }
+        Semaphore freeConnections = freeChannelsPerHost.get(partitionKey);
+        if (freeConnections != null) {
+            freeConnections.release();
+            releaseHostReservation(partitionKey, freeConnections);
+        }
     }
 
     protected Semaphore getFreeConnectionsForHost(Object partitionKey) {
         return maxConnectionsPerHost > 0 ?
-                freeChannelsPerHost.computeIfAbsent(partitionKey, pk -> new Semaphore(maxConnectionsPerHost)) :
+                freeChannelsPerHost.computeIfAbsent(partitionKey, pk -> new TrackedSemaphore(maxConnectionsPerHost)) :
                 InfiniteSemaphore.INSTANCE;
     }
+
+    final Semaphore reserveFreeConnectionsForHost(Object partitionKey) {
+        if (maxConnectionsPerHost <= 0) {
+            return InfiniteSemaphore.INSTANCE;
+        }
+        return freeChannelsPerHost.compute(partitionKey, (pk, current) -> {
+            TrackedSemaphore semaphore = current == null
+                    ? new TrackedSemaphore(maxConnectionsPerHost)
+                    : (TrackedSemaphore) current;
+            semaphore.references++;
+            return semaphore;
+        });
+    }
+
+    final void releaseHostReservation(Object partitionKey, Semaphore expected) {
+        if (maxConnectionsPerHost <= 0) {
+            return;
+        }
+        freeChannelsPerHost.computeIfPresent(partitionKey, (pk, current) -> {
+            if (current != expected) {
+                return current;
+            }
+            TrackedSemaphore semaphore = (TrackedSemaphore) current;
+            semaphore.references--;
+            return semaphore.references == 0 ? null : semaphore;
+        });
+    }
 }

client/src/test/java/org/asynchttpclient/netty/channel/SemaphoreTest.java

@@ -28,12 +28,18 @@
 import java.io.IOException;
 import java.util.List;
 import java.util.Objects;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 @TestInstance(TestInstance.Lifecycle.PER_CLASS)
@@ -173,6 +179,55 @@ private void checkRelease(ConnectionSemaphore semaphore) throws IOException {
         assertFalse(tooManyCaught);
     }
 
+    @Test
+    public void perHostRemovesEntriesAfterLastRelease() throws IOException {
+        removesEntriesAfterLastRelease(new PerHostConnectionSemaphore(1, 0));
+    }
+
+    @Test
+    public void combinedRemovesEntriesAfterLastRelease() throws IOException {
+        removesEntriesAfterLastRelease(new CombinedConnectionSemaphore(1, 1, 0));
+    }
+
+    private static void removesEntriesAfterLastRelease(PerHostConnectionSemaphore semaphore) throws IOException {
+        for (int i = 0; i < 100; i++) {
+            String partitionKey = "host-" + i;
+            semaphore.acquireChannelLock(partitionKey);
+            semaphore.releaseChannelLock(partitionKey);
+        }
+        assertTrue(semaphore.freeChannelsPerHost.isEmpty());
+    }
+
+    @Test
+    @Timeout(unit = TimeUnit.SECONDS, value = 5)
+    public void perHostRetainsEntryDuringWaitingAcquire() throws Exception {
+        PerHostConnectionSemaphore semaphore = new PerHostConnectionSemaphore(1, 30_000);
+        semaphore.acquireChannelLock(PK);
+        Semaphore hostSemaphore = semaphore.freeChannelsPerHost.get(PK);
+        ExecutorService executor = Executors.newSingleThreadExecutor();
+        try {
+            Future<Void> waitingAcquire = executor.submit(() -> {
+                semaphore.acquireChannelLock(PK);
+                return null;
+            });
+            while (!hostSemaphore.hasQueuedThreads()) {
+                assertFalse(waitingAcquire.isDone(), "waiting acquire completed before the permit was released");
+                Thread.yield();
+            }
+
+            semaphore.releaseChannelLock(PK);
+            waitingAcquire.get(1, TimeUnit.SECONDS);
+
+            assertSame(hostSemaphore, semaphore.freeChannelsPerHost.get(PK));
+            assertThrows(TooManyConnectionsPerHostException.class,
+                    () -> semaphore.acquireChannelLock(PK, true));
+            semaphore.releaseChannelLock(PK);
+            assertTrue(semaphore.freeChannelsPerHost.isEmpty());
+        } finally {
+            executor.shutdownNow();
+        }
+    }
+
     // ---- non-blocking acquire (event-loop path): must fail fast, never wait for acquireTimeout ----
 
     // Deliberately large: a blocking acquire on an exhausted semaphore would blow the 1s @Timeout below,