Merge pull request #19373 from javanna/enhancement/rest_client_builder_callback

Rest Client: add callback to customize http client settings

client/rest/src/main/java/org/elasticsearch/client/RestClient.java

@@ -37,8 +37,6 @@ import org.apache.http.client.methods.HttpPut;
 import org.apache.http.client.methods.HttpRequestBase;
 import org.apache.http.client.methods.HttpTrace;
 import org.apache.http.client.utils.URIBuilder;
-import org.apache.http.config.Registry;
-import org.apache.http.conn.socket.ConnectionSocketFactory;
 import org.apache.http.entity.ContentType;
 import org.apache.http.impl.client.CloseableHttpClient;
 import org.apache.http.impl.client.HttpClientBuilder;
@@ -91,7 +89,7 @@ public final class RestClient implements Closeable {
     private final ConcurrentMap<HttpHost, DeadHostState> blacklist = new ConcurrentHashMap<>();
     private final FailureListener failureListener;
 
-    private RestClient(CloseableHttpClient client, long maxRetryTimeoutMillis, Header[] defaultHeaders,
+    RestClient(CloseableHttpClient client, long maxRetryTimeoutMillis, Header[] defaultHeaders,
                        HttpHost[] hosts, FailureListener failureListener) {
         this.client = client;
         this.maxRetryTimeoutMillis = maxRetryTimeoutMillis;
@@ -393,10 +391,11 @@ public final class RestClient implements Closeable {
         private static final Header[] EMPTY_HEADERS = new Header[0];
 
         private final HttpHost[] hosts;
-        private CloseableHttpClient httpClient;
         private int maxRetryTimeout = DEFAULT_MAX_RETRY_TIMEOUT_MILLIS;
         private Header[] defaultHeaders = EMPTY_HEADERS;
         private FailureListener failureListener;
+        private HttpClientConfigCallback httpClientConfigCallback;
+        private RequestConfigCallback requestConfigCallback;
 
         /**
          * Creates a new builder instance and sets the hosts that the client will send requests to.
@@ -408,17 +407,6 @@ public final class RestClient implements Closeable {
             this.hosts = hosts;
         }
 
-        /**
-         * Sets the http client. A new default one will be created if not
-         * specified, by calling {@link #createDefaultHttpClient(Registry)})}.
-         *
-         * @see CloseableHttpClient
-         */
-        public Builder setHttpClient(CloseableHttpClient httpClient) {
-            this.httpClient = httpClient;
-            return this;
-        }
-
         /**
          * Sets the maximum timeout (in milliseconds) to honour in case of multiple retries of the same request.
          * {@link #DEFAULT_MAX_RETRY_TIMEOUT_MILLIS} if not specified.
@@ -434,12 +422,10 @@ public final class RestClient implements Closeable {
         }
 
         /**
-         * Sets the default request headers, to be used when creating the default http client instance.
-         * In case the http client is set through {@link #setHttpClient(CloseableHttpClient)}, the default headers need to be
-         * set to it externally during http client construction.
+         * Sets the default request headers, to be used sent with every request unless overridden on a per request basis
          */
         public Builder setDefaultHeaders(Header[] defaultHeaders) {
-            Objects.requireNonNull(defaultHeaders, "default headers must not be null");
+            Objects.requireNonNull(defaultHeaders, "defaultHeaders must not be null");
             for (Header defaultHeader : defaultHeaders) {
                 Objects.requireNonNull(defaultHeader, "default header must not be null");
             }
@@ -451,48 +437,94 @@ public final class RestClient implements Closeable {
          * Sets the {@link FailureListener} to be notified for each request failure
          */
         public Builder setFailureListener(FailureListener failureListener) {
-            Objects.requireNonNull(failureListener, "failure listener must not be null");
+            Objects.requireNonNull(failureListener, "failureListener must not be null");
             this.failureListener = failureListener;
             return this;
         }
 
+        /**
+         * Sets the {@link HttpClientConfigCallback} to be used to customize http client configuration
+         */
+        public Builder setHttpClientConfigCallback(HttpClientConfigCallback httpClientConfigCallback) {
+            Objects.requireNonNull(httpClientConfigCallback, "httpClientConfigCallback must not be null");
+            this.httpClientConfigCallback = httpClientConfigCallback;
+            return this;
+        }
+
+        /**
+         * Sets the {@link RequestConfigCallback} to be used to customize http client configuration
+         */
+        public Builder setRequestConfigCallback(RequestConfigCallback requestConfigCallback) {
+            Objects.requireNonNull(requestConfigCallback, "requestConfigCallback must not be null");
+            this.requestConfigCallback = requestConfigCallback;
+            return this;
+        }
+
         /**
          * Creates a new {@link RestClient} based on the provided configuration.
          */
         public RestClient build() {
-            if (httpClient == null) {
-                httpClient = createDefaultHttpClient(null);
-            }
             if (failureListener == null) {
                 failureListener = new FailureListener();
             }
+            CloseableHttpClient httpClient = createHttpClient();
             return new RestClient(httpClient, maxRetryTimeout, defaultHeaders, hosts, failureListener);
         }
 
-        /**
-         * Creates a {@link CloseableHttpClient} with default settings. Used when the http client instance is not provided.
-         *
-         * @see CloseableHttpClient
-         */
-        public static CloseableHttpClient createDefaultHttpClient(Registry<ConnectionSocketFactory> socketFactoryRegistry) {
-            PoolingHttpClientConnectionManager connectionManager;
-            if (socketFactoryRegistry == null) {
-                connectionManager = new PoolingHttpClientConnectionManager();
-            } else {
-                connectionManager = new PoolingHttpClientConnectionManager(socketFactoryRegistry);
+        private CloseableHttpClient createHttpClient() {
+            //default timeouts are all infinite
+            RequestConfig.Builder requestConfigBuilder = RequestConfig.custom().setConnectTimeout(DEFAULT_CONNECT_TIMEOUT_MILLIS)
+                    .setSocketTimeout(DEFAULT_SOCKET_TIMEOUT_MILLIS)
+                    .setConnectionRequestTimeout(DEFAULT_CONNECTION_REQUEST_TIMEOUT_MILLIS);
+
+            if (requestConfigCallback != null) {
+                requestConfigCallback.customizeRequestConfig(requestConfigBuilder);
             }
+            RequestConfig requestConfig = requestConfigBuilder.build();
+
+            PoolingHttpClientConnectionManager connectionManager = new PoolingHttpClientConnectionManager();
             //default settings may be too constraining
             connectionManager.setDefaultMaxPerRoute(10);
             connectionManager.setMaxTotal(30);
 
-            //default timeouts are all infinite
-            RequestConfig requestConfig = RequestConfig.custom().setConnectTimeout(DEFAULT_CONNECT_TIMEOUT_MILLIS)
-                    .setSocketTimeout(DEFAULT_SOCKET_TIMEOUT_MILLIS)
-                    .setConnectionRequestTimeout(DEFAULT_CONNECTION_REQUEST_TIMEOUT_MILLIS).build();
-            return HttpClientBuilder.create().setConnectionManager(connectionManager).setDefaultRequestConfig(requestConfig).build();
+            HttpClientBuilder httpClientBuilder = HttpClientBuilder.create().setConnectionManager(connectionManager)
+                    .setDefaultRequestConfig(requestConfig);
+
+            if (httpClientConfigCallback != null) {
+                httpClientConfigCallback.customizeHttpClient(httpClientBuilder);
+            }
+            return httpClientBuilder.build();
         }
     }
 
+    /**
+     * Callback used the default {@link RequestConfig} being set to the {@link CloseableHttpClient}
+     * @see HttpClientBuilder#setDefaultRequestConfig
+     */
+    public interface RequestConfigCallback {
+        /**
+         * Allows to customize the {@link RequestConfig} that will be used with each request.
+         * It is common to customize the different timeout values through this method without losing any other useful default
+         * value that the {@link RestClient.Builder} internally sets.
+         */
+        void customizeRequestConfig(RequestConfig.Builder requestConfigBuilder);
+    }
+
+    /**
+     * Callback used to customize the {@link CloseableHttpClient} instance used by a {@link RestClient} instance.
+     * Allows to customize default {@link RequestConfig} being set to the client and any parameter that
+     * can be set through {@link HttpClientBuilder}
+     */
+    public interface HttpClientConfigCallback {
+        /**
+         * Allows to customize the {@link CloseableHttpClient} being created and used by the {@link RestClient}.
+         * It is common to customzie the default {@link org.apache.http.client.CredentialsProvider} through this method,
+         * without losing any other useful default value that the {@link RestClient.Builder} internally sets.
+         * Also useful to setup ssl through {@link SSLSocketFactoryHttpConfigCallback}.
+         */
+        void customizeHttpClient(HttpClientBuilder httpClientBuilder);
+    }
+
     /**
      * Listener that allows to be notified whenever a failure happens. Useful when sniffing is enabled, so that we can sniff on failure.
      * The default implementation is a no-op.

client/rest/src/main/java/org/elasticsearch/client/SSLSocketFactoryHttpConfigCallback.java

@@ -0,0 +1,53 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.client;
+
+import org.apache.http.config.Registry;
+import org.apache.http.config.RegistryBuilder;
+import org.apache.http.conn.socket.ConnectionSocketFactory;
+import org.apache.http.conn.socket.PlainConnectionSocketFactory;
+import org.apache.http.conn.ssl.SSLConnectionSocketFactory;
+import org.apache.http.impl.client.HttpClientBuilder;
+import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
+
+/**
+ * Helps configuring the http client when needing to communicate over ssl. It effectively replaces the connection manager
+ * with one that has ssl properly configured thanks to the provided {@link SSLConnectionSocketFactory}.
+ */
+public class SSLSocketFactoryHttpConfigCallback implements RestClient.HttpClientConfigCallback {
+
+    private final SSLConnectionSocketFactory sslSocketFactory;
+
+    public SSLSocketFactoryHttpConfigCallback(SSLConnectionSocketFactory sslSocketFactory) {
+        this.sslSocketFactory = sslSocketFactory;
+    }
+
+    @Override
+    public void customizeHttpClient(HttpClientBuilder httpClientBuilder) {
+        Registry<ConnectionSocketFactory> socketFactoryRegistry = RegistryBuilder.<ConnectionSocketFactory>create()
+                .register("http", PlainConnectionSocketFactory.getSocketFactory())
+                .register("https", sslSocketFactory).build();
+        PoolingHttpClientConnectionManager connectionManager = new PoolingHttpClientConnectionManager(socketFactoryRegistry);
+        //default settings may be too constraining
+        connectionManager.setDefaultMaxPerRoute(10);
+        connectionManager.setMaxTotal(30);
+        httpClientBuilder.setConnectionManager(connectionManager);
+    }
+}

client/rest/src/test/java/org/elasticsearch/client/RestClientBuilderTests.java

@@ -22,6 +22,7 @@ package org.elasticsearch.client;
 import com.carrotsearch.randomizedtesting.generators.RandomInts;
 import org.apache.http.Header;
 import org.apache.http.HttpHost;
+import org.apache.http.client.config.RequestConfig;
 import org.apache.http.impl.client.HttpClientBuilder;
 import org.apache.http.message.BasicHeader;
 
@@ -67,7 +68,7 @@ public class RestClientBuilderTests extends RestClientTestCase {
             RestClient.builder(new HttpHost("localhost", 9200)).setDefaultHeaders(null);
             fail("should have failed");
         } catch(NullPointerException e) {
-            assertEquals("default headers must not be null", e.getMessage());
+            assertEquals("defaultHeaders must not be null", e.getMessage());
         }
 
         try {
@@ -81,7 +82,21 @@ public class RestClientBuilderTests extends RestClientTestCase {
             RestClient.builder(new HttpHost("localhost", 9200)).setFailureListener(null);
             fail("should have failed");
         } catch(NullPointerException e) {
-            assertEquals("failure listener must not be null", e.getMessage());
+            assertEquals("failureListener must not be null", e.getMessage());
+        }
+
+        try {
+            RestClient.builder(new HttpHost("localhost", 9200)).setHttpClientConfigCallback(null);
+            fail("should have failed");
+        } catch(NullPointerException e) {
+            assertEquals("httpClientConfigCallback must not be null", e.getMessage());
+        }
+
+        try {
+            RestClient.builder(new HttpHost("localhost", 9200)).setRequestConfigCallback(null);
+            fail("should have failed");
+        } catch(NullPointerException e) {
+            assertEquals("requestConfigCallback must not be null", e.getMessage());
         }
 
         int numNodes = RandomInts.randomIntBetween(getRandom(), 1, 5);
@@ -91,7 +106,18 @@ public class RestClientBuilderTests extends RestClientTestCase {
         }
         RestClient.Builder builder = RestClient.builder(hosts);
         if (getRandom().nextBoolean()) {
-            builder.setHttpClient(HttpClientBuilder.create().build());
+            builder.setHttpClientConfigCallback(new RestClient.HttpClientConfigCallback() {
+                @Override
+                public void customizeHttpClient(HttpClientBuilder httpClientBuilder) {
+                }
+            });
+        }
+        if (getRandom().nextBoolean()) {
+            builder.setRequestConfigCallback(new RestClient.RequestConfigCallback() {
+                @Override
+                public void customizeRequestConfig(RequestConfig.Builder requestConfigBuilder) {
+                }
+            });
         }
         if (getRandom().nextBoolean()) {
             int numHeaders = RandomInts.randomIntBetween(getRandom(), 1, 5);

client/rest/src/test/java/org/elasticsearch/client/RestClientMultipleHostsTests.java

@@ -20,6 +20,7 @@
 package org.elasticsearch.client;
 
 import com.carrotsearch.randomizedtesting.generators.RandomInts;
+import org.apache.http.Header;
 import org.apache.http.HttpHost;
 import org.apache.http.HttpRequest;
 import org.apache.http.ProtocolVersion;
@@ -91,7 +92,7 @@ public class RestClientMultipleHostsTests extends RestClientTestCase {
             httpHosts[i] = new HttpHost("localhost", 9200 + i);
         }
         failureListener = new TrackingFailureListener();
-        restClient = RestClient.builder(httpHosts).setHttpClient(httpClient).setFailureListener(failureListener).build();
+        restClient = new RestClient(httpClient, 10000, new Header[0], httpHosts, failureListener);
     }
 
     public void testRoundRobinOkStatusCodes() throws Exception {

client/rest/src/test/java/org/elasticsearch/client/RestClientSingleHostTests.java

@@ -128,8 +128,7 @@ public class RestClientSingleHostTests extends RestClientTestCase {
         }
         httpHost = new HttpHost("localhost", 9200);
         failureListener = new TrackingFailureListener();
-        restClient = RestClient.builder(httpHost).setHttpClient(httpClient).setDefaultHeaders(defaultHeaders)
-                .setFailureListener(failureListener).build();
+        restClient = new RestClient(httpClient, 10000, defaultHeaders, new HttpHost[]{httpHost}, failureListener);
     }
 
     /**

qa/smoke-test-http/src/test/java/org/elasticsearch/http/HttpCompressionIT.java

@@ -23,8 +23,7 @@ import org.apache.http.HttpException;
 import org.apache.http.HttpHeaders;
 import org.apache.http.HttpResponseInterceptor;
 import org.apache.http.entity.StringEntity;
-import org.apache.http.impl.client.CloseableHttpClient;
-import org.apache.http.impl.client.HttpClients;
+import org.apache.http.impl.client.HttpClientBuilder;
 import org.apache.http.message.BasicHeader;
 import org.apache.http.protocol.HttpContext;
 import org.elasticsearch.client.Response;
@@ -50,7 +49,7 @@ public class HttpCompressionIT extends ESIntegTestCase {
         ensureGreen();
         // we need to intercept early, otherwise internal logic in HttpClient will just remove the header and we cannot verify it
         ContentEncodingHeaderExtractor headerExtractor = new ContentEncodingHeaderExtractor();
-        try (RestClient client = createRestClient(HttpClients.custom().addInterceptorFirst(headerExtractor).build())) {
+        try (RestClient client = createRestClient(new ContentEncodingHeaderExtractorConfigCallback(headerExtractor))) {
             try (Response response = client.performRequest("GET", "/", new BasicHeader(HttpHeaders.ACCEPT_ENCODING, GZIP_ENCODING))) {
                 assertEquals(200, response.getStatusLine().getStatusCode());
                 assertTrue(headerExtractor.hasContentEncodingHeader());
@@ -62,8 +61,7 @@ public class HttpCompressionIT extends ESIntegTestCase {
     public void testUncompressedResponseByDefault() throws Exception {
         ensureGreen();
         ContentEncodingHeaderExtractor headerExtractor = new ContentEncodingHeaderExtractor();
-        CloseableHttpClient httpClient = HttpClients.custom().disableContentCompression().addInterceptorFirst(headerExtractor).build();
-        try (RestClient client = createRestClient(httpClient)) {
+        try (RestClient client = createRestClient(new NoContentCompressionConfigCallback(headerExtractor))) {
             try (Response response = client.performRequest("GET", "/")) {
                 assertEquals(200, response.getStatusLine().getStatusCode());
                 assertFalse(headerExtractor.hasContentEncodingHeader());
@@ -75,8 +73,7 @@ public class HttpCompressionIT extends ESIntegTestCase {
         ensureGreen();
         ContentEncodingHeaderExtractor headerExtractor = new ContentEncodingHeaderExtractor();
         // this disable content compression in both directions (request and response)
-        CloseableHttpClient httpClient = HttpClients.custom().disableContentCompression().addInterceptorFirst(headerExtractor).build();
-        try (RestClient client = createRestClient(httpClient)) {
+        try (RestClient client = createRestClient(new NoContentCompressionConfigCallback(headerExtractor))) {
             try (Response response = client.performRequest("POST", "/company/employees/1",
                     Collections.emptyMap(), SAMPLE_DOCUMENT)) {
                 assertEquals(201, response.getStatusLine().getStatusCode());
@@ -89,7 +86,7 @@ public class HttpCompressionIT extends ESIntegTestCase {
         ensureGreen();
         ContentEncodingHeaderExtractor headerExtractor = new ContentEncodingHeaderExtractor();
         // we don't call #disableContentCompression() hence the client will send the content compressed
-        try (RestClient client = createRestClient(HttpClients.custom().addInterceptorFirst(headerExtractor).build())) {
+        try (RestClient client = createRestClient(new ContentEncodingHeaderExtractorConfigCallback(headerExtractor))) {
             try (Response response = client.performRequest("POST", "/company/employees/2",
                     Collections.emptyMap(), SAMPLE_DOCUMENT)) {
                 assertEquals(201, response.getStatusLine().getStatusCode());
@@ -119,4 +116,30 @@ public class HttpCompressionIT extends ESIntegTestCase {
             return contentEncodingHeader;
         }
     }
+
+    private static class NoContentCompressionConfigCallback extends ContentEncodingHeaderExtractorConfigCallback {
+        NoContentCompressionConfigCallback(ContentEncodingHeaderExtractor contentEncodingHeaderExtractor) {
+            super(contentEncodingHeaderExtractor);
+        }
+
+        @Override
+        public void customizeHttpClient(HttpClientBuilder httpClientBuilder) {
+            super.customizeHttpClient(httpClientBuilder);
+            httpClientBuilder.disableContentCompression();
+        }
+    }
+
+    private static class ContentEncodingHeaderExtractorConfigCallback implements RestClient.HttpClientConfigCallback {
+
+        private final ContentEncodingHeaderExtractor contentEncodingHeaderExtractor;
+
+        ContentEncodingHeaderExtractorConfigCallback(ContentEncodingHeaderExtractor contentEncodingHeaderExtractor) {
+            this.contentEncodingHeaderExtractor = contentEncodingHeaderExtractor;
+        }
+
+        @Override
+        public void customizeHttpClient(HttpClientBuilder httpClientBuilder) {
+            httpClientBuilder.addInterceptorFirst(contentEncodingHeaderExtractor);
+        }
+    }
 }

test/framework/src/main/java/org/elasticsearch/test/ESIntegTestCase.java

@@ -23,7 +23,6 @@ import com.carrotsearch.randomizedtesting.annotations.TestGroup;
 import com.carrotsearch.randomizedtesting.generators.RandomInts;
 import com.carrotsearch.randomizedtesting.generators.RandomPicks;
 import org.apache.http.HttpHost;
-import org.apache.http.impl.client.CloseableHttpClient;
 import org.apache.lucene.util.IOUtils;
 import org.apache.lucene.util.LuceneTestCase;
 import org.apache.lucene.util.TestUtil;
@@ -2075,11 +2074,11 @@ public abstract class ESIntegTestCase extends ESTestCase {
         return restClient;
     }
 
-    protected static RestClient createRestClient(CloseableHttpClient httpClient) {
-        return createRestClient(httpClient, "http");
+    protected static RestClient createRestClient(RestClient.HttpClientConfigCallback httpClientConfigCallback) {
+        return createRestClient(httpClientConfigCallback, "http");
     }
 
-    protected static RestClient createRestClient(CloseableHttpClient httpClient, String protocol) {
+    protected static RestClient createRestClient(RestClient.HttpClientConfigCallback httpClientConfigCallback, String protocol) {
         final NodesInfoResponse nodeInfos = client().admin().cluster().prepareNodesInfo().get();
         final List<NodeInfo> nodes = nodeInfos.getNodes();
         assertFalse(nodeInfos.hasFailures());
@@ -2093,8 +2092,8 @@ public abstract class ESIntegTestCase extends ESTestCase {
             }
         }
         RestClient.Builder builder = RestClient.builder(hosts.toArray(new HttpHost[hosts.size()]));
-        if (httpClient != null) {
-            builder.setHttpClient(httpClient);
+        if (httpClientConfigCallback != null) {
+            builder.setHttpClientConfigCallback(httpClientConfigCallback);
         }
         return builder.build();
     }

test/framework/src/main/java/org/elasticsearch/test/rest/client/RestTestClient.java

@@ -22,16 +22,8 @@ import com.carrotsearch.randomizedtesting.RandomizedTest;
 import org.apache.http.Header;
 import org.apache.http.HttpEntity;
 import org.apache.http.HttpHost;
-import org.apache.http.client.config.RequestConfig;
-import org.apache.http.config.Registry;
-import org.apache.http.config.RegistryBuilder;
-import org.apache.http.conn.socket.ConnectionSocketFactory;
-import org.apache.http.conn.socket.PlainConnectionSocketFactory;
 import org.apache.http.conn.ssl.SSLConnectionSocketFactory;
 import org.apache.http.entity.StringEntity;
-import org.apache.http.impl.client.CloseableHttpClient;
-import org.apache.http.impl.client.HttpClientBuilder;
-import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
 import org.apache.http.message.BasicHeader;
 import org.apache.http.ssl.SSLContexts;
 import org.apache.lucene.util.IOUtils;
@@ -39,6 +31,7 @@ import org.elasticsearch.Version;
 import org.elasticsearch.client.Response;
 import org.elasticsearch.client.ResponseException;
 import org.elasticsearch.client.RestClient;
+import org.elasticsearch.client.SSLSocketFactoryHttpConfigCallback;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.PathUtils;
 import org.elasticsearch.common.logging.ESLogger;
@@ -82,11 +75,6 @@ public class RestTestClient implements Closeable {
     public static final String TRUSTSTORE_PATH = "truststore.path";
     public static final String TRUSTSTORE_PASSWORD = "truststore.password";
 
-    public static final int CONNECT_TIMEOUT_MILLIS = 1000;
-    public static final int SOCKET_TIMEOUT_MILLIS = 30000;
-    public static final int MAX_RETRY_TIMEOUT_MILLIS = SOCKET_TIMEOUT_MILLIS;
-    public static final int CONNECTION_REQUEST_TIMEOUT_MILLIS = 500;
-
     private static final ESLogger logger = Loggers.getLogger(RestTestClient.class);
     //query_string params that don't need to be declared in the spec, thay are supported by default
     private static final Set<String> ALWAYS_ACCEPTED_QUERY_STRING_PARAMS = Sets.newHashSet("pretty", "source", "filter_path");
@@ -274,7 +262,15 @@ public class RestTestClient implements Closeable {
     }
 
     private static RestClient createRestClient(URL[] urls, Settings settings) throws IOException {
-        SSLConnectionSocketFactory sslsf;
+        String protocol = settings.get(PROTOCOL, "http");
+        HttpHost[] hosts = new HttpHost[urls.length];
+        for (int i = 0; i < hosts.length; i++) {
+            URL url = urls[i];
+            hosts[i] = new HttpHost(url.getHost(), url.getPort(), protocol);
+        }
+        RestClient.Builder builder = RestClient.builder(hosts).setMaxRetryTimeoutMillis(30000)
+                .setRequestConfigCallback(requestConfigBuilder -> requestConfigBuilder.setSocketTimeout(30000));
+
         String keystorePath = settings.get(TRUSTSTORE_PATH);
         if (keystorePath != null) {
             final String keystorePass = settings.get(TRUSTSTORE_PASSWORD);
@@ -291,38 +287,13 @@ public class RestTestClient implements Closeable {
                     keyStore.load(is, keystorePass.toCharArray());
                 }
                 SSLContext sslcontext = SSLContexts.custom().loadTrustMaterial(keyStore, null).build();
-                sslsf = new SSLConnectionSocketFactory(sslcontext);
+                SSLConnectionSocketFactory sslConnectionSocketFactory = new SSLConnectionSocketFactory(sslcontext);
+                builder.setHttpClientConfigCallback(new SSLSocketFactoryHttpConfigCallback(sslConnectionSocketFactory));
             } catch (KeyStoreException|NoSuchAlgorithmException|KeyManagementException|CertificateException e) {
                 throw new RuntimeException(e);
             }
-        } else {
-            sslsf = SSLConnectionSocketFactory.getSocketFactory();
-        }
-
-        Registry<ConnectionSocketFactory> socketFactoryRegistry = RegistryBuilder.<ConnectionSocketFactory>create()
-                .register("http", PlainConnectionSocketFactory.getSocketFactory())
-                .register("https", sslsf)
-                .build();
-        PoolingHttpClientConnectionManager connectionManager = new PoolingHttpClientConnectionManager(socketFactoryRegistry);
-        //default settings may be too constraining
-        connectionManager.setDefaultMaxPerRoute(10);
-        connectionManager.setMaxTotal(30);
-
-        //default timeouts are all infinite
-        RequestConfig requestConfig = RequestConfig.custom().setConnectTimeout(CONNECT_TIMEOUT_MILLIS)
-            .setSocketTimeout(SOCKET_TIMEOUT_MILLIS)
-            .setConnectionRequestTimeout(CONNECTION_REQUEST_TIMEOUT_MILLIS).build();
-        CloseableHttpClient httpClient = HttpClientBuilder.create()
-            .setConnectionManager(connectionManager).setDefaultRequestConfig(requestConfig).build();
-
-        String protocol = settings.get(PROTOCOL, "http");
-        HttpHost[] hosts = new HttpHost[urls.length];
-        for (int i = 0; i < hosts.length; i++) {
-            URL url = urls[i];
-            hosts[i] = new HttpHost(url.getHost(), url.getPort(), protocol);
         }
 
-        RestClient.Builder builder = RestClient.builder(hosts).setHttpClient(httpClient).setMaxRetryTimeoutMillis(MAX_RETRY_TIMEOUT_MILLIS);
         try (ThreadContext threadContext = new ThreadContext(settings)) {
             Header[] defaultHeaders = new Header[threadContext.getHeaders().size()];
             int i = 0;