Backport ConnectionPoolKeyStrategy for 1.7.x branch

Author: Stephane Landelle

src/main/java/com/ning/http/client/ConnectionPoolKeyStrategy.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright 2010 Ning, Inc.
+ *
+ * Ning licenses this file to you under the Apache License, version 2.0
+ * (the "License"); you may not use this file except in compliance with the
+ * License.  You may obtain a copy of the License at:
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+package com.ning.http.client;
+
+import java.net.URI;
+
+public interface ConnectionPoolKeyStrategy {
+
+	String getKey(URI uri);
+}

src/main/java/com/ning/http/client/DefaultConnectionPoolStrategy.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2010 Ning, Inc.
+ *
+ * Ning licenses this file to you under the Apache License, version 2.0
+ * (the "License"); you may not use this file except in compliance with the
+ * License.  You may obtain a copy of the License at:
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+package com.ning.http.client;
+
+import java.net.URI;
+
+import com.ning.http.util.AsyncHttpProviderUtils;
+
+public enum DefaultConnectionPoolStrategy implements ConnectionPoolKeyStrategy {
+
+	INSTANCE;
+	
+	public String getKey(URI uri) {
+		return AsyncHttpProviderUtils.getBaseUrl(uri);
+	}
+}

src/main/java/com/ning/http/client/Request.java

@@ -232,4 +232,5 @@ public static interface EntityWriter {
 
     public boolean isUseRawUrl();
 
+    ConnectionPoolKeyStrategy getConnectionPoolKeyStrategy();
 }

src/main/java/com/ning/http/client/RequestBuilderBase.java

