Disable connection reuse for VIPs

2 files changed, 186 insertions, 18 deletions

diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/proxy/ConfigServerRestExecutorImpl.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/proxy/ConfigServerRestExecutorImpl.java
index fcfc4eca32b..f5dcae9c961 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/proxy/ConfigServerRestExecutorImpl.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/proxy/ConfigServerRestExecutorImpl.java
@@ -10,6 +10,7 @@ import com.yahoo.vespa.athenz.identity.ServiceIdentityProvider;
 import com.yahoo.vespa.athenz.tls.AthenzIdentityVerifier;
 import com.yahoo.vespa.hosted.controller.api.integration.zone.ZoneRegistry;
 import org.apache.http.Header;
+import org.apache.http.HttpResponse;
 import org.apache.http.client.config.RequestConfig;
 import org.apache.http.client.methods.CloseableHttpResponse;
 import org.apache.http.client.methods.HttpDelete;
@@ -19,11 +20,15 @@ import org.apache.http.client.methods.HttpPost;
 import org.apache.http.client.methods.HttpPut;
 import org.apache.http.client.methods.HttpRequestBase;
 import org.apache.http.entity.InputStreamEntity;
+import org.apache.http.impl.DefaultConnectionReuseStrategy;
 import org.apache.http.impl.client.CloseableHttpClient;
 import org.apache.http.impl.client.HttpClientBuilder;
+import org.apache.http.protocol.HttpContext;
+import org.apache.http.protocol.HttpCoreContext;
 import org.apache.http.util.EntityUtils;
 
 import javax.net.ssl.HostnameVerifier;
+import javax.net.ssl.SSLContext;
 import javax.net.ssl.SSLSession;
 import java.io.IOException;
 import java.io.InputStream;
@@ -63,14 +68,16 @@ public class ConfigServerRestExecutorImpl extends AbstractComponent implements C
 
     @Inject
     public ConfigServerRestExecutorImpl(ZoneRegistry zoneRegistry, ServiceIdentityProvider sslContextProvider) {
-        RequestConfig config = RequestConfig.custom()
-                .setConnectTimeout((int) PROXY_REQUEST_TIMEOUT.toMillis())
-                .setConnectionRequestTimeout((int) PROXY_REQUEST_TIMEOUT.toMillis())
-                .setSocketTimeout((int) PROXY_REQUEST_TIMEOUT.toMillis()).build();
+        this(zoneRegistry, sslContextProvider.getIdentitySslContext(), new Sleeper.Default(),
+             new ConnectionReuseStrategy(zoneRegistry));
+    }
 
-        this.client = createHttpClient(config, sslContextProvider,
-                new ControllerOrConfigserverHostnameVerifier(zoneRegistry));
-        this.sleeper = new Sleeper.Default();
+    ConfigServerRestExecutorImpl(ZoneRegistry zoneRegistry, SSLContext sslContext,
+                                 Sleeper sleeper, ConnectionReuseStrategy connectionReuseStrategy) {
+        this.client = createHttpClient(sslContext,
+                                       new ControllerOrConfigserverHostnameVerifier(zoneRegistry),
+                                       connectionReuseStrategy);
+        this.sleeper = sleeper;
     }
 
     @Override
@@ -218,18 +225,25 @@ public class ConfigServerRestExecutorImpl extends AbstractComponent implements C
         }
     }
 
