Rest Client: add callback to customize http client settings

The callback replaces the ability to fully replace the http client instance. By doing that, one used to lose any default that the RestClient had set for the underlying http client. Given that you'd usually override one or two things only, like a couple of timeout values, the ssl factory or the default credentials providers, it is not uder friendly if by doing that users end up replacing the whole http client instance and lose any default set by us.

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,10 @@ 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;
 
         /**
          * Creates a new builder instance and sets the hosts that the client will send requests to.
@@ -408,17 +406,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 +421,10 @@ public final class RestClient implements Closeable {
         }
 
         /**
-         * Sets the default request headers, to be used when creating the default http client instance.
+         * Sets the default request headers, to be used sent with every request unless overridden on a per request basis
-         * 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.
          */
         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 +436,79 @@ 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;
+        }
+
         /**
          * 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);
         }
 
-        /**
+        private CloseableHttpClient createHttpClient() {
-         * Creates a {@link CloseableHttpClient} with default settings. Used when the http client instance is not provided.
+            //default timeouts are all infinite
-         *
+            RequestConfig.Builder requestConfigBuilder = RequestConfig.custom().setConnectTimeout(DEFAULT_CONNECT_TIMEOUT_MILLIS)
-         * @see CloseableHttpClient
+                    .setSocketTimeout(DEFAULT_SOCKET_TIMEOUT_MILLIS)
-         */
+                    .setConnectionRequestTimeout(DEFAULT_CONNECTION_REQUEST_TIMEOUT_MILLIS);
-        public static CloseableHttpClient createDefaultHttpClient(Registry<ConnectionSocketFactory> socketFactoryRegistry) {
+
-            PoolingHttpClientConnectionManager connectionManager;
+            if (httpClientConfigCallback != null) {
-            if (socketFactoryRegistry == null) {
+                httpClientConfigCallback.customizeDefaultRequestConfig(requestConfigBuilder);
-                connectionManager = new PoolingHttpClientConnectionManager();
-            } else {
-                connectionManager = new PoolingHttpClientConnectionManager(socketFactoryRegistry);
             }
+            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
+            HttpClientBuilder httpClientBuilder = HttpClientBuilder.create().setConnectionManager(connectionManager)
-            RequestConfig requestConfig = RequestConfig.custom().setConnectTimeout(DEFAULT_CONNECT_TIMEOUT_MILLIS)
+                    .setDefaultRequestConfig(requestConfig);
-                    .setSocketTimeout(DEFAULT_SOCKET_TIMEOUT_MILLIS)
+
-                    .setConnectionRequestTimeout(DEFAULT_CONNECTION_REQUEST_TIMEOUT_MILLIS).build();
+            if (httpClientConfigCallback != null) {
-            return HttpClientBuilder.create().setConnectionManager(connectionManager).setDefaultRequestConfig(requestConfig).build();
+                httpClientConfigCallback.customizeHttpClient(httpClientBuilder);
+            }
+            return httpClientBuilder.build();
         }
     }
 
+    /**
+     * 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 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} sets internally.
+         */
+        void customizeDefaultRequestConfig(RequestConfig.Builder requestConfigBuilder);
+
+        /**
+         * 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,
+         * or the {@link org.apache.http.conn.HttpClientConnectionManager} being used without losing any other useful default
+         * value that the {@link RestClient.Builder} sets internally.
+         */
+        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/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,14 @@ 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());
         }
 
         int numNodes = RandomInts.randomIntBetween(getRandom(), 1, 5);
@@ -91,7 +99,17 @@ 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 customizeDefaultRequestConfig(RequestConfig.Builder requestConfigBuilder) {
+
+                }
+
+                @Override
+                public void customizeHttpClient(HttpClientBuilder httpClientBuilder) {
+
+                }
+            });
         }
         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)
