Respect TTL of a DNS record for proxy config

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

@@ -21,6 +21,9 @@
 import java.net.InetSocketAddress;
 import java.util.function.BiConsumer;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import com.linecorp.armeria.client.HttpChannelPool.PoolKey;
 import com.linecorp.armeria.client.endpoint.EmptyEndpointGroupException;
 import com.linecorp.armeria.client.proxy.HAProxyConfig;
@@ -50,6 +53,7 @@
 
 final class HttpClientDelegate implements HttpClient {
 
+    private static final Logger logger = LoggerFactory.getLogger(HttpClientDelegate.class);
     private final HttpClientFactory factory;
     private final AddressResolverGroup<InetSocketAddress> addressResolverGroup;
 
@@ -88,24 +92,37 @@ 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) {
+                    earlyFailed(thrown, ctx);
+                    res.close(thrown);
+                } 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);
+            final Throwable t = earlyFailed(cancellationCause, ctx);
+            res.close(t);
+            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 +142,6 @@ public HttpResponse execute(ClientRequestContext ctx, HttpRequest req) throws Ex
                 }
             });
         }
-
-        return res;
     }
 
     private static void updateCancellationTask(ClientRequestContext ctx, HttpRequest req,
@@ -215,24 +230,58 @@ 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 Future<InetSocketAddress> resolveFuture = addressResolverGroup
+                    .getResolver(ctx.eventLoop().withoutContext())
+                    .resolve(proxyConfig.proxyAddress());
+
+            resolveFuture.addListener(future -> {
+                if (future.isSuccess()) {
+                    final InetSocketAddress resolvedAddress = (InetSocketAddress) future.getNow();
+                    if (resolvedAddress != null && !resolvedAddress.isUnresolved()) {
+                        final ProxyConfig newProxyConfig = proxyConfig.withNewProxyAddress(resolvedAddress);
+                        onComplete.accept(maybeHAProxy(newProxyConfig, capturedProxiedAddresses), null);
+                    } else {
+                        logger.warn("Resolved address is invalid or unresolved: {}. " +
+                                    "Using the previous address instead.", resolvedAddress);
+                        onComplete.accept(maybeHAProxy(proxyConfig, capturedProxiedAddresses), null);
+                    }
+                } else {
+                    final Throwable cause = future.cause();
+                    logger.warn("Failed to refresh {}'s ip address. " +
+                                "Using the previous address instead. reason: {}",
+                                proxyConfig.proxyAddress(),
+                                cause != null ? cause.getMessage() : "Unknown Error");
+                    onComplete.accept(maybeHAProxy(proxyConfig, capturedProxiedAddresses), null);
+                }
+            });
+        } 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;
     }
 
@@ -247,9 +296,14 @@ private static void logSession(ClientRequestContext ctx, @Nullable PooledChannel
         }
     }
 
-    private static HttpResponse earlyFailedResponse(Throwable t, ClientRequestContext ctx) {
+    private static Throwable earlyFailed(Throwable t, ClientRequestContext ctx) {
         final UnprocessedRequestException cause = UnprocessedRequestException.of(t);
         ctx.cancel(cause);
+        return cause;
+    }
+
+    private static HttpResponse earlyFailedResponse(Throwable t, ClientRequestContext ctx) {
+        final Throwable cause = earlyFailed(t, ctx);
         return HttpResponse.ofFailure(cause);
     }
 

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

@@ -90,6 +90,12 @@ public ProxyType proxyType() {
         return ProxyType.CONNECT;
     }
 
+    @Override
+    public ProxyConfig withNewProxyAddress(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,11 @@ public InetSocketAddress proxyAddress() {
         return null;
     }
 
+    @Override
+    public ProxyConfig withNewProxyAddress(InetSocketAddress newProxyAddress) {
+        return new DirectProxyConfig();
+    }
+
     @Override
     public String toString() {
         return "DirectProxyConfig{proxyType=DIRECT}";

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

@@ -67,6 +67,12 @@ public InetSocketAddress sourceAddress() {
         return sourceAddress;
     }
 
+    @Override
+    public ProxyConfig withNewProxyAddress(InetSocketAddress newProxyAddress) {
+        return this.sourceAddress == null ? new HAProxyConfig(proxyAddress)
+                                          : new HAProxyConfig(proxyAddress, this.sourceAddress);
+    }
+
     @Override
     public boolean equals(@Nullable Object o) {
         if (this == o) {

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

@@ -201,6 +201,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 withNewProxyAddress(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 withNewProxyAddress(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 withNewProxyAddress(InetSocketAddress newProxyAddress) {
+        return new Socks5ProxyConfig(newProxyAddress, this.username, this.password);
+    }
+
     @Override
     public boolean equals(@Nullable Object o) {
         if (this == o) {