Merge pull request #6571 from vespa-engine/bjorncs/sia-closeable-http-client

Bjorncs/sia closeable http client

2 files changed, 137 insertions, 56 deletions

diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java
index 951794798bf..8d18d2f26f3 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/client/zts/DefaultZtsClient.java
@@ -26,6 +26,7 @@ import com.yahoo.vespa.athenz.identity.SiaBackedApacheHttpClient;
 import com.yahoo.vespa.athenz.tls.Pkcs10Csr;
 import com.yahoo.vespa.athenz.tls.Pkcs10CsrBuilder;
 import org.apache.http.HttpResponse;
+import org.apache.http.client.ResponseHandler;
 import org.apache.http.client.config.RequestConfig;
 import org.apache.http.client.methods.HttpUriRequest;
 import org.apache.http.client.methods.RequestBuilder;
@@ -63,7 +64,7 @@ public class DefaultZtsClient implements ZtsClient {
     private static final ObjectMapper objectMapper = new ObjectMapper().registerModule(new JavaTimeModule());
 
     private final URI ztsUrl;
-    private final SiaBackedApacheHttpClient client;
+    private final CloseableHttpClient client;
     private final AthenzIdentity identity;
 
     public DefaultZtsClient(URI ztsUrl, AthenzIdentity identity, SSLContext sslContext) {
@@ -93,7 +94,7 @@ public class DefaultZtsClient implements ZtsClient {
                 .setUri(ztsUrl.resolve("instance/"))
                 .setEntity(toJsonStringEntity(payload))
                 .build();
-        return client.execute(request, DefaultZtsClient::getInstanceIdentity);
+        return execute(request, DefaultZtsClient::getInstanceIdentity);
     }
 
     @Override
@@ -113,7 +114,7 @@ public class DefaultZtsClient implements ZtsClient {
                 .setUri(uri)
                 .setEntity(toJsonStringEntity(payload))
                 .build();
-        return client.execute(request, DefaultZtsClient::getInstanceIdentity);
+        return execute(request, DefaultZtsClient::getInstanceIdentity);
     }
 
     @Override
@@ -123,7 +124,7 @@ public class DefaultZtsClient implements ZtsClient {
                 .setUri(uri)
                 .setEntity(toJsonStringEntity(new IdentityRefreshRequestEntity(csr, keyId)))
                 .build();
-        return client.execute(request, response -> {
+        return execute(request, response -> {
             IdentityResponseEntity entity = readEntity(response, IdentityResponseEntity.class);
             return new Identity(entity.certificate(), entity.caCertificateBundle());
         });
@@ -153,7 +154,7 @@ public class DefaultZtsClient implements ZtsClient {
             requestBuilder.addParameter("role", roleName);
         }
         HttpUriRequest request = requestBuilder.build();
-        return client.execute(request, response -> {
+        return execute(request, response -> {
             RoleTokenResponseEntity roleTokenResponseEntity = readEntity(response, RoleTokenResponseEntity.class);
             return roleTokenResponseEntity.token;
         });
@@ -174,7 +175,7 @@ public class DefaultZtsClient implements ZtsClient {
         HttpUriRequest request = RequestBuilder.post(uri)
                 .setEntity(toJsonStringEntity(requestEntity))
                 .build();
-        return client.execute(request, response -> {
+        return execute(request, response -> {
             RoleCertificateResponseEntity responseEntity = readEntity(response, RoleCertificateResponseEntity.class);
             return responseEntity.certificate;
         });
@@ -195,12 +196,20 @@ public class DefaultZtsClient implements ZtsClient {
                 .addParameter("roleName", roleName)
                 .addParameter("serviceName", providerIdentity.getName())
                 .build();
-        return client.execute(request, response -> {
+        return execute(request, response -> {
             TenantDomainsResponseEntity entity = readEntity(response, TenantDomainsResponseEntity.class);
             return entity.tenantDomainNames.stream().map(AthenzDomain::new).collect(toList());
         });
     }
 
+    private <T> T execute(HttpUriRequest request, ResponseHandler<T> responseHandler) {
+        try {
+            return client.execute(request, responseHandler);
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
     private static InstanceIdentity getInstanceIdentity(HttpResponse response) throws IOException {
         InstanceIdentityCredentials entity = readEntity(response, InstanceIdentityCredentials.class);
         return entity.getServiceToken() != null
@@ -247,7 +256,11 @@ public class DefaultZtsClient implements ZtsClient {
 
     @Override
     public void close() {
-        this.client.close();
+        try {
+            this.client.close();
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
     }
 
 }
diff --git a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/SiaBackedApacheHttpClient.java b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/SiaBackedApacheHttpClient.java
index 588ed5c7abd..c8e405bf886 100644
--- a/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/SiaBackedApacheHttpClient.java
+++ b/vespa-athenz/src/main/java/com/yahoo/vespa/athenz/identity/SiaBackedApacheHttpClient.java
@@ -1,15 +1,23 @@
 // Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.athenz.identity;
 
+import org.apache.http.Header;
+import org.apache.http.HeaderIterator;
+import org.apache.http.HttpEntity;
+import org.apache.http.HttpHost;
+import org.apache.http.HttpRequest;
+import org.apache.http.ProtocolVersion;
+import org.apache.http.StatusLine;
+import org.apache.http.client.ClientProtocolException;
 import org.apache.http.client.HttpClient;
-import org.apache.http.client.ResponseHandler;
-import org.apache.http.client.methods.HttpUriRequest;
+import org.apache.http.client.methods.CloseableHttpResponse;
 import org.apache.http.impl.client.CloseableHttpClient;
 import org.apache.http.impl.client.HttpClientBuilder;
+import org.apache.http.protocol.HttpContext;
 
 import javax.net.ssl.SSLContext;
 import java.io.IOException;
-import java.io.UncheckedIOException;
+import java.util.Locale;
 import java.util.function.Supplier;
 
 /**
@@ -17,7 +25,16 @@ import java.util.function.Supplier;
  *
  * @author bjorncs
  */
-public class SiaBackedApacheHttpClient implements AutoCloseable {
+public class SiaBackedApacheHttpClient extends CloseableHttpClient {
+
+    /**
+     * A factory that returns a new instance of {@link CloseableHttpClient}.
+     * The implementor is responsible for configuring the {@link SSLContext}, e.g. using {@link HttpClientBuilder#setSSLContext(SSLContext)}.
+     */
+    @FunctionalInterface
+    public interface HttpClientFactory {
+        CloseableHttpClient createHttpClient(SSLContext sslContext);
+    }
 
     private final Object clientLock = new Object();
     private final Supplier<SSLContext> sslContextSupplier;
@@ -39,31 +56,19 @@ public class SiaBackedApacheHttpClient implements AutoCloseable {
                                      HttpClientFactory httpClientFactory) {
         this.sslContextSupplier = sslContextSupplier;
         this.httpClientFactory = httpClientFactory;
-        this.client = new HttpClientHolder(httpClientFactory, sslContextSupplier);
+        this.client = new HttpClientHolder(httpClientFactory, sslContextSupplier, clientLock);
+        this.client.acquire(); // owner ref
     }
 
-    public <T> T execute(HttpUriRequest request, ResponseHandler<T> responseHandler) {
+    @Override
+    protected CloseableHttpResponse doExecute(HttpHost target, HttpRequest request, HttpContext context) throws IOException, ClientProtocolException {
         HttpClientHolder client = getClient();
+        client.acquire(); // request ref
         try {
-            // Note: Apache http client will handle response/connection cleanup for this overload of execute().
-            return client.apacheClient.execute(request, responseHandler);
-        } catch (IOException e) {
-            throw new UncheckedIOException(e);
+            CloseableHttpResponse response = client.apacheClient.execute(target, request, context);
+            return new ResponseHolder(response, client);
         } finally {
-            synchronized (clientLock) {
-                client.release();
-            }
-        }
-    }
-
-    private HttpClientHolder getClient() {
-        synchronized (clientLock) {
-            if (closed) throw new IllegalStateException("Client has been closed!");
-            if (sslContextSupplier.get() != client.sslContext) {
-                client.release();
-                client = new HttpClientHolder(httpClientFactory, sslContextSupplier);
-            }
-            return client.refer();
+            client.release(); // request ref
         }
     }
 
@@ -71,51 +76,114 @@ public class SiaBackedApacheHttpClient implements AutoCloseable {
      * Thread-safe and idempotent.
      */
     @Override
-    public void close() {
+    public void close() throws IOException {
         synchronized (clientLock) {
             if (closed) return;
-            client.release();
+            client.release(); // owner ref
             closed = true;
         }
     }
 
+    @Override
+    @SuppressWarnings("deprecation")
+    public org.apache.http.params.HttpParams getParams() {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
+    @SuppressWarnings("deprecation")
+    public org.apache.http.conn.ClientConnectionManager getConnectionManager() {
+        throw new UnsupportedOperationException();
+    }
+
+    private HttpClientHolder getClient() throws IOException {
+        synchronized (clientLock) {
+            if (closed) throw new IllegalStateException("Client has been closed!");
+            if (sslContextSupplier.get() != client.sslContext) {
+                client.release(); // owner ref
+                client = new HttpClientHolder(httpClientFactory, sslContextSupplier, clientLock);
+                client.acquire(); // owner ref
+            }
+            return client;
+        }
+    }
+
     private static class HttpClientHolder {
-        final Supplier<SSLContext> sslContextSupplier;
+        final Object clientLock;
         final CloseableHttpClient apacheClient;
         final SSLContext sslContext;
-        int referenceCount = 1; // Owner's reference implicitly counted
+        int referenceCount = 0;
+        boolean closed = false;
 
-        HttpClientHolder(HttpClientFactory clientBuilder, Supplier<SSLContext> sslContextSupplier) {
+        HttpClientHolder(HttpClientFactory httpClientFactory, Supplier<SSLContext> sslContextSupplier, Object clientLock) {
             SSLContext sslContext = sslContextSupplier.get();
-            this.sslContextSupplier = sslContextSupplier;
-            this.apacheClient = clientBuilder.createHttpClient(sslContext);
+            this.apacheClient = httpClientFactory.createHttpClient(sslContext);
             this.sslContext = sslContext;
+            this.clientLock = clientLock;
         }
 
-        HttpClientHolder refer() {
-            if (referenceCount == 0) throw new IllegalStateException("Client already closed!");
-            ++referenceCount;
-            return this;
+        void acquire() {
+            synchronized (clientLock) {
+                if (closed) throw new IllegalStateException("Client already closed!");
+                ++referenceCount;
+            }
         }
 
-        void release() {
-            if (referenceCount == 0) throw new IllegalStateException("Client already closed!");
-            --referenceCount;
-            if (referenceCount == 0) {
-                try {
+        void release() throws IOException {
+            synchronized (clientLock) {
+                if (closed) throw new IllegalStateException("Client already closed!");
+                --referenceCount;
+                if (referenceCount == 0) {
                     apacheClient.close();
-                } catch (IOException e) {
-                    throw new UncheckedIOException(e);
+                    closed = true;
                 }
             }
         }
     }
 
-    /**
-     * A factory that returns a new instance of {@link CloseableHttpClient}. The implementor is responsible for configuring the {@link SSLContext}, e.g. using {@link HttpClientBuilder#setSSLContext(SSLContext)}.
-     */
-    @FunctionalInterface
-    public interface HttpClientFactory {
-        CloseableHttpClient createHttpClient(SSLContext sslContext);
+    private static class ResponseHolder implements CloseableHttpResponse {
+        final CloseableHttpResponse response;
+        final HttpClientHolder clientHolder;
+
+        ResponseHolder(CloseableHttpResponse response, HttpClientHolder clientHolder) {
+            clientHolder.acquire(); // response ref
+            this.response = response;
+            this.clientHolder = clientHolder;
+        }
+
+        @Override
+        public void close() throws IOException {
+            response.close(); // response ref
+            clientHolder.release();
+        }
+
+        // Proxy methods
+        @Override public StatusLine getStatusLine() { return response.getStatusLine(); }
+        @Override public void setStatusLine(StatusLine statusline) { response.setStatusLine(statusline); }
+        @Override public void setStatusLine(ProtocolVersion ver, int code) { response.setStatusLine(ver, code); }
+        @Override public void setStatusLine(ProtocolVersion ver, int code, String reason) { response.setStatusLine(ver, code, reason); }
+        @Override public void setStatusCode(int code) throws IllegalStateException { response.setStatusCode(code);}
+        @Override public void setReasonPhrase(String reason) throws IllegalStateException { response.setReasonPhrase(reason); }
+        @Override public HttpEntity getEntity() { return response.getEntity(); }
+        @Override public void setEntity(HttpEntity entity) { response.setEntity(entity); }
+        @Override public Locale getLocale() { return response.getLocale(); }
+        @Override public void setLocale(Locale loc) { response.setLocale(loc); }
+        @Override public ProtocolVersion getProtocolVersion() { return response.getProtocolVersion(); }
+        @Override public boolean containsHeader(String name) { return response.containsHeader(name); }
+        @Override public Header[] getHeaders(String name) { return response.getHeaders(name); }
+        @Override public Header getFirstHeader(String name) { return response.getFirstHeader(name); }
+        @Override public Header getLastHeader(String name) { return response.getLastHeader(name); }
+        @Override public Header[] getAllHeaders() { return response.getAllHeaders(); }
+        @Override public void addHeader(Header header) { response.addHeader(header); }
+        @Override public void addHeader(String name, String value) { response.addHeader(name, value); }
+        @Override public void setHeader(Header header) { response.setHeader(header); }
+        @Override public void setHeader(String name, String value) { response.setHeader(name, value); }
+        @Override public void setHeaders(Header[] headers) { response.setHeaders(headers); }
+        @Override public void removeHeader(Header header) { response.removeHeader(header); }
+        @Override public void removeHeaders(String name) { response.removeHeaders(name); }
+        @Override public HeaderIterator headerIterator() { return response.headerIterator(); }
+        @Override public HeaderIterator headerIterator(String name) { return response.headerIterator(name); }
+        @Override @SuppressWarnings("deprecation") public org.apache.http.params.HttpParams getParams() { return response.getParams(); }
+        @Override @SuppressWarnings("deprecation") public void setParams(org.apache.http.params.HttpParams params) { response.setParams(params); }
     }
 }