-    private static CloseableHttpClient createHttpClient(RequestConfig config,
-                                                        ServiceIdentityProvider sslContextProvider,
-                                                        HostnameVerifier hostnameVerifier) {
+    private static CloseableHttpClient createHttpClient(SSLContext sslContext,
+                                                        HostnameVerifier hostnameVerifier,
+                                                        org.apache.http.ConnectionReuseStrategy connectionReuseStrategy) {
+
+        RequestConfig config = RequestConfig.custom()
+                                            .setConnectTimeout((int) PROXY_REQUEST_TIMEOUT.toMillis())
+                                            .setConnectionRequestTimeout((int) PROXY_REQUEST_TIMEOUT.toMillis())
+                                            .setSocketTimeout((int) PROXY_REQUEST_TIMEOUT.toMillis()).build();
         return HttpClientBuilder.create()
-                .setUserAgent("config-server-proxy-client")
-                .setSSLContext(sslContextProvider.getIdentitySslContext())
-                .setSSLHostnameVerifier(hostnameVerifier)
-                .setDefaultRequestConfig(config)
-                .setMaxConnPerRoute(10)
-                .setMaxConnTotal(500)
-                .setConnectionTimeToLive(1, TimeUnit.MINUTES)
-                .build();
+                                .setUserAgent("config-server-proxy-client")
+                                .setSSLContext(sslContext)
+                                .setSSLHostnameVerifier(hostnameVerifier)
+                                .setDefaultRequestConfig(config)
+                                .setMaxConnPerRoute(10)
+                                .setMaxConnTotal(500)
+                                .setConnectionReuseStrategy(connectionReuseStrategy)
+                                .setConnectionTimeToLive(1, TimeUnit.MINUTES)
+                                .build();
+
     }
 
     private static class ControllerOrConfigserverHostnameVerifier implements HostnameVerifier {
@@ -253,4 +267,35 @@ public class ConfigServerRestExecutorImpl extends AbstractComponent implements C
         }
     }
 
