Refactor Sniffer and make it testable

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

@@ -61,6 +61,7 @@
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -132,7 +133,7 @@ public synchronized void setHosts(HttpHost... hosts) {
         if (hosts == null || hosts.length == 0) {
             throw new IllegalArgumentException("hosts must not be null nor empty");
         }
-        Set<HttpHost> httpHosts = new HashSet<>();
+        Set<HttpHost> httpHosts = new LinkedHashSet<>();
         AuthCache authCache = new BasicAuthCache();
         for (HttpHost host : hosts) {
             Objects.requireNonNull(host, "host cannot be null");
@@ -143,6 +144,13 @@ public synchronized void setHosts(HttpHost... hosts) {
         this.blacklist.clear();
     }
 
+    /**
+     * Returns the configured hosts
+     */
+    public List<HttpHost> getHosts() {
+        return new ArrayList<>(hostTuple.hosts);
+    }
+
     /**
      * Sends a request to the Elasticsearch cluster that the client points to and waits for the corresponding response
      * to be returned. Shortcut to {@link #performRequest(String, String, Map, HttpEntity, Header...)} but without parameters

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

@@ -25,6 +25,7 @@
 
 import java.io.IOException;
 import java.net.URI;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.concurrent.CountDownLatch;
 
@@ -175,6 +176,37 @@ public void testSetHostsWrongArguments() throws IOException {
         }
     }
 
+    public void testSetHostsPreservesOrdering() throws Exception {
+        try (RestClient restClient = createRestClient()) {
+            HttpHost[] hosts = randomHosts();
+            restClient.setHosts(hosts);
+            assertEquals(Arrays.asList(hosts), restClient.getHosts());
+        }
+    }
+
+    private static HttpHost[] randomHosts() {
+        int numHosts = randomIntBetween(1, 10);
+        HttpHost[] hosts = new HttpHost[numHosts];
+        for (int i = 0; i < hosts.length; i++) {
+            hosts[i] = new HttpHost("host-" + i, 9200);
+        }
+        return hosts;
+    }
+
+    public void testSetHostsDuplicatedHosts() throws Exception {
+        try (RestClient restClient = createRestClient()) {
+            int numHosts = randomIntBetween(1, 10);
+            HttpHost[] hosts = new HttpHost[numHosts];
+            HttpHost host = new HttpHost("host", 9200);
+            for (int i = 0; i < hosts.length; i++) {
+                hosts[i] = host;
+            }
+            restClient.setHosts(hosts);
+            assertEquals(1, restClient.getHosts().size());
+            assertEquals(host, restClient.getHosts().get(0));
+        }
+    }
+
     public void testNullPath() throws IOException {
         try (RestClient restClient = createRestClient()) {
             for (String method : getHttpMethods()) {

client/sniffer/src/main/java/org/elasticsearch/client/sniff/SniffOnFailureListener.java

@@ -58,7 +58,6 @@ public void onFailure(HttpHost host) {
         if (sniffer == null) {
             throw new IllegalStateException("sniffer was not set, unable to sniff on failure");
         }
-        //re-sniff immediately but take out the node that failed
-        sniffer.sniffOnFailure(host);
+        sniffer.sniffOnFailure();
     }
 }

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

@@ -31,12 +31,13 @@
 import java.security.PrivilegedAction;
 import java.util.List;
 import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * Class responsible for sniffing nodes from some source (default is elasticsearch itself) and setting them to a provided instance of
@@ -51,103 +52,84 @@ public class Sniffer implements Closeable {
     private static final Log logger = LogFactory.getLog(Sniffer.class);
     private static final String SNIFFER_THREAD_NAME = "es_rest_client_sniffer";
 
-    private final Task task;
+    private final HostsSniffer hostsSniffer;
+    private final RestClient restClient;
+
+    private final long sniffIntervalMillis;
+    private final long sniffAfterFailureDelayMillis;
+    private final Scheduler scheduler;
+
+    private final AtomicReference<Future> nextTask = new AtomicReference<>();
 
     Sniffer(RestClient restClient, HostsSniffer hostsSniffer, long sniffInterval, long sniffAfterFailureDelay) {
-        this.task = new Task(hostsSniffer, restClient, sniffInterval, sniffAfterFailureDelay);
+        this(restClient, hostsSniffer, new DefaultScheduler(), sniffInterval, sniffAfterFailureDelay);
+    }
+
+    Sniffer(RestClient restClient, HostsSniffer hostsSniffer, Scheduler scheduler,  long sniffInterval, long sniffAfterFailureDelay) {
+        this.hostsSniffer = hostsSniffer;
+        this.restClient = restClient;
+        this.sniffIntervalMillis = sniffInterval;
+        this.sniffAfterFailureDelayMillis = sniffAfterFailureDelay;
+        this.scheduler = scheduler;
+        //first sniffing round is immediately executed, next one will be executed depending on the configured sniff interval
+        scheduleNextRound(0L, sniffIntervalMillis, false);
     }
 
     /**
-     * Triggers a new sniffing round and explicitly takes out the failed host provided as argument
+     * Triggers a new immediate sniffing round, which will schedule a new round in sniffAfterFailureDelayMillis ms
      */
-    public void sniffOnFailure(HttpHost failedHost) {
-        this.task.sniffOnFailure(failedHost);
+    public final void sniffOnFailure() {
+        scheduleNextRound(0L, sniffAfterFailureDelayMillis, true);
     }
 
-    @Override
-    public void close() throws IOException {
-        task.shutdown();
+    private void scheduleNextRound(long delay, long nextDelay, boolean mustCancelNextRound) {
+        Task task = new Task(nextDelay);
+        Future<?> nextFuture = task.schedule(delay);
+        Future<?> previousFuture = nextTask.getAndSet(nextFuture);
+        if (mustCancelNextRound) {
+            previousFuture.cancel(false);
+        }
     }
 
-    private static class Task implements Runnable {
-        private final HostsSniffer hostsSniffer;
-        private final RestClient restClient;
-
-        private final long sniffIntervalMillis;
-        private final long sniffAfterFailureDelayMillis;
-        private final ScheduledExecutorService scheduledExecutorService;
-        private final AtomicBoolean running = new AtomicBoolean(false);
-        private ScheduledFuture<?> scheduledFuture;
-
-        private Task(HostsSniffer hostsSniffer, RestClient restClient, long sniffIntervalMillis, long sniffAfterFailureDelayMillis) {
-            this.hostsSniffer = hostsSniffer;
-            this.restClient = restClient;
-            this.sniffIntervalMillis = sniffIntervalMillis;
-            this.sniffAfterFailureDelayMillis = sniffAfterFailureDelayMillis;
-            SnifferThreadFactory threadFactory = new SnifferThreadFactory(SNIFFER_THREAD_NAME);
-            this.scheduledExecutorService = Executors.newScheduledThreadPool(1, threadFactory);
-            scheduleNextRun(0);
+    final class Task implements Runnable {
+        final long nextTaskDelay;
+
+        Task(long nextTaskDelay) {
+            this.nextTaskDelay = nextTaskDelay;
         }
 
-        synchronized void scheduleNextRun(long delayMillis) {
-            if (scheduledExecutorService.isShutdown() == false) {
-                try {
-                    if (scheduledFuture != null) {
-                        //regardless of when the next sniff is scheduled, cancel it and schedule a new one with updated delay
-                        this.scheduledFuture.cancel(false);
-                    }
-                    logger.debug("scheduling next sniff in " + delayMillis + " ms");
-                    this.scheduledFuture = this.scheduledExecutorService.schedule(this, delayMillis, TimeUnit.MILLISECONDS);
-                } catch(Exception e) {
-                    logger.error("error while scheduling next sniffer task", e);
-                }
-            }
+        Future<?> schedule(long delay) {
+            return scheduler.schedule(this, delay);
         }
 
         @Override
         public void run() {
-            sniff(null, sniffIntervalMillis);
-        }
-
-        void sniffOnFailure(HttpHost failedHost) {
-            sniff(failedHost, sniffAfterFailureDelayMillis);
-        }
-
-        void sniff(HttpHost excludeHost, long nextSniffDelayMillis) {
-            if (running.compareAndSet(false, true)) {
-                try {
-                    List<HttpHost> sniffedHosts = hostsSniffer.sniffHosts();
-                    logger.debug("sniffed hosts: " + sniffedHosts);
-                    if (excludeHost != null) {
-                        sniffedHosts.remove(excludeHost);
-                    }
-                    if (sniffedHosts.isEmpty()) {
-                        logger.warn("no hosts to set, hosts will be updated at the next sniffing round");
-                    } else {
-                        this.restClient.setHosts(sniffedHosts.toArray(new HttpHost[sniffedHosts.size()]));
-                    }
-                } catch (Exception e) {
-                    logger.error("error while sniffing nodes", e);
-                } finally {
-                    scheduleNextRun(nextSniffDelayMillis);
-                    running.set(false);
-                }
+            try {
+                sniff();
+            } catch (Exception e) {
+                logger.error("error while sniffing nodes", e);
+            } finally {
+                scheduleNextRound(nextTaskDelay, sniffIntervalMillis, false);
             }
         }
+    }
 
-        synchronized void shutdown() {
-            scheduledExecutorService.shutdown();
-            try {
-                if (scheduledExecutorService.awaitTermination(1000, TimeUnit.MILLISECONDS)) {
-                    return;
-                }
-                scheduledExecutorService.shutdownNow();
-            } catch (InterruptedException e) {
-                Thread.currentThread().interrupt();
-            }
+    final void sniff() throws IOException {
+        List<HttpHost> sniffedHosts = hostsSniffer.sniffHosts();
+        logger.debug("sniffed hosts: " + sniffedHosts);
+        if (sniffedHosts.isEmpty()) {
+            logger.warn("no hosts to set, hosts will be updated at the next sniffing round");
+        } else {
+            restClient.setHosts(sniffedHosts.toArray(new HttpHost[sniffedHosts.size()]));
         }
     }
 
+    @Override
+    public void close() {
+        nextTask.get().cancel(false);
+        this.scheduler.shutdown();
+    }
+
     /**
      * Returns a new {@link SnifferBuilder} to help with {@link Sniffer} creation.
      *
@@ -158,8 +140,62 @@ public static SnifferBuilder builder(RestClient restClient) {
         return new SnifferBuilder(restClient);
     }
 
-    private static class SnifferThreadFactory implements ThreadFactory {
+    /**
+     * The Scheduler interface allows to isolate the sniffing scheduling aspects so that we can test
+     * the sniffer by injecting when needed a custom scheduler that is more suited for testing.
+     */
+    interface Scheduler {
+        /**
+         * Schedules the provided {@link Runnable} to be executed in <code>delayMillis</code> milliseconds
+         */
+        Future<?> schedule(Task task, long delayMillis);
+
+        /**
+         * Shuts this scheduler down
+         */
+        void shutdown();
+    }
+
+    /**
+     * Default implementation of {@link Scheduler}, based on {@link ScheduledExecutorService}
+     */
+    static final class DefaultScheduler implements Scheduler {
+        final ScheduledExecutorService executor;
+
+        DefaultScheduler() {
+            this(initScheduledExecutorService());
+        }
+
+        DefaultScheduler(ScheduledExecutorService executor) {
+            this.executor = executor;
+        }
+
+        private static ScheduledExecutorService initScheduledExecutorService() {
+            ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(1, new SnifferThreadFactory(SNIFFER_THREAD_NAME));
+            executor.setRemoveOnCancelPolicy(true);
+            return executor;
+        }
+
+        @Override
+        public Future<?> schedule(Task task, long delayMillis) {
+            return executor.schedule(task, delayMillis, TimeUnit.MILLISECONDS);
+        }
+
+        @Override
+        public void shutdown() {
+            executor.shutdown();
+            try {
+                if (executor.awaitTermination(1000, TimeUnit.MILLISECONDS)) {
+                    return;
+                }
+                executor.shutdownNow();
+            } catch (InterruptedException ignore) {
+                Thread.currentThread().interrupt();
+            }
+        }
+    }
 
+    static class SnifferThreadFactory implements ThreadFactory {
         private final AtomicInteger threadNumber = new AtomicInteger(1);
         private final String namePrefix;
         private final ThreadFactory originalThreadFactory;

client/sniffer/src/test/java/org/elasticsearch/client/sniff/MockHostsSniffer.java

@@ -21,7 +21,6 @@
 
 import org.apache.http.HttpHost;
 
-import java.io.IOException;
 import java.util.Collections;
 import java.util.List;
 
@@ -30,7 +29,7 @@
  */
 class MockHostsSniffer implements HostsSniffer {
     @Override
-    public List<HttpHost> sniffHosts() throws IOException {
+    public List<HttpHost> sniffHosts() {
         return Collections.singletonList(new HttpHost("localhost", 9200));
     }
 }