Allow to pass socket facttry registry to createDefaultHttpClient method

client-sniffer/src/main/java/org/elasticsearch/client/sniff/Sniffer.java

@@ -22,6 +22,7 @@ package org.elasticsearch.client.sniff;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.http.HttpHost;
+import org.apache.http.config.Registry;
 import org.apache.http.impl.client.CloseableHttpClient;
 import org.elasticsearch.client.RestClient;
 
@@ -203,7 +204,7 @@ public final class Sniffer extends RestClient.FailureListener implements Closeab
         /**
          * Sets the http client. Mandatory argument. Best practice is to use the same client used
          * within {@link org.elasticsearch.client.RestClient} which can be created manually or
-         * through {@link RestClient.Builder#createDefaultHttpClient()}.
+         * through {@link RestClient.Builder#createDefaultHttpClient(Registry)}.
          * @see CloseableHttpClient
          */
         public Builder setRestClient(RestClient restClient) {

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

@@ -37,6 +37,8 @@ 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;
@@ -374,7 +376,7 @@ public final class RestClient implements Closeable {
 
         /**
          * Sets the http client. A new default one will be created if not
-         * specified, by calling {@link #createDefaultHttpClient()}.
+         * specified, by calling {@link #createDefaultHttpClient(Registry)})}.
          *
          * @see CloseableHttpClient
          */
@@ -427,7 +429,7 @@ public final class RestClient implements Closeable {
          */
         public RestClient build() {
             if (httpClient == null) {
-                httpClient = createDefaultHttpClient();
+                httpClient = createDefaultHttpClient(null);
             }
             if (hosts == null || hosts.length == 0) {
                 throw new IllegalArgumentException("no hosts provided");
@@ -436,12 +438,17 @@ public final class RestClient implements Closeable {
         }
 
         /**
-         * Creates an http client with default settings
+         * Creates a {@link CloseableHttpClient} with default settings. Used when the http client instance is not provided.
          *
          * @see CloseableHttpClient
          */
-        public static CloseableHttpClient createDefaultHttpClient() {
+        public static CloseableHttpClient createDefaultHttpClient(Registry<ConnectionSocketFactory> socketFactoryRegistry) {
-            PoolingHttpClientConnectionManager connectionManager = new PoolingHttpClientConnectionManager();
+            PoolingHttpClientConnectionManager connectionManager;
+            if (socketFactoryRegistry == null) {
+                connectionManager = new PoolingHttpClientConnectionManager();
+            } else {
+                connectionManager = new PoolingHttpClientConnectionManager(socketFactoryRegistry);
+            }
             //default settings may be too constraining
             connectionManager.setDefaultMaxPerRoute(10);
             connectionManager.setMaxTotal(30);

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

@@ -29,8 +29,6 @@ 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;
@@ -71,7 +69,8 @@ import java.util.Objects;
 import java.util.Set;
 
 /**
- * REST client used to test the elasticsearch REST layer
+ * REST client used to test the elasticsearch REST layer.
+ * Wraps a {@link RestClient} instance used to send the REST requests.
  * Holds the {@link RestSpec} used to translate api calls into REST calls
  */
 public class RestTestClient implements Closeable {
@@ -136,7 +135,7 @@ public class RestTestClient implements Closeable {
 
         if ("raw".equals(apiName)) {
             // Raw requests are bit simpler....
-            HashMap<String, String> queryStringParams = new HashMap<>(params);
+            Map<String, String> queryStringParams = new HashMap<>(params);
             String method = Objects.requireNonNull(queryStringParams.remove("method"), "Method must be set to use raw request");
             String path = "/"+ Objects.requireNonNull(queryStringParams.remove("path"), "Path must be set to use raw request");
             HttpEntity entity = null;
@@ -267,7 +266,7 @@ public class RestTestClient implements Closeable {
         return restApi;
     }
 
-    protected RestClient createRestClient(URL[] urls, Settings settings) throws IOException {
+    private static RestClient createRestClient(URL[] urls, Settings settings) throws IOException {
         SSLConnectionSocketFactory sslsf;
         String keystorePath = settings.get(TRUSTSTORE_PATH);
         if (keystorePath != null) {
@@ -297,16 +296,7 @@ public class RestTestClient implements Closeable {
                 .register("http", PlainConnectionSocketFactory.getSocketFactory())
                 .register("https", sslsf)
                 .build();
-
+        CloseableHttpClient httpClient = RestClient.Builder.createDefaultHttpClient(socketFactoryRegistry);
-        List<Header> headers = new ArrayList<>();
-        try (ThreadContext threadContext = new ThreadContext(settings)) {
-            for (Map.Entry<String, String> entry : threadContext.getHeaders().entrySet()) {
-                headers.add(new BasicHeader(entry.getKey(), entry.getValue()));
-            }
-        }
-
-        CloseableHttpClient httpClient = HttpClientBuilder.create().setDefaultHeaders(headers)
-                .setConnectionManager(new PoolingHttpClientConnectionManager(socketFactoryRegistry)).build();
 
         String protocol = settings.get(PROTOCOL, "http");
         HttpHost[] hosts = new HttpHost[urls.length];
@@ -314,7 +304,17 @@ public class RestTestClient implements Closeable {
             URL url = urls[i];
             hosts[i] = new HttpHost(url.getHost(), url.getPort(), protocol);
         }
-        return RestClient.builder().setHttpClient(httpClient).setHosts(hosts).build();
+
+        RestClient.Builder builder = RestClient.builder().setHttpClient(httpClient).setHosts(hosts);
+        try (ThreadContext threadContext = new ThreadContext(settings)) {
+            Header[] defaultHeaders = new Header[threadContext.getHeaders().size()];
+            int i = 0;
+            for (Map.Entry<String, String> entry : threadContext.getHeaders().entrySet()) {
+                defaultHeaders[i++] = new BasicHeader(entry.getKey(), entry.getValue());
+            }
+            builder.setDefaultHeaders(defaultHeaders);
+        }
+        return builder.build();
     }
 
     /**