Respect TTL of a DNS record for proxy config

core/src/main/java/com/linecorp/armeria/client/HttpClientDelegate.java

@@ -88,24 +88,35 @@ public HttpResponse execute(ClientRequestContext ctx, HttpRequest req) throws Ex
         }
 
         final SessionProtocol protocol = ctx.sessionProtocol();
-        final ProxyConfig proxyConfig;
+
+        final Endpoint endpointWithPort = endpoint.withDefaultPort(ctx.sessionProtocol());
+        final EventLoop eventLoop = ctx.eventLoop().withoutContext();
+        // TODO(ikhoon) Use ctx.exchangeType() to create an optimized HttpResponse for non-streaming response.
+        final DecodedHttpResponse res = new DecodedHttpResponse(eventLoop);
+        updateCancellationTask(ctx, req, res);
+
         try {
-            proxyConfig = getProxyConfig(protocol, endpoint);
+            resolveProxyConfig(protocol, endpoint, ctx, (proxyConfig, thrown) -> {
+                if (thrown != null) {
+                    earlyFailedResponse(thrown, ctx, res);
+                } else {
+                    execute0(ctx, endpointWithPort, req, res, proxyConfig);
+                }
+            });
         } catch (Throwable t) {
             return earlyFailedResponse(t, ctx);
         }
+        return res;
+    }
 
+    private void execute0(ClientRequestContext ctx, Endpoint endpointWithPort, HttpRequest req,
+                          DecodedHttpResponse res, ProxyConfig proxyConfig) {
         final Throwable cancellationCause = ctx.cancellationCause();
         if (cancellationCause != null) {
-            return earlyFailedResponse(cancellationCause, ctx);
+            earlyFailedResponse(cancellationCause, ctx, res);
+            return;
         }
 
-        final Endpoint endpointWithPort = endpoint.withDefaultPort(ctx.sessionProtocol());
-        final EventLoop eventLoop = ctx.eventLoop().withoutContext();
-        // TODO(ikhoon) Use ctx.exchangeType() to create an optimized HttpResponse for non-streaming response.
-        final DecodedHttpResponse res = new DecodedHttpResponse(eventLoop);
-        updateCancellationTask(ctx, req, res);
-
         final ClientConnectionTimingsBuilder timingsBuilder = ClientConnectionTimings.builder();
 
         if (endpointWithPort.hasIpAddr() ||
@@ -125,8 +136,6 @@ public HttpResponse execute(ClientRequestContext ctx, HttpRequest req) throws Ex
                 }
             });
         }
-
-        return res;
     }
 
     private static void updateCancellationTask(ClientRequestContext ctx, HttpRequest req,
@@ -215,24 +224,53 @@ private void acquireConnectionAndExecute0(ClientRequestContext ctx, Endpoint end
         }
     }
 
