Add support for ServiceUnavailableRetryStrategy (dropwizard#2185)

* Add support for ServiceUnavailableRetryStrategy Resolves dropwizard#2093. * Remove unused imports * Remove unused field * @nickbabcock: expand glob imports
sankate · Nov 21, 2017 · eba2536 · eba2536
1 parent 68be347
commit eba2536
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 1 deletion.
diff --git a/dropwizard-client/src/main/java/io/dropwizard/client/HttpClientBuilder.java b/dropwizard-client/src/main/java/io/dropwizard/client/HttpClientBuilder.java
@@ -13,7 +13,6 @@
 import io.dropwizard.lifecycle.Managed;
 import io.dropwizard.setup.Environment;
 import io.dropwizard.util.Duration;
-import javax.net.ssl.HostnameVerifier;
 import org.apache.http.ConnectionReuseStrategy;
 import org.apache.http.Header;
 import org.apache.http.HttpHost;
@@ -26,6 +25,7 @@
 import org.apache.http.client.HttpClient;
 import org.apache.http.client.HttpRequestRetryHandler;
 import org.apache.http.client.RedirectStrategy;
+import org.apache.http.client.ServiceUnavailableRetryStrategy;
 import org.apache.http.client.config.CookieSpecs;
 import org.apache.http.client.config.RequestConfig;
 import org.apache.http.config.Registry;
@@ -46,6 +46,7 @@
 import org.apache.http.protocol.HttpContext;
 import org.apache.http.protocol.HttpProcessor;
 
+import javax.net.ssl.HostnameVerifier;
 import java.util.List;
 
 /**
@@ -78,6 +79,7 @@ public class HttpClientBuilder {
     private boolean disableContentCompression;
     private List<? extends Header> defaultHeaders;
     private HttpProcessor httpProcessor;
+    private ServiceUnavailableRetryStrategy serviceUnavailableRetryStrategy;
 
     public HttpClientBuilder(MetricRegistry metricRegistry) {
         this.metricRegistry = metricRegistry;
@@ -221,6 +223,17 @@ public HttpClientBuilder using(HttpProcessor httpProcessor) {
         return this;
     }
 
+    /**
+     * Use the given {@link ServiceUnavailableRetryStrategy} instance
+     *
+     * @param serviceUnavailableRetryStrategy a {@link ServiceUnavailableRetryStrategy} instance
+     * @return {@code} this
+     */
+    public HttpClientBuilder using(ServiceUnavailableRetryStrategy serviceUnavailableRetryStrategy) {
+        this.serviceUnavailableRetryStrategy = serviceUnavailableRetryStrategy;
+        return this;
+    }
+
     /**
      * Disable support of decompression of responses
      *
@@ -389,6 +402,10 @@ public long getKeepAliveDuration(HttpResponse response, HttpContext context) {
             builder.setHttpProcessor(httpProcessor);
         }
 
+        if (serviceUnavailableRetryStrategy != null) {
+            builder.setServiceUnavailableRetryStrategy(serviceUnavailableRetryStrategy);
+        }
+
         return new ConfiguredCloseableHttpClient(builder.build(), requestConfig);
     }
 

diff --git a/dropwizard-client/src/main/java/io/dropwizard/client/JerseyClientBuilder.java b/dropwizard-client/src/main/java/io/dropwizard/client/JerseyClientBuilder.java
@@ -14,6 +14,7 @@
 import io.dropwizard.util.Duration;
 import org.apache.http.client.CredentialsProvider;
 import org.apache.http.client.HttpRequestRetryHandler;
+import org.apache.http.client.ServiceUnavailableRetryStrategy;
 import org.apache.http.config.Registry;
 import org.apache.http.conn.DnsResolver;
 import org.apache.http.conn.routing.HttpRoutePlanner;
@@ -313,6 +314,17 @@ public JerseyClientBuilder using(CredentialsProvider credentialsProvider) {
         return this;
     }
 
+    /**
+     * Use the given {@link ServiceUnavailableRetryStrategy} instance.
+     *
+     * @param serviceUnavailableRetryStrategy a {@link ServiceUnavailableRetryStrategy} instance
+     * @return {@code this}
+     */
+    public JerseyClientBuilder using(ServiceUnavailableRetryStrategy serviceUnavailableRetryStrategy) {
+        apacheHttpClientBuilder.using(serviceUnavailableRetryStrategy);
+        return this;
+    }
+
     /**
      * Builds the {@link RxClient} instance.
      *

diff --git a/dropwizard-client/src/test/java/io/dropwizard/client/HttpClientBuilderTest.java b/dropwizard-client/src/test/java/io/dropwizard/client/HttpClientBuilderTest.java
@@ -35,6 +35,7 @@
 import org.apache.http.client.CredentialsProvider;
 import org.apache.http.client.HttpRequestRetryHandler;
 import org.apache.http.client.RedirectStrategy;
+import org.apache.http.client.ServiceUnavailableRetryStrategy;
 import org.apache.http.client.config.CookieSpecs;
 import org.apache.http.client.config.RequestConfig;
 import org.apache.http.client.methods.HttpGet;
@@ -678,6 +679,19 @@ public void usesHttpProcessor() throws Exception {
             .isSameAs(httpProcessor);
     }
 
+    @Test
+    public void usesServiceUnavailableRetryStrategy() throws Exception {
+        ServiceUnavailableRetryStrategy serviceUnavailableRetryStrategy = mock(ServiceUnavailableRetryStrategy.class);
+        final ConfiguredCloseableHttpClient client =
+            builder.using(serviceUnavailableRetryStrategy)
+                .createClient(apacheBuilder, connectionManager, "test");
+        assertThat(client).isNotNull();
+        assertThat(FieldUtils.getField(httpClientBuilderClass,
+            "serviceUnavailStrategy", true)
+            .get(apacheBuilder))
+            .isSameAs(serviceUnavailableRetryStrategy);
+    }
+
     @Test
     public void allowsCustomBuilderConfiguration() throws Exception {
         CustomBuilder builder = new CustomBuilder(new MetricRegistry());

diff --git a/dropwizard-client/src/test/java/io/dropwizard/client/JerseyClientBuilderTest.java b/dropwizard-client/src/test/java/io/dropwizard/client/JerseyClientBuilderTest.java
@@ -13,6 +13,7 @@
 import io.dropwizard.lifecycle.setup.LifecycleEnvironment;
 import io.dropwizard.setup.Environment;
 import org.apache.http.client.CredentialsProvider;
+import org.apache.http.client.ServiceUnavailableRetryStrategy;
 import org.apache.http.config.Registry;
 import org.apache.http.config.RegistryBuilder;
 import org.apache.http.conn.DnsResolver;
@@ -272,6 +273,13 @@ public void usesACustomHostnameVerifier() {
         verify(apacheHttpClientBuilder).using(customHostnameVerifier);
     }
 
+    @Test
+    public void usesACustomServiceUnavailableRetryStrategy() {
+        final ServiceUnavailableRetryStrategy customServiceUnavailableRetryStrategy = mock(ServiceUnavailableRetryStrategy.class);
+        builder.using(customServiceUnavailableRetryStrategy);
+        verify(apacheHttpClientBuilder).using(customServiceUnavailableRetryStrategy);
+    }
+
     @Test
     public void usesACustomConnectionFactoryRegistry() throws Exception {
         final SSLContext ctx = SSLContext.getInstance(SSLConnectionSocketFactory.TLS);