Allow AhcEndpoint address override in grizzly provider

Author: elrodro83

src/main/java/com/ning/http/client/providers/grizzly/ConnectionManager.java

@@ -129,12 +129,12 @@ void openAsync(final Request request,
             port = getPort(scheme, uri.getPort());
         }
 
-        final String partitionId = getPartitionId(request, proxy);
+        final String partitionId = getPartitionId(request.getInetAddress(), request, proxy);
         Endpoint endpoint = endpointMap.get(partitionId);
         if (endpoint == null) {
             final boolean isSecure = Utils.isSecure(scheme);
             endpoint = new AhcEndpoint(partitionId,
-                    isSecure, host, port, request.getLocalAddress(),
+                    isSecure, request.getInetAddress(), host, port, request.getLocalAddress(),
                     defaultConnectionHandler);
 
             endpointMap.put(partitionId, endpoint);
@@ -164,11 +164,11 @@ Connection openSync(final Request request)
 
         final boolean isSecure = Utils.isSecure(scheme);
 
-        final String partitionId = getPartitionId(request, proxy);
+        final String partitionId = getPartitionId(request.getInetAddress(), request, proxy);
         Endpoint endpoint = endpointMap.get(partitionId);
         if (endpoint == null) {
             endpoint = new AhcEndpoint(partitionId,
-                    isSecure, host, port, request.getLocalAddress(),
+                    isSecure, request.getInetAddress(), host, port, request.getLocalAddress(),
                     defaultConnectionHandler);
 
             endpointMap.put(partitionId, endpoint);
@@ -219,10 +219,11 @@ static boolean isKeepAlive(final Connection connection) {
         return !IS_NOT_KEEP_ALIVE.isSet(connection);
     }
 
-    private static String getPartitionId(Request request,
+    private static String getPartitionId(InetAddress overrideAddress, Request request,
             ProxyServer proxyServer) {
-        return request.getConnectionPoolPartitioning()
-                .getPartitionKey(request.getUri(), proxyServer).toString();
+        return (overrideAddress != null ? overrideAddress.toString() + "_" : "") + 
+            request.getConnectionPoolPartitioning()
+             .getPartitionKey(request.getUri(), proxyServer).toString();
     }
 
     private static int getPort(final String scheme, final int p) {
@@ -244,19 +245,21 @@ private class AhcEndpoint extends Endpoint<SocketAddress> {
 
         private final String partitionId;
         private final boolean isSecure;
+        private final InetAddress remoteOverrideAddress;
         private final String host;
         private final int port;
         private final InetAddress localAddress;
         private final ConnectorHandler<SocketAddress> connectorHandler;
 
         private AhcEndpoint(final String partitionId,
                 final boolean isSecure,
-                final String host, final int port,
+                final InetAddress remoteOverrideAddress, final String host, final int port,
                 final InetAddress localAddress,
                 final ConnectorHandler<SocketAddress> connectorHandler) {
 
             this.partitionId = partitionId;
             this.isSecure = isSecure;
+            this.remoteOverrideAddress = remoteOverrideAddress;
             this.host = host;
             this.port = port;
             this.localAddress = localAddress;
@@ -275,12 +278,19 @@ public Object getId() {
         @Override
         public GrizzlyFuture<Connection> connect() {
             return (GrizzlyFuture<Connection>) connectorHandler.connect(
-                    new InetSocketAddress(host, port),
+                    buildRemoteSocketAddress(),
                     localAddress != null
                             ? new InetSocketAddress(localAddress, 0)
                             : null);
         }
 
+        private InetSocketAddress buildRemoteSocketAddress()
+        {
+            return remoteOverrideAddress != null
+                    ? new InetSocketAddress(remoteOverrideAddress, port)
+                    : new InetSocketAddress(host, port);
+        }
+
         @Override
         protected void onConnect(final Connection connection,
                 final SingleEndpointPool<SocketAddress> pool) {