@@ -67,6 +67,7 @@ private static final class RequestImpl implements Request {
         private long rangeOffset = 0;
         public String charset;
         private boolean useRawUrl = false;
+        private ConnectionPoolKeyStrategy connectionPoolKeyStrategy = DefaultConnectionPoolStrategy.INSTANCE;
 
         public RequestImpl(boolean useRawUrl) {
             this.useRawUrl = useRawUrl;
@@ -100,6 +101,7 @@ public RequestImpl(Request prototype) {
                 this.rangeOffset = prototype.getRangeOffset();
                 this.charset = prototype.getBodyEncoding();
                 this.useRawUrl = prototype.isUseRawUrl();
+                this.connectionPoolKeyStrategy = prototype.getConnectionPoolKeyStrategy();
             }
         }
 
@@ -287,6 +289,10 @@ public String getBodyEncoding() {
             return charset;
         }
 
+        public ConnectionPoolKeyStrategy getConnectionPoolKeyStrategy() {
+        	return connectionPoolKeyStrategy;
+       }
+        
         @Override
         public String toString() {
             StringBuilder sb = new StringBuilder(url);
@@ -603,6 +609,11 @@ public T setBodyEncoding(String charset) {
         return derived.cast(this);
     }
 
+    public T setConnectionPoolKeyStrategy(ConnectionPoolKeyStrategy connectionPoolKeyStrategy) {
+    	request.connectionPoolKeyStrategy = connectionPoolKeyStrategy;
+    	return derived.cast(this);
+    }
+
     public Request build() {
         if ((request.length < 0) && (request.streamData == null) && allowBody(request.getMethod())) {
             // can't concatenate content-length

src/main/java/com/ning/http/client/providers/netty/NettyAsyncHttpProvider.java

@@ -21,6 +21,7 @@
 import com.ning.http.client.AsyncHttpProvider;
 import com.ning.http.client.Body;
 import com.ning.http.client.BodyGenerator;
+import com.ning.http.client.ConnectionPoolKeyStrategy;
 import com.ning.http.client.ConnectionsPool;
 import com.ning.http.client.Cookie;
 import com.ning.http.client.FluentCaseInsensitiveStringsMap;
@@ -381,8 +382,8 @@ public ChannelPipeline getPipeline() throws Exception {
         }
     }
 
-    private Channel lookupInCache(URI uri) {
-        final Channel channel = connectionsPool.poll(AsyncHttpProviderUtils.getBaseUrl(uri));
+    private Channel lookupInCache(URI uri, ConnectionPoolKeyStrategy connectionPoolKeyStrategy) {
+        final Channel channel = connectionsPool.poll(connectionPoolKeyStrategy.getKey(uri));
 
         if (channel != null) {
             log.debug("Using cached Channel {}\n for uri {}\n", channel, uri);
@@ -955,7 +956,7 @@ private <T> ListenableFuture<T> doConnect(final Request request, final AsyncHand
             if (f != null && f.reuseChannel() && f.channel() != null) {
                 channel = f.channel();
             } else {
-                channel = lookupInCache(uri);
+                channel = lookupInCache(uri, request.getConnectionPoolKeyStrategy());
             }
         }
 
@@ -1318,7 +1319,7 @@ private Realm ntlmProxyChallenge(List<String> wwwAuth,
     private void drainChannel(final ChannelHandlerContext ctx, final NettyResponseFuture<?> future, final boolean keepAlive, final URI uri) {
         ctx.setAttachment(new AsyncCallable(future) {
             public Object call() throws Exception {
-                if (keepAlive && ctx.getChannel().isReadable() && connectionsPool.offer(AsyncHttpProviderUtils.getBaseUrl(uri), ctx.getChannel())) {
+                if (keepAlive && ctx.getChannel().isReadable() && connectionsPool.offer(future.getConnectionPoolKeyStrategy().getKey(uri), ctx.getChannel())) {
                     return null;
                 }
 
@@ -1514,7 +1515,7 @@ private void finishUpdate(final NettyResponseFuture<?> future, final ChannelHand
             drainChannel(ctx, future, future.getKeepAlive(), future.getURI());
         } else {
             if (future.getKeepAlive() && ctx.getChannel().isReadable() &&
-                    connectionsPool.offer(AsyncHttpProviderUtils.getBaseUrl(future.getURI()), ctx.getChannel())) {
+                    connectionsPool.offer(future.getConnectionPoolKeyStrategy().getKey(future.getURI()), ctx.getChannel())) {
                 markAsDone(future, ctx);
                 return;
             }
@@ -1715,7 +1716,7 @@ public static <T> NettyResponseFuture<T> newFuture(URI uri,
                                                        NettyAsyncHttpProvider provider) {
 
         NettyResponseFuture<T> f = new NettyResponseFuture<T>(uri, request, asyncHandler, nettyRequest,
-                requestTimeout(config, request.getPerRequestConfig()), config.getIdleConnectionTimeoutInMs(), provider);
+                requestTimeout(config, request.getPerRequestConfig()), config.getIdleConnectionTimeoutInMs(), provider, request.getConnectionPoolKeyStrategy());
 
         if (request.getHeaders().getFirstValue("Expect") != null
                 && request.getHeaders().getFirstValue("Expect").equalsIgnoreCase("100-Continue")) {
@@ -2085,10 +2086,11 @@ private boolean redirect(Request request,
                         nBuilder.addOrReplaceCookie(c);
                     }
 
+                    final String connectionPoolKey = future.getConnectionPoolKeyStrategy().getKey(initialConnectionUri);
                     AsyncCallable ac = new AsyncCallable(future) {
                         public Object call() throws Exception {
                             if (initialConnectionKeepAlive && ctx.getChannel().isReadable() &&
-                                    connectionsPool.offer(AsyncHttpProviderUtils.getBaseUrl(initialConnectionUri), ctx.getChannel())) {
+                                    connectionsPool.offer(connectionPoolKey, ctx.getChannel())) {
                                 return null;
                             }
                             finishChannel(ctx);

src/main/java/com/ning/http/client/providers/netty/NettyResponseFuture.java

@@ -16,6 +16,7 @@
 package com.ning.http.client.providers.netty;
 
 import com.ning.http.client.AsyncHandler;
+import com.ning.http.client.ConnectionPoolKeyStrategy;
 import com.ning.http.client.Request;
 import com.ning.http.client.listenable.AbstractListenableFuture;
 import org.jboss.netty.channel.Channel;
@@ -85,14 +86,16 @@ enum STATE {
     private boolean writeBody;
     private final AtomicBoolean throwableCalled = new AtomicBoolean(false);
     private boolean allowConnect = false;
+    private final ConnectionPoolKeyStrategy connectionPoolKeyStrategy;
 
     public NettyResponseFuture(URI uri,
                                Request request,
                                AsyncHandler<V> asyncHandler,
                                HttpRequest nettyRequest,
                                int responseTimeoutInMs,
                                int idleConnectionTimeoutInMs,
-                               NettyAsyncHttpProvider asyncHttpProvider) {
+                               NettyAsyncHttpProvider asyncHttpProvider,
+                                  ConnectionPoolKeyStrategy connectionPoolKeyStrategy) {
 
         this.asyncHandler = asyncHandler;
         this.responseTimeoutInMs = responseTimeoutInMs;
@@ -101,6 +104,7 @@ public NettyResponseFuture(URI uri,
         this.nettyRequest = nettyRequest;
         this.uri = uri;
         this.asyncHttpProvider = asyncHttpProvider;
+        this.connectionPoolKeyStrategy = connectionPoolKeyStrategy;
 
         if (System.getProperty(MAX_RETRY) != null) {
             maxRetry = Integer.valueOf(System.getProperty(MAX_RETRY));
@@ -119,6 +123,10 @@ protected void setURI(URI uri) {
         this.uri = uri;
     }
 
+    public ConnectionPoolKeyStrategy getConnectionPoolKeyStrategy() {
+        return connectionPoolKeyStrategy;
+    }
+
     /**
      * {@inheritDoc}
      */