diff --git a/client/src/main/java/org/asynchttpclient/netty/channel/CombinedConnectionSemaphore.java b/client/src/main/java/org/asynchttpclient/netty/channel/CombinedConnectionSemaphore.java index 5d626c59c..6cb62c967 100644 --- a/client/src/main/java/org/asynchttpclient/netty/channel/CombinedConnectionSemaphore.java +++ b/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); + } } } diff --git a/client/src/main/java/org/asynchttpclient/netty/channel/PerHostConnectionSemaphore.java b/client/src/main/java/org/asynchttpclient/netty/channel/PerHostConnectionSemaphore.java index 73468a9ee..041b45ac4 100644 --- a/client/src/main/java/org/asynchttpclient/netty/channel/PerHostConnectionSemaphore.java +++ b/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 freeChannelsPerHost = new ConcurrentHashMap<>(); protected final int maxConnectionsPerHost; protected final IOException tooManyConnectionsPerHost; @@ -48,11 +60,12 @@ 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) { @@ -60,17 +73,55 @@ public void acquireChannelLock(Object partitionKey, boolean nonBlocking) throws } } 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; + }); + } } diff --git a/client/src/test/java/org/asynchttpclient/netty/channel/SemaphoreTest.java b/client/src/test/java/org/asynchttpclient/netty/channel/SemaphoreTest.java index 32faeb43c..2cd85a01b 100644 --- a/client/src/test/java/org/asynchttpclient/netty/channel/SemaphoreTest.java +++ b/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 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,