+        restClient = new RestClient(httpClient, 10000, defaultHeaders, new HttpHost[]{httpHost}, failureListener);
-                .setFailureListener(failureListener).build();
     }
 
     /**

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

@@ -22,9 +22,9 @@ import org.apache.http.Header;
 import org.apache.http.HttpException;
 import org.apache.http.HttpHeaders;
 import org.apache.http.HttpResponseInterceptor;
+import org.apache.http.client.config.RequestConfig;
 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.client.HttpClients;
 import org.apache.http.message.BasicHeader;
 import org.apache.http.protocol.HttpContext;
 import org.elasticsearch.client.Response;
@@ -50,7 +50,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 +62,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(new NoContentCompressionConfigCallback(headerExtractor))) {
-        try (RestClient client = createRestClient(httpClient)) {
             try (Response response = client.performRequest("GET", "/")) {
                 assertEquals(200, response.getStatusLine().getStatusCode());
                 assertFalse(headerExtractor.hasContentEncodingHeader());
@@ -75,8 +74,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(new NoContentCompressionConfigCallback(headerExtractor))) {
-        try (RestClient client = createRestClient(httpClient)) {
             try (Response response = client.performRequest("POST", "/company/employees/1",
                     Collections.emptyMap(), SAMPLE_DOCUMENT)) {
                 assertEquals(201, response.getStatusLine().getStatusCode());
@@ -89,7 +87,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 +117,34 @@ 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 customizeDefaultRequestConfig(RequestConfig.Builder requestConfigBuilder) {
+        }
+
+        @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) {
+    protected static RestClient createRestClient(RestClient.HttpClientConfigCallback httpClientConfigCallback) {
-        return createRestClient(httpClient, "http");
+        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) {
+        if (httpClientConfigCallback != null) {
-            builder.setHttpClient(httpClient);
+            builder.setHttpClientConfigCallback(httpClientConfigCallback);
         }
         return builder.build();
     }

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

@@ -26,10 +26,8 @@ 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;
@@ -82,11 +80,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 +267,7 @@ public class RestTestClient implements Closeable {
     }
 
     private static RestClient createRestClient(URL[] urls, Settings settings) throws IOException {
-        SSLConnectionSocketFactory sslsf;
+        PoolingHttpClientConnectionManager connectionManager;
         String keystorePath = settings.get(TRUSTSTORE_PATH);
         if (keystorePath != null) {
             final String keystorePass = settings.get(TRUSTSTORE_PASSWORD);
@@ -291,30 +284,19 @@ public class RestTestClient implements Closeable {
                     keyStore.load(is, keystorePass.toCharArray());
                 }
                 SSLContext sslcontext = SSLContexts.custom().loadTrustMaterial(keyStore, null).build();
-                sslsf = new SSLConnectionSocketFactory(sslcontext);
+                Registry<ConnectionSocketFactory> socketFactoryRegistry = RegistryBuilder.<ConnectionSocketFactory>create()
+                        .register("https", new SSLConnectionSocketFactory(sslcontext)).build();
+                connectionManager = new PoolingHttpClientConnectionManager(socketFactoryRegistry);
             } catch (KeyStoreException|NoSuchAlgorithmException|KeyManagementException|CertificateException e) {
                 throw new RuntimeException(e);
             }
         } else {
-            sslsf = SSLConnectionSocketFactory.getSocketFactory();
+            connectionManager = new PoolingHttpClientConnectionManager();
         }
-
-        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++) {
@@ -322,7 +304,18 @@ public class RestTestClient implements Closeable {
             hosts[i] = new HttpHost(url.getHost(), url.getPort(), protocol);
         }
 
-        RestClient.Builder builder = RestClient.builder(hosts).setHttpClient(httpClient).setMaxRetryTimeoutMillis(MAX_RETRY_TIMEOUT_MILLIS);
+        RestClient.Builder builder = RestClient.builder(hosts).setMaxRetryTimeoutMillis(30000)
+                .setHttpClientConfigCallback(new RestClient.HttpClientConfigCallback() {
+            @Override
+            public void customizeDefaultRequestConfig(RequestConfig.Builder requestConfigBuilder) {
+                requestConfigBuilder.setSocketTimeout(30000);
+            }
+
+            @Override
+            public void customizeHttpClient(HttpClientBuilder httpClientBuilder) {
+                httpClientBuilder.setConnectionManager(connectionManager);
+            }
+        });
         try (ThreadContext threadContext = new ThreadContext(settings)) {
             Header[] defaultHeaders = new Header[threadContext.getHeaders().size()];
             int i = 0;