LLClient: Support host selection

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

@@ -654,18 +654,18 @@ static List<Node> selectHosts(NodeTuple<List<Node>> nodeTuple,
 
         if (false == livingNodes.isEmpty()) {
             /*
-             * Normal state: there is at least one living node. Rotate the
-             * list so subsequent requests will prefer the nodes in a
-             * different order then run them through the NodeSelector so it
-             * can have its say in which nodes are ok and their ordering. If
-             * the selector is ok with any over the living nodes then use
-             * them for the request.
+             * Normal state: there is at least one living node. If the
+             * selector is ok with any over the living nodes then use them
+             * for the request.
              */
-            // TODO this is going to send more requests to nodes right *after* a node that the selector removes
             List<Node> selectedLivingNodes = new ArrayList<>(livingNodes);
-            Collections.rotate(selectedLivingNodes, lastNodeIndex.getAndIncrement());
             nodeSelector.select(selectedLivingNodes);
             if (false == selectedLivingNodes.isEmpty()) {
+                /*
+                 * Rotate the list so subsequent requests will prefer the
+                 * nodes in a different order.
+                 */
+                Collections.rotate(selectedLivingNodes, lastNodeIndex.getAndIncrement());
                 return selectedLivingNodes;
             }
         }
@@ -682,21 +682,54 @@ static List<Node> selectHosts(NodeTuple<List<Node>> nodeTuple,
          * node.
          */
         if (false == deadNodes.isEmpty()) {
-            Collections.sort(deadNodes);
-
-            List<Node> selectedDeadNodes = new ArrayList<>(deadNodes.size());
-            for (DeadNodeAndRevival n : deadNodes) {
-                selectedDeadNodes.add(n.node);
-            }
-            nodeSelector.select(selectedDeadNodes);
+            final List<DeadNodeAndRevival> selectedDeadNodes = new ArrayList<>(deadNodes);
+            /*
+             * We'd like NodeSelectors to remove items directly from deadNodes
+             * so we can find the minimum after it is filtered without having
+             * to compare many things. This saves us a sort on the unfiltered
+             * list.
+             */
+            nodeSelector.select(new Iterable<Node>() {
+                @Override
+                public Iterator<Node> iterator() {
+                    return new Adapter(selectedDeadNodes.iterator());
+                }
+            });
             if (false == selectedDeadNodes.isEmpty()) {
-                return singletonList(selectedDeadNodes.get(0));
+                return singletonList(Collections.min(selectedDeadNodes).node);
             }
         }
         throw new IOException("NodeSelector [" + nodeSelector + "] rejected all nodes, "
                 + "living " + livingNodes + " and dead " + deadNodes);
     }
 
+    /**
+     * Adapts an <code>Iterator<DeadNodeAndRevival></code> into an
+     * <code>Iterator<Node></code>.
+     */
+    private static class Adapter implements Iterator<Node> {
+        private final Iterator<DeadNodeAndRevival> itr;
+
+        private Adapter(Iterator<DeadNodeAndRevival> itr) {
+            this.itr = itr;
+        }
+
+        @Override
+        public boolean hasNext() {
+            return itr.hasNext();
+        }
+
+        @Override
+        public Node next() {
+            return itr.next().node;
+        }
+
+        @Override
+        public void remove() {
+            itr.remove();
+        }
+    }
+
     /**
      * Called after each successful request call.
      * Receives as an argument the host that was used for the successful request.

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

@@ -340,11 +340,9 @@ public void testNullPath() throws IOException {
     }
 
     public void testSelectHosts() throws IOException {
-        int iterations = 1000;
         Node n1 = new Node(new HttpHost("1"), null, null, "1", null);
         Node n2 = new Node(new HttpHost("2"), null, null, "2", null);
         Node n3 = new Node(new HttpHost("3"), null, null, "3", null);
-        List<Node> nodes = Arrays.asList(n1, n2, n3);
 
         NodeSelector not1 = new NodeSelector() {
             @Override
@@ -376,106 +374,101 @@ public String toString() {
             }
         };
 
-        NodeTuple<List<Node>> nodeTuple = new NodeTuple<>(nodes, null);
-        Map<HttpHost, DeadHostState> blacklist = new HashMap<>();
-        AtomicInteger lastNodeIndex = new AtomicInteger(0);
-        long now = 0;
+        NodeTuple<List<Node>> nodeTuple = new NodeTuple<>(Arrays.asList(n1, n2, n3), null);
 
-        // Normal case
-        List<Node> expectedNodes = Arrays.asList(n1, n2, n3);
-        assertEquals(expectedNodes, RestClient.selectHosts(nodeTuple, blacklist,
-                lastNodeIndex, now, NodeSelector.ANY));
-        // Calling it again rotates the set of results
-        for (int i = 0; i < iterations; i++) {
-            Collections.rotate(expectedNodes, 1);
-            assertEquals(expectedNodes, RestClient.selectHosts(nodeTuple, blacklist,
-                    lastNodeIndex, now, NodeSelector.ANY));
-        }
+        Map<HttpHost, DeadHostState> emptyBlacklist = Collections.emptyMap();
 
-        // Exclude some node
-        lastNodeIndex.set(0);
-        // h1 excluded
-        assertEquals(Arrays.asList(n2, n3), RestClient.selectHosts(nodeTuple, blacklist,
-                lastNodeIndex, now, not1));
-        // Calling it again rotates the set of results
-        assertEquals(Arrays.asList(n3, n2), RestClient.selectHosts(nodeTuple, blacklist,
-                lastNodeIndex, now, not1));
-        // And again, same
-        assertEquals(Arrays.asList(n2, n3), RestClient.selectHosts(nodeTuple, blacklist,
-                lastNodeIndex, now, not1));
-        /*
-         * But this time it doesn't because the list being filtered changes
-         * from (h1, h2, h3) to (h2, h3, h1) which both look the same when
-         * you filter out h1.
-         */
-        assertEquals(Arrays.asList(n2, n3), RestClient.selectHosts(nodeTuple, blacklist,
-                lastNodeIndex, now, not1));
+        // Normal cases where the node selector doesn't reject all living nodes
+        assertSelectLivingHosts(Arrays.asList(n1, n2, n3), nodeTuple, emptyBlacklist, 0, NodeSelector.ANY);
+        assertSelectLivingHosts(Arrays.asList(n2, n3), nodeTuple, emptyBlacklist, 0, not1);
 
         /*
          * Try a NodeSelector that excludes all nodes. This should
          * throw an exception
          */
-        lastNodeIndex.set(0);
-        try {
-            RestClient.selectHosts(nodeTuple, blacklist, lastNodeIndex, now, noNodes);
-            fail("expected selectHosts to fail");
-        } catch (IOException e) {
+        {
             String message = "NodeSelector [NONE] rejected all nodes, living ["
                     + "[host=http://1, version=1], [host=http://2, version=2], "
                     + "[host=http://3, version=3]] and dead []";
-            assertEquals(message, e.getMessage());
+            assertEquals(message, assertSelectAllRejected(nodeTuple, emptyBlacklist, 0, noNodes));
         }
 
-        /*
-         * Mark all nodes as dead and look up at a time *after* the
-         * revival time. This should return all nodes.
-         */
-        blacklist.put(n1.getHost(), new DeadHostState(1));
-        blacklist.put(n2.getHost(), new DeadHostState(new DeadHostState(1), 1));
-        blacklist.put(n3.getHost(), new DeadHostState(new DeadHostState(new DeadHostState(1), 1), 1));
-        lastNodeIndex.set(0);
-        now = DeadHostState.MAX_CONNECTION_TIMEOUT_NANOS + 1;
-        expectedNodes = Arrays.asList(n1, n2, n3);
-        assertEquals(expectedNodes, RestClient.selectHosts(nodeTuple, blacklist, lastNodeIndex,
-                now, NodeSelector.ANY));
+        // Mark all the nodes dead for a few test cases
+        {
+            long now = 0;
+            Map<HttpHost, DeadHostState> blacklist = new HashMap<>();
+            blacklist.put(n1.getHost(), new DeadHostState(now));
+            blacklist.put(n2.getHost(), new DeadHostState(new DeadHostState(now), now));
+            blacklist.put(n3.getHost(), new DeadHostState(new DeadHostState(new DeadHostState(now), now), now));
+
+            /*
+             * selectHosts will revive a single host if regardless of
+             * blacklist time. It'll revive the node that is closest
+             * to being revived that the NodeSelector is ok with.
+             */
+            assertEquals(singletonList(n1), RestClient.selectHosts(nodeTuple, blacklist, new AtomicInteger(), now, NodeSelector.ANY));
+            assertEquals(singletonList(n2), RestClient.selectHosts(nodeTuple, blacklist, new AtomicInteger(), now, not1));
+
+            /*
+             * Try a NodeSelector that excludes all nodes. This should
+             * return a failure, but a different failure than when the
+             * blacklist is empty so that the caller knows that all of
+             * their nodes are blacklisted AND blocked.
+             */
+            String message = "NodeSelector [NONE] rejected all nodes, living [] and dead ["
+                    + "[host=http://1, version=1], [host=http://2, version=2], "
+                    + "[host=http://3, version=3]]";
+            assertEquals(message, assertSelectAllRejected(nodeTuple, blacklist, now, noNodes));
+
+            /*
+             * Now lets wind the clock forward, past the timeout for one of
+             * the dead nodes. We should return it.
+             */
+            now = new DeadHostState(now).getDeadUntilNanos();
+            assertSelectLivingHosts(Arrays.asList(n1), nodeTuple, blacklist, 0, NodeSelector.ANY);
+
+            /*
+             * But if the NodeSelector rejects that node then we'll pick the
+             * first on that the NodeSelector doesn't reject.
+             */
+            assertSelectLivingHosts(Arrays.asList(n2), nodeTuple, blacklist, 0, not1);
+
+            /*
+             * If we wind the clock way into the future, past any of the
+             * blacklist timeouts then we function as though the nodes aren't
+             * in the blacklist at all.
+             */
+            now += DeadHostState.MAX_CONNECTION_TIMEOUT_NANOS;
+            assertSelectLivingHosts(Arrays.asList(n1, n2, n3), nodeTuple, blacklist, now, NodeSelector.ANY);
+            assertSelectLivingHosts(Arrays.asList(n2, n3), nodeTuple, blacklist, now, not1);
+        }
+    }
+
+    private void assertSelectLivingHosts(List<Node> expectedNodes, NodeTuple<List<Node>> nodeTuple,
+            Map<HttpHost, DeadHostState> blacklist, long now, NodeSelector nodeSelector) throws IOException {
+        int iterations = 1000;
+        AtomicInteger lastNodeIndex = new AtomicInteger(0);
+        assertEquals(expectedNodes, RestClient.selectHosts(nodeTuple, blacklist,
+                lastNodeIndex, now, nodeSelector));
         // Calling it again rotates the set of results
-        for (int i = 0; i < iterations; i++) {
+        for (int i = 1; i < iterations; i++) {
             Collections.rotate(expectedNodes, 1);
-            assertEquals(expectedNodes, RestClient.selectHosts(nodeTuple, blacklist,
-                    lastNodeIndex, now, NodeSelector.ANY));
+            assertEquals("iteration " + i, expectedNodes, RestClient.selectHosts(nodeTuple, blacklist,
+                    lastNodeIndex, now, nodeSelector));
         }
+    }
 
-        /*
-         * Now try with the nodes dead and *not* past their dead time.
-         * Only the node closest to revival should come back.
-         */
-        now = 0;
-        assertEquals(singletonList(n1), RestClient.selectHosts(nodeTuple, blacklist, lastNodeIndex,
-                now, NodeSelector.ANY));
-
-        /*
-         * Now try with the nodes dead and *not* past their dead time
-         * *and* a node selector that removes the node that is closest
-         * to being revived. The second closest node should come back.
-         */
-        assertEquals(singletonList(n2), RestClient.selectHosts(nodeTuple, blacklist,
-                lastNodeIndex, now, not1));
-
-        /*
-         * Try a NodeSelector that excludes all nodes. This should
-         * return a failure, but a different failure than normal
-         * because it'll block revival rather than outright reject
-         * healthy nodes.
-         */
-        lastNodeIndex.set(0);
+    /**
+     * Assert that {@link RestClient#selectHosts} fails on the provided arguments.
+     * @return the message in the exception thrown by the failure
+     */
+    private String assertSelectAllRejected( NodeTuple<List<Node>> nodeTuple,
+            Map<HttpHost, DeadHostState> blacklist, long now, NodeSelector nodeSelector) {
         try {
-            RestClient.selectHosts(nodeTuple, blacklist, lastNodeIndex, now, noNodes);
-            fail("expected selectHosts to fail");
+            RestClient.selectHosts(nodeTuple, blacklist, new AtomicInteger(0), now, nodeSelector);
+            throw new AssertionError("expected selectHosts to fail");
         } catch (IOException e) {
-            String message = "NodeSelector [NONE] rejected all nodes, living [] and dead ["
-                    + "[host=http://1, version=1], [host=http://2, version=2], "
-                    + "[host=http://3, version=3]]";
-            assertEquals(message, e.getMessage());
+            return e.getMessage();
         }
     }