+    /**
+     * A connection reuse strategy which avoids reusing connections to VIPs. Since VIPs are TCP-level load balancers,
+     * a reconnect is needed to (potentially) switch real server.
+     */
+    public static class ConnectionReuseStrategy extends DefaultConnectionReuseStrategy {
+
+        private final Set<String> vips;
+
+        public ConnectionReuseStrategy(ZoneRegistry zoneRegistry) {
+            this(zoneRegistry.zones().all().ids().stream()
+                             .map(zoneRegistry::getConfigServerVipUri)
+                             .map(URI::getHost)
+                             .collect(Collectors.toUnmodifiableSet()));
+        }
+
+        public ConnectionReuseStrategy(Set<String> vips) {
+            this.vips = Set.copyOf(vips);
+        }
+
+        @Override
+        public boolean keepAlive(HttpResponse response, HttpContext context) {
+            HttpCoreContext coreContext = HttpCoreContext.adapt(context);
+            String host = coreContext.getTargetHost().getHostName();
+            if (vips.contains(host)) {
+                return false;
+            }
+            return super.keepAlive(response, context);
+        }
+
+    }
+
 }
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/proxy/ConfigServerRestExecutorImplTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/proxy/ConfigServerRestExecutorImplTest.java
new file mode 100644
index 00000000000..1fce7ba5695
--- /dev/null
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/proxy/ConfigServerRestExecutorImplTest.java
@@ -0,0 +1,123 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.controller.proxy;
+
+import com.github.tomakehurst.wiremock.junit.WireMockRule;
+import com.github.tomakehurst.wiremock.stubbing.Scenario;
+import com.yahoo.config.provision.SystemName;
+import com.yahoo.container.jdisc.HttpRequest;
+import com.yahoo.container.jdisc.HttpResponse;
+import com.yahoo.vespa.hosted.controller.integration.ZoneRegistryMock;
+import org.apache.http.protocol.HttpContext;
+import org.apache.http.protocol.HttpCoreContext;
+import org.junit.Rule;
+import org.junit.Test;
+
+import javax.net.ssl.SSLContext;
+import java.io.ByteArrayOutputStream;
+import java.net.URI;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+
+import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
+import static com.github.tomakehurst.wiremock.client.WireMock.get;
+import static com.github.tomakehurst.wiremock.client.WireMock.getRequestedFor;
+import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo;
+import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.options;
+import static org.junit.Assert.assertEquals;
+
+/**
+ * @author mpolden
+ */
+public class ConfigServerRestExecutorImplTest {
+
+    @Rule
+    public final WireMockRule wireMock = new WireMockRule(options().dynamicPort(), true);
+
+    @Test
+    public void proxy_with_retries() throws Exception {
+        var connectionReuseStrategy = new CountingConnectionReuseStrategy(Set.of("127.0.0.1"));
+        var proxy = new ConfigServerRestExecutorImpl(new ZoneRegistryMock(SystemName.cd), SSLContext.getDefault(),
+                                                     (duration) -> {}, connectionReuseStrategy);
+
+        URI url = url();
+        String path = url.getPath();
+        stubRequests(path);
+
+        HttpRequest request = HttpRequest.createTestRequest(url.toString(), com.yahoo.jdisc.http.HttpRequest.Method.GET);
+        ProxyRequest proxyRequest = ProxyRequest.tryOne(url, path, request);
+
+        // Request is retried
+        HttpResponse response = proxy.handle(proxyRequest);
+        wireMock.verify(3, getRequestedFor(urlEqualTo(path)));
+        assertEquals(200, response.getStatus());
+        ByteArrayOutputStream out = new ByteArrayOutputStream();
+        response.render(out);
+        assertEquals("OK", out.toString());
+
+        // No connections are reused as host is a VIP
+        assertEquals(0, connectionReuseStrategy.reusedConnections.get(url.getHost()).intValue());
+    }
+
+    @Test
+    public void proxy_without_connection_reuse() throws Exception {
+        var connectionReuseStrategy = new CountingConnectionReuseStrategy(Set.of());
+        var proxy = new ConfigServerRestExecutorImpl(new ZoneRegistryMock(SystemName.cd), SSLContext.getDefault(),
+                                                     (duration) -> {}, connectionReuseStrategy);
+
+        URI url = url();
+        String path = url.getPath();
+        stubRequests(path);
+
+        HttpRequest request = HttpRequest.createTestRequest(url.toString(), com.yahoo.jdisc.http.HttpRequest.Method.GET);
+        ProxyRequest proxyRequest = ProxyRequest.tryOne(url, path, request);
+
+        // Connections are reused
+        assertEquals(200, proxy.handle(proxyRequest).getStatus());
+        assertEquals(3, connectionReuseStrategy.reusedConnections.get(url.getHost()).intValue());
+    }
+
+    private URI url() {
+        return URI.create("http://127.0.0.1:" + wireMock.port() + "/");
+    }
+
+    private void stubRequests(String path) {
+        String retryScenario = "Retry scenario";
+        String retryRequest = "Retry request 1";
+        String retryRequestAgain = "Retry request 2";
+
+        wireMock.stubFor(get(urlEqualTo(path)).inScenario(retryScenario)
+                                              .whenScenarioStateIs(Scenario.STARTED)
+                                              .willSetStateTo(retryRequest)
+                                              .willReturn(aResponse().withStatus(500)));
+
+        wireMock.stubFor(get(urlEqualTo(path)).inScenario(retryScenario)
+                                              .whenScenarioStateIs(retryRequest)
+                                              .willSetStateTo(retryRequestAgain)
+                                              .willReturn(aResponse().withStatus(500)));
+
+        wireMock.stubFor(get(urlEqualTo(path)).inScenario(retryScenario)
+                                              .whenScenarioStateIs(retryRequestAgain)
+                                              .willReturn(aResponse().withBody("OK")));
+    }
+
+    private static class CountingConnectionReuseStrategy extends ConfigServerRestExecutorImpl.ConnectionReuseStrategy {
+
+        private final Map<String, Integer> reusedConnections = new HashMap<>();
+
+        public CountingConnectionReuseStrategy(Set<String> vips) {
+            super(vips);
+        }
+
+        @Override
+        public boolean keepAlive(org.apache.http.HttpResponse response, HttpContext context) {
+            boolean keepAlive = super.keepAlive(response, context);
+            String host = HttpCoreContext.adapt(context).getTargetHost().getHostName();
+            reusedConnections.putIfAbsent(host, 0);
+            if (keepAlive) reusedConnections.compute(host, (ignored, count) -> ++count);
+            return keepAlive;
+        }
+
+    }
+
+}