Fix idle pool cleaner generation race

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

@@ -37,7 +37,7 @@
 import java.util.concurrent.ConcurrentLinkedDeque;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
+import java.util.concurrent.atomic.AtomicLongFieldUpdater;
 import java.util.function.Predicate;
 
 import static org.asynchttpclient.util.DateUtils.unpreciseMillisTime;
@@ -52,8 +52,8 @@ public final class DefaultChannelPool implements ChannelPool {
     private static final AttributeKey<IdleState> IDLE_STATE_ATTRIBUTE_KEY = AttributeKey.valueOf("channelIdleState");
 
     // The partition deques hold the bare Channel; per-checkout idle state (start timestamp + the
-    // owned/tombstone CAS flag) lives on the channel's IDLE_STATE_ATTRIBUTE_KEY attribute, which is
-    // allocated once per physical connection and reused across every pool cycle (no per-offer holder).
+    // generation/ownership CAS state) lives on the channel's IDLE_STATE_ATTRIBUTE_KEY attribute,
+    // which is allocated once per physical connection and reused across every pool cycle.
     private final ConcurrentHashMap<Object, ConcurrentLinkedDeque<Channel>> partitions = new ConcurrentHashMap<>();
     private final AtomicBoolean isClosed = new AtomicBoolean(false);
     private final Timer nettyTimer;
@@ -132,16 +132,18 @@ private boolean offer0(Channel channel, Object partitionKey, long now) {
         if (partition == null) {
             partition = partitions.computeIfAbsent(partitionKey, pk -> new ConcurrentLinkedDeque<>());
         }
-        // Reuse the channel's IdleState instead of allocating a holder per offer; reset() stamps the
-        // idle start and clears the owned flag (must happen-before offerFirst publishes the channel,
-        // so any thread that observes it in the deque also observes owned == 0).
+        // Reuse the channel's IdleState instead of allocating a holder per offer; reset() transfers an
+        // owned generation to a new leasable generation before offerFirst publishes the channel.
         Attribute<IdleState> idleStateAttribute = channel.attr(IDLE_STATE_ATTRIBUTE_KEY);
         IdleState idleState = idleStateAttribute.get();
         if (idleState == null) {
-            idleState = new IdleState();
-            idleStateAttribute.set(idleState);
+            IdleState newIdleState = new IdleState();
+            IdleState existingIdleState = idleStateAttribute.setIfAbsent(newIdleState);
+            idleState = existingIdleState == null ? newIdleState : existingIdleState;
+        }
+        if (!idleState.reset(now)) {
+            return false;
         }
-        idleState.reset(now);
         return partition.offerFirst(channel);
     }
 
@@ -293,52 +295,68 @@ private static final class ChannelCreation {
      * {@link #IDLE_STATE_ATTRIBUTE_KEY} attribute, then reused across every pool checkout so no holder
      * is allocated per offer.
      *
-     * <p>{@code owned} is a single CAS flag with two roles, both meaning "this idle entry is claimed,
-     * do not lease it": a successful {@code poll()} lease, or a {@code removeAll()} tombstone. The pool
-     * upholds the invariant that a channel sitting in a partition deque has {@code owned == 0} unless it
-     * was tombstoned, because {@link #reset(long)} clears the flag before {@code offerFirst} publishes
-     * the channel and {@code poll()} unlinks a channel from the deque before claiming it. {@code start}
-     * doubles as a generation token: it changes on every offer, letting the cleaner detect a channel
-     * that was leased and re-offered between its expiry check and its claim.
+     * <p>The low bit of {@code state} is the ownership flag; the remaining bits advance by two on every
+     * offer because bit zero is reserved. A successful {@code poll()} lease and a {@code removeAll()}
+     * tombstone both claim the current generation. The cleaner can therefore claim only the exact
+     * generation whose expiry it evaluated, without temporarily claiming a newer entry while checking
+     * whether {@code start} changed.
+     *
+     * <p>A generation can be reset only while owned. New state starts owned, and {@code poll()} claims
+     * a generation before returning it, so {@code offer()} transfers ownership instead of performing a
+     * non-atomic update from a leasable generation.
      */
     static final class IdleState {
 
-        private static final AtomicIntegerFieldUpdater<IdleState> OWNED_UPDATER =
-                AtomicIntegerFieldUpdater.newUpdater(IdleState.class, "owned");
+        private static final long OWNED_MASK = 1L;
+        private static final AtomicLongFieldUpdater<IdleState> STATE_UPDATER =
+                AtomicLongFieldUpdater.newUpdater(IdleState.class, "state");
 
         private volatile long start;
-        @SuppressWarnings("unused")
-        private volatile int owned;
+        private volatile long state = OWNED_MASK;
 
         long start() {
             return start;
         }
 
+        long snapshot() {
+            return state;
+        }
+
         boolean isOwned() {
-            return owned != 0;
+            return isOwned(state);
         }
 
-        /** Atomically claim this entry; returns true only for the caller that transitions 0 -> 1. */
+        static boolean isOwned(long stateSnapshot) {
+            return (stateSnapshot & OWNED_MASK) != 0;
+        }
+
+        /** Atomically claim the current generation. */
         boolean takeOwnership() {
-            return OWNED_UPDATER.getAndSet(this, 1) == 0;
+            return tryTakeOwnership(state);
         }
 
-        /** Undo a claim taken via {@link #takeOwnership()} (used only on the cleaner re-offer race). */
-        void releaseOwnership() {
-            owned = 0;
+        /** Atomically claim {@code stateSnapshot}, provided it is still the current generation. */
+        boolean tryTakeOwnership(long stateSnapshot) {
+            return !isOwned(stateSnapshot)
+                    && STATE_UPDATER.compareAndSet(this, stateSnapshot, stateSnapshot | OWNED_MASK);
         }
 
-        /** Stamp the idle start and mark the channel leasable again. Called on every offer. */
-        void reset(long now) {
+        /** Transfer an owned generation to a new leasable generation. Called on every offer. */
+        synchronized boolean reset(long now) {
+            long stateSnapshot = state;
+            if (!isOwned(stateSnapshot)) {
+                return false;
+            }
+            long nextGeneration = (stateSnapshot & ~OWNED_MASK) + 2;
             start = now;
-            owned = 0;
+            return STATE_UPDATER.compareAndSet(this, stateSnapshot, nextGeneration);
         }
     }
 
     private final class IdleChannelDetector implements TimerTask {
 
-        private boolean isIdleTimeoutExpired(IdleState idleState, long now) {
-            return maxIdleTimeEnabled && now - idleState.start() >= maxIdleTime;
+        private boolean isIdleTimeoutExpired(long idleStart, long now) {
+            return maxIdleTimeEnabled && now - idleStart >= maxIdleTime;
         }
 
         @Override
@@ -410,7 +428,8 @@ private int reapPartition(ConcurrentLinkedDeque<Channel> partition, long now) {
                     continue;
                 }
 
-                if (idleState.isOwned()) {
+                long stateSnapshot = idleState.snapshot();
+                if (IdleState.isOwned(stateSnapshot)) {
                     // In-deque + owned ==> a removeAll(Channel) tombstone, or a node a concurrent poll()
                     // has already leased and unlinked. Either way: unlink, never close — the owner of the
                     // claim is responsible for closing it. Unlinking an already-unlinked node through the
@@ -419,22 +438,17 @@ private int reapPartition(ConcurrentLinkedDeque<Channel> partition, long now) {
                     continue;
                 }
 
-                boolean isIdleTimeoutExpired = isIdleTimeoutExpired(idleState, now);
+                long idleStart = idleState.start();
+                boolean isIdleTimeoutExpired = isIdleTimeoutExpired(idleStart, now);
                 boolean isRemotelyClosed = !Channels.isChannelActive(channel);
                 boolean isTtlExpired = isTtlExpired(channel, now);
                 if (!isIdleTimeoutExpired && !isRemotelyClosed && !isTtlExpired) {
                     continue; // healthy idle channel, leave it for poll()
                 }
 
-                long startSnapshot = idleState.start();
-                // Claim before closing so we never close a channel poll() is leasing concurrently.
-                if (!idleState.takeOwnership()) {
-                    continue; // poll() (or removeAll(Channel)) won the claim; that owner now handles the channel
-                }
-                if (idleState.start() != startSnapshot) {
-                    // The channel was leased and re-offered (fresh start) between the expiry check and
-                    // the claim, so it is leasable again — release it instead of closing it.
-                    idleState.releaseOwnership();
+                // Claim the generation evaluated above. A lease and re-offer changes the generation,
+                // so this CAS cannot claim the fresh entry or interfere with a concurrent poll of it.
+                if (!idleState.tryTakeOwnership(stateSnapshot)) {
                     continue;
                 }
 

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

@@ -31,6 +31,7 @@
 import java.util.Collections;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedDeque;
 import java.util.concurrent.ConcurrentLinkedQueue;
@@ -86,6 +87,22 @@ public void offerThenPollReturnsSameChannelThenEmpty() {
         pool.destroy();
     }
 
+    @Test
+    public void duplicateOfferDoesNotResetLeasableGeneration() throws Exception {
+        DefaultChannelPool pool = noReaperPool();
+        Channel channel = new EmbeddedChannel();
+
+        assertTrue(pool.offer(channel, KEY));
+        long generation = idleState(channel).snapshot();
+
+        assertFalse(pool.offer(channel, KEY), "an already leasable channel must not be offered again");
+        assertEquals(generation, idleState(channel).snapshot());
+        assertEquals(1, partitionSize(pool, KEY));
+        assertSame(channel, pool.poll(KEY));
+
+        pool.destroy();
+    }
+
     @Test
     public void reofferingReusesTheSameIdleStateInstance() throws Exception {
         DefaultChannelPool pool = noReaperPool();
@@ -245,6 +262,46 @@ public void channelReofferedAfterExpiryIsNotReaped() throws Exception {
         pool.destroy();
     }
 
+    @Test
+    public void cleanerDoesNotCloseChannelReofferedAfterExpiryCheck() throws Exception {
+        CapturingTimer timer = new CapturingTimer();
+        DefaultChannelPool pool = idlePool(timer, Duration.ofMillis(1));
+        BlockingActiveChannel channel = new BlockingActiveChannel();
+
+        pool.offer(channel, KEY);
+        idleState(channel).reset(0);
+        channel.blockNextActiveCheck();
+
+        CompletableFuture<Void> cleanerRun = CompletableFuture.runAsync(timer::fire);
+        try {
+            assertTrue(channel.awaitActiveCheck(), "cleaner must reach the active-channel check");
+            assertSame(channel, pool.poll(KEY));
+            assertTrue(pool.offer(channel, KEY));
+        } finally {
+            channel.resumeActiveCheck();
+        }
+        cleanerRun.get(5, TimeUnit.SECONDS);
+
+        assertTrue(channel.isActive(), "cleaner must not close a freshly re-offered channel");
+        assertEquals(1, partitionSize(pool, KEY), "cleaner must not unlink the fresh generation");
+        assertSame(channel, pool.poll(KEY), "freshly re-offered channel must remain leasable");
+
+        pool.destroy();
+    }
+
+    @Test
+    public void idleGenerationChangesWhenTimestampIsReused() {
+        DefaultChannelPool.IdleState idleState = new DefaultChannelPool.IdleState();
+        assertTrue(idleState.reset(1));
+        long firstGeneration = idleState.snapshot();
+
+        assertTrue(idleState.takeOwnership());
+        assertTrue(idleState.reset(1));
+
+        assertFalse(idleState.tryTakeOwnership(firstGeneration));
+        assertFalse(idleState.isOwned());
+    }
+
     // ---- reap pass unlinks many channels in a single tick (O(n) iterator-remove) ----
 
     @Test
@@ -415,11 +472,15 @@ public void idleCountPerHostCountsOnlyLeasableChannels() {
 
     @Test
     public void concurrentOfferPollRemoveAllIsConsistent() throws Exception {
-        // Real timer so the cleaner reaps tombstones concurrently with offer/poll/removeAll.
-        // TTL only (idle disabled) so the cleaner never closes our shared EmbeddedChannels cross-thread.
+        // Real timer so the cleaner claims expired entries and reaps tombstones concurrently with
+        // offer/poll/removeAll.
         HashedWheelTimer timer = new HashedWheelTimer(10, TimeUnit.MILLISECONDS);
-        DefaultChannelPool pool = new DefaultChannelPool(Duration.ZERO, Duration.ofHours(1),
+        DefaultChannelPool pool = new DefaultChannelPool(Duration.ofMillis(20), Duration.ofHours(1),
                 PoolLeaseStrategy.LIFO, timer, Duration.ofMillis(10));
+        CountDownLatch cleanerClosedChannel = new CountDownLatch(1);
+        Channel expiringChannel = new EmbeddedChannel();
+        expiringChannel.closeFuture().addListener(future -> cleanerClosedChannel.countDown());
+        assertTrue(pool.offer(expiringChannel, "expiring-partition"));
 
         final int channelCount = 16;
         Channel[] channels = new Channel[channelCount];
@@ -473,6 +534,8 @@ public void concurrentOfferPollRemoveAllIsConsistent() throws Exception {
         if (failure.get() != null) {
             fail("worker threw: " + failure.get(), failure.get());
         }
+        assertTrue(cleanerClosedChannel.await(5, TimeUnit.SECONDS),
+                "soak must exercise the cleaner ownership and close path");
         assertTrue(leasedInactive.isEmpty(), "poll must never lease an inactive channel");
 
         // Drain leases, then let the cleaner run a couple of ticks and confirm no tombstone leak:
@@ -490,14 +553,52 @@ public void concurrentOfferPollRemoveAllIsConsistent() throws Exception {
 
     // ---- helpers ----
 
-    private static Object idleState(Channel channel) throws Exception {
+    private static DefaultChannelPool.IdleState idleState(Channel channel) throws Exception {
         Field keyField = DefaultChannelPool.class.getDeclaredField("IDLE_STATE_ATTRIBUTE_KEY");
         keyField.setAccessible(true);
         @SuppressWarnings("unchecked")
-        io.netty.util.AttributeKey<Object> key = (io.netty.util.AttributeKey<Object>) keyField.get(null);
+        io.netty.util.AttributeKey<DefaultChannelPool.IdleState> key =
+                (io.netty.util.AttributeKey<DefaultChannelPool.IdleState>) keyField.get(null);
         return channel.attr(key).get();
     }
 
+    private static final class BlockingActiveChannel extends EmbeddedChannel {
+
+        private final AtomicBoolean blockNextActiveCheck = new AtomicBoolean();
+        private final CountDownLatch activeCheck = new CountDownLatch(1);
+        private final CountDownLatch resumeActiveCheck = new CountDownLatch(1);
+
+        @Override
+        public boolean isActive() {
+            // EmbeddedChannel calls this from its constructor before subclass fields are initialized.
+            AtomicBoolean block = blockNextActiveCheck;
+            if (block != null && block.compareAndSet(true, false)) {
+                activeCheck.countDown();
+                try {
+                    if (!resumeActiveCheck.await(5, TimeUnit.SECONDS)) {
+                        throw new AssertionError("active-channel check was not resumed");
+                    }
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    throw new AssertionError("active-channel check was interrupted", e);
+                }
+            }
+            return super.isActive();
+        }
+
+        void blockNextActiveCheck() {
+            blockNextActiveCheck.set(true);
+        }
+
+        boolean awaitActiveCheck() throws InterruptedException {
+            return activeCheck.await(5, TimeUnit.SECONDS);
+        }
+
+        void resumeActiveCheck() {
+            resumeActiveCheck.countDown();
+        }
+    }
+
     private static Channel channelWithRemoteAddress(String host) {
         return new EmbeddedChannel() {