-    private ProxyConfig getProxyConfig(SessionProtocol protocol, Endpoint endpoint) {
+    private void resolveProxyConfig(SessionProtocol protocol, Endpoint endpoint, ClientRequestContext ctx,
+                                    BiConsumer<@Nullable ProxyConfig, @Nullable Throwable> onComplete) {
         final ProxyConfig proxyConfig = factory.proxyConfigSelector().select(protocol, endpoint);
         requireNonNull(proxyConfig, "proxyConfig");
 
+        final ServiceRequestContext serviceCtx = ServiceRequestContext.currentOrNull();
+        final ProxiedAddresses capturedProxiedAddresses = serviceCtx == null ? null
+                                                          : serviceCtx.proxiedAddresses();
+        // DirectProxyConfig does not have proxyAddress as its field.
+        if (proxyConfig.proxyAddress() != null) {
+            final InetSocketAddress currentAddr = proxyConfig.proxyAddress();
+            if (!currentAddr.isUnresolved() && (currentAddr.getHostName() == currentAddr.getHostString())) {
+                // If InetSocket is created from IP Address, we don't need to refresh DNS.
+                onComplete.accept(maybeHAProxy(proxyConfig, capturedProxiedAddresses), null);
+            } else {
+                final Future<InetSocketAddress> resolveFuture = addressResolverGroup
+                        .getResolver(ctx.eventLoop().withoutContext())
+                        .resolve(createUnresolvedAddressForRefreshing(proxyConfig.proxyAddress()));
+
+                resolveFuture.addListener(future -> {
+                    if (future.isSuccess()) {
+                        final InetSocketAddress resolvedAddress = (InetSocketAddress) future.getNow();
+                        final ProxyConfig newProxyConfig = proxyConfig.withProxyAddress(resolvedAddress);
+                        onComplete.accept(maybeHAProxy(newProxyConfig, capturedProxiedAddresses), null);
+                    } else {
+                        final Throwable cause = future.cause();
+                        onComplete.accept(null, cause);
+                    }
+                });
+            }
+        } else {
+            onComplete.accept(maybeHAProxy(proxyConfig, capturedProxiedAddresses), null);
+        }
+    }
+
+    private ProxyConfig maybeHAProxy(ProxyConfig proxyConfig, @Nullable ProxiedAddresses proxiedAddresses) {
         // special behavior for haproxy when sourceAddress is null
         if (proxyConfig.proxyType() == ProxyType.HAPROXY &&
             ((HAProxyConfig) proxyConfig).sourceAddress() == null) {
             final InetSocketAddress proxyAddress = proxyConfig.proxyAddress();
             assert proxyAddress != null;
 
             // use proxy information in context if available
-            final ServiceRequestContext serviceCtx = ServiceRequestContext.currentOrNull();
-            if (serviceCtx != null) {
-                final ProxiedAddresses proxiedAddresses = serviceCtx.proxiedAddresses();
+            if (proxiedAddresses != null) {
                 return ProxyConfig.haproxy(proxyAddress, proxiedAddresses.sourceAddress());
             }
         }
-
         return proxyConfig;
     }
 
@@ -253,11 +291,24 @@ private static HttpResponse earlyFailedResponse(Throwable t, ClientRequestContex
         return HttpResponse.ofFailure(cause);
     }
 
+    private static HttpResponse earlyFailedResponse(Throwable t,
+                                                    ClientRequestContext ctx,
+                                                    DecodedHttpResponse response) {
+        final UnprocessedRequestException cause = UnprocessedRequestException.of(t);
+        ctx.cancel(cause);
+        response.close(cause);
+        return response;
+    }
+
     private static void doExecute(PooledChannel pooledChannel, ClientRequestContext ctx,
                                   HttpRequest req, DecodedHttpResponse res) {
         final Channel channel = pooledChannel.get();
         final HttpSession session = HttpSession.get(channel);
         res.init(session.inboundTrafficController());
         session.invoke(pooledChannel, ctx, req, res);
     }
+
+    private static InetSocketAddress createUnresolvedAddressForRefreshing(InetSocketAddress previousAddress) {
+        return InetSocketAddress.createUnresolved(previousAddress.getHostString(), previousAddress.getPort());
+    }
 }

core/src/main/java/com/linecorp/armeria/client/proxy/ConnectProxyConfig.java

@@ -90,6 +90,12 @@ public ProxyType proxyType() {
         return ProxyType.CONNECT;
     }
 
+    @Override
+    public ProxyConfig withProxyAddress(InetSocketAddress newProxyAddress) {
+        return new ConnectProxyConfig(newProxyAddress, this.username,
+                                      this.password, this.headers, this.useTls);
+    }
+
     @Override
     public boolean equals(@Nullable Object o) {
         if (this == o) {

core/src/main/java/com/linecorp/armeria/client/proxy/DirectProxyConfig.java

@@ -40,6 +40,12 @@ public InetSocketAddress proxyAddress() {
         return null;
     }
 
+    @Override
+    public ProxyConfig withProxyAddress(InetSocketAddress newProxyAddress) {
+        throw new UnsupportedOperationException(
+                "A proxy address can't be set to DirectProxyConfig.");
+    }
+
     @Override
     public String toString() {
         return "DirectProxyConfig{proxyType=DIRECT}";

core/src/main/java/com/linecorp/armeria/client/proxy/HAProxyConfig.java

@@ -17,10 +17,15 @@
 package com.linecorp.armeria.client.proxy;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static java.util.Objects.hash;
+import static java.util.Objects.requireNonNull;
 
 import java.net.InetSocketAddress;
 import java.util.Objects;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import com.google.common.base.MoreObjects;
 
 import com.linecorp.armeria.common.annotation.Nullable;
@@ -31,6 +36,8 @@
  */
 public final class HAProxyConfig extends ProxyConfig {
 
+    private static final Logger logger = LoggerFactory.getLogger(HAProxyConfig.class);
+
     private final InetSocketAddress proxyAddress;
 
     @Nullable
@@ -42,8 +49,18 @@ public final class HAProxyConfig extends ProxyConfig {
     }
 
     HAProxyConfig(InetSocketAddress proxyAddress, InetSocketAddress sourceAddress) {
-        checkArgument(sourceAddress.getAddress().getClass() == proxyAddress.getAddress().getClass(),
-                      "sourceAddress and proxyAddress should be the same type");
+        if (sourceAddress.getAddress() != null &&
+            proxyAddress.getAddress() != null) {
+            checkArgument(sourceAddress.getAddress().getClass() == proxyAddress.getAddress().getClass(),
+                          "sourceAddress and proxyAddress should be the same type");
+        } else {
+            logger.warn("Either the source or proxy address could not be resolved. " +
+                        "The proxy address may be resolved later. " +
+                        "proxyAddress: {}, sourceAddress: {}",
+                        proxyAddress.getAddress(),
+                        sourceAddress.getAddress());
+        }
+
         this.proxyAddress = proxyAddress;
         this.sourceAddress = sourceAddress;
     }
@@ -67,6 +84,13 @@ public InetSocketAddress sourceAddress() {
         return sourceAddress;
     }
 
+    @Override
+    public ProxyConfig withProxyAddress(InetSocketAddress newProxyAddress) {
+        requireNonNull(newProxyAddress, "newProxyAddress");
+        return this.sourceAddress == null ? new HAProxyConfig(proxyAddress)
+                                          : new HAProxyConfig(proxyAddress, this.sourceAddress);
+    }
+
     @Override
     public boolean equals(@Nullable Object o) {
         if (this == o) {
@@ -82,7 +106,7 @@ public boolean equals(@Nullable Object o) {
 
     @Override
     public int hashCode() {
-        return Objects.hash(proxyAddress, sourceAddress);
+        return hash(proxyAddress, sourceAddress);
     }
 
     @Override

core/src/main/java/com/linecorp/armeria/client/proxy/ProxyConfig.java

@@ -40,7 +40,6 @@ public abstract class ProxyConfig {
      */
     public static Socks4ProxyConfig socks4(InetSocketAddress proxyAddress) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         return new Socks4ProxyConfig(proxyAddress, null);
     }
 
@@ -52,7 +51,6 @@ public static Socks4ProxyConfig socks4(InetSocketAddress proxyAddress) {
      */
     public static Socks4ProxyConfig socks4(InetSocketAddress proxyAddress, String username) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         return new Socks4ProxyConfig(proxyAddress, requireNonNull(username, "username"));
     }
 
@@ -63,7 +61,6 @@ public static Socks4ProxyConfig socks4(InetSocketAddress proxyAddress, String us
      */
     public static Socks5ProxyConfig socks5(InetSocketAddress proxyAddress) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         return new Socks5ProxyConfig(proxyAddress, null, null);
     }
 
@@ -77,7 +74,6 @@ public static Socks5ProxyConfig socks5(InetSocketAddress proxyAddress) {
     public static Socks5ProxyConfig socks5(
             InetSocketAddress proxyAddress, String username, String password) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         return new Socks5ProxyConfig(proxyAddress, requireNonNull(username, "username"),
                                      requireNonNull(password, "password"));
     }
@@ -89,7 +85,6 @@ public static Socks5ProxyConfig socks5(
      */
     public static ConnectProxyConfig connect(InetSocketAddress proxyAddress) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         return new ConnectProxyConfig(proxyAddress, null, null, HttpHeaders.of(), false);
     }
 
@@ -101,7 +96,6 @@ public static ConnectProxyConfig connect(InetSocketAddress proxyAddress) {
      */
     public static ConnectProxyConfig connect(InetSocketAddress proxyAddress, boolean useTls) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         return new ConnectProxyConfig(proxyAddress, null, null, HttpHeaders.of(), useTls);
     }
 
@@ -129,7 +123,6 @@ public static ConnectProxyConfig connect(
     public static ConnectProxyConfig connect(
             InetSocketAddress proxyAddress, HttpHeaders headers, boolean useTls) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         return new ConnectProxyConfig(proxyAddress, null, null, headers, useTls);
     }
 
@@ -146,7 +139,6 @@ public static ConnectProxyConfig connect(
     public static ConnectProxyConfig connect(InetSocketAddress proxyAddress, String username, String password,
                                              HttpHeaders headers, boolean useTls) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         requireNonNull(username, "username");
         requireNonNull(password, "password");
         requireNonNull(headers, "headers");
@@ -162,7 +154,6 @@ public static ConnectProxyConfig connect(InetSocketAddress proxyAddress, String
     public static HAProxyConfig haproxy(
             InetSocketAddress proxyAddress, InetSocketAddress sourceAddress) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         requireNonNull(sourceAddress, "sourceAddress");
         checkArgument(!sourceAddress.isUnresolved(), "sourceAddress must be resolved");
         return new HAProxyConfig(proxyAddress, sourceAddress);
@@ -176,7 +167,6 @@ public static HAProxyConfig haproxy(
      */
     public static ProxyConfig haproxy(InetSocketAddress proxyAddress) {
         requireNonNull(proxyAddress, "proxyAddress");
-        checkArgument(!proxyAddress.isUnresolved(), "proxyAddress must be resolved");
         return new HAProxyConfig(proxyAddress);
     }
 
@@ -201,6 +191,12 @@ public static ProxyConfig direct() {
     @Nullable
     public abstract InetSocketAddress proxyAddress();
 
+    /**
+     * Returns a new proxy address instance that respects DNS TTL.
+     * @param newProxyAddress the inet socket address
+     */
+    public abstract ProxyConfig withProxyAddress(InetSocketAddress newProxyAddress);
+
     @Nullable
     static String maskPassword(@Nullable String username, @Nullable String password) {
         return username != null ? "****" : null;

core/src/main/java/com/linecorp/armeria/client/proxy/Socks4ProxyConfig.java

@@ -56,6 +56,11 @@ public ProxyType proxyType() {
         return ProxyType.SOCKS4;
     }
 
+    @Override
+    public ProxyConfig withProxyAddress(InetSocketAddress newProxyAddress) {
+        return new Socks4ProxyConfig(newProxyAddress, this.username);
+    }
+
     @Override
     public boolean equals(@Nullable Object o) {
         if (this == o) {

core/src/main/java/com/linecorp/armeria/client/proxy/Socks5ProxyConfig.java

@@ -69,6 +69,11 @@ public ProxyType proxyType() {
         return ProxyType.SOCKS5;
     }
 
+    @Override
+    public ProxyConfig withProxyAddress(InetSocketAddress newProxyAddress) {
+        return new Socks5ProxyConfig(newProxyAddress, this.username, this.password);
+    }
+
     @Override
     public boolean equals(@Nullable Object o) {
         if (this == o) {