Prevent Netty connection pool from killing active connections, backport AsyncHttpClient#1114

Author: slandelle

src/main/java/com/ning/http/client/providers/netty/channel/ChannelManager.java

@@ -296,9 +296,13 @@ public final void tryToOfferChannelToPool(Channel channel, boolean keepAlive, Ob
         if (channel.isConnected() && keepAlive && channel.isReadable()) {
             LOGGER.debug("Adding key: {} for channel {}", partitionKey, channel);
             Channels.setDiscard(channel);
-            channelPool.offer(channel, partitionKey);
-            if (maxConnectionsPerHostEnabled)
-                channelId2PartitionKey.putIfAbsent(channel.getId(), partitionKey);
+            if (channelPool.offer(channel, partitionKey)) {
+                if (maxConnectionsPerHostEnabled)
+                    channelId2PartitionKey.putIfAbsent(channel.getId(), partitionKey);
+            } else {
+                // rejected by pool
+                closeChannel(channel);
+            }
         } else {
             // not offered
             closeChannel(channel);

src/main/java/com/ning/http/client/providers/netty/channel/pool/DefaultChannelPool.java

@@ -26,7 +26,6 @@
 import com.ning.http.client.AsyncHttpClientConfig;
 import com.ning.http.client.providers.netty.channel.Channels;
 import com.ning.http.client.providers.netty.chmv8.ConcurrentHashMapV8;
-import com.ning.http.client.providers.netty.future.NettyResponseFuture;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -104,6 +103,7 @@ private static final class ChannelCreation {
     private static final class IdleChannel {
         final Channel channel;
         final long start;
+        final AtomicBoolean owned = new AtomicBoolean(false);
 
         IdleChannel(Channel channel, long start) {
             if (channel == null)
@@ -112,6 +112,10 @@ private static final class IdleChannel {
             this.start = start;
         }
 
+        public boolean takeOwnership() {
+            return owned.compareAndSet(false, true);
+        }
+        
         @Override
         // only depends on channel
         public boolean equals(Object o) {
@@ -154,26 +158,16 @@ private List<IdleChannel> expiredChannels(ConcurrentLinkedQueue<IdleChannel> par
             return idleTimeoutChannels != null ? idleTimeoutChannels : Collections.<IdleChannel> emptyList();
         }
 
-        private boolean isChannelCloseable(Channel channel) {
-            Object attribute = Channels.getAttribute(channel);
-            if (attribute instanceof NettyResponseFuture) {
-                NettyResponseFuture<?> future = (NettyResponseFuture<?>) attribute;
-                if (!future.isDone()) {
-                    LOGGER.error("Future not in appropriate state %s, not closing", future);
-                    return false;
-                }
-            }
-            return true;
-        }
-
         private final List<IdleChannel> closeChannels(List<IdleChannel> candidates) {
 
             // lazy create, only if we have a non-closeable channel
             List<IdleChannel> closedChannels = null;
 
             for (int i = 0; i < candidates.size(); i++) {
+                // We call takeOwnership here to avoid closing a channel that has just been taken out
+                // of the pool, otherwise we risk closing an active connection.
                 IdleChannel idleChannel = candidates.get(i);
-                if (isChannelCloseable(idleChannel.channel)) {
+                if (idleChannel.takeOwnership()) {
                     LOGGER.debug("Closing Idle Channel {}", idleChannel.channel);
                     close(idleChannel.channel);
                     if (closedChannels != null) {
@@ -260,12 +254,15 @@ public Channel poll(Object partitionKey) {
             while (idleChannel == null) {
                 idleChannel = partition.poll();
 
-                if (idleChannel == null)
+                if (idleChannel == null) {
                     // pool is empty
                     break;
-                else if (!Channels.isChannelValid(idleChannel.channel)) {
+                } else if (!Channels.isChannelValid(idleChannel.channel)) {
                     idleChannel = null;
                     LOGGER.trace("Channel not connected or not opened, probably remotely closed!");
+                } else if (!idleChannel.takeOwnership()) {
+                    idleChannel = null;
+                    LOGGER.trace("Couldn't take ownership of channel, probably in the process of being expired!");
                 }
             }
         }