Ensure node search path index follows config order

3 files changed, 20 insertions, 10 deletions

diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/searchcluster/Node.java b/container-search/src/main/java/com/yahoo/search/dispatch/searchcluster/Node.java
index c93f2f4c491..31e02f910ee 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/searchcluster/Node.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/searchcluster/Node.java
@@ -35,7 +35,7 @@ public class Node {
 
     /** Give a monotonically increasing sequence number.*/
     public long createPingSequenceId() { return pingSequence.incrementAndGet(); }
-    /** Checks if this pong is received in line and accepted, or out of band and should be ignored..*/
+    /** Checks if this pong is received in line and accepted, or out of band and should be ignored. */
     public boolean isLastReceivedPong(long pingId ) {
         long last = lastPong.get();
         while ((pingId > last) && ! lastPong.compareAndSet(last, pingId)) {
diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/searchcluster/SearchCluster.java b/container-search/src/main/java/com/yahoo/search/dispatch/searchcluster/SearchCluster.java
index ce602ffcdaa..5005c772e0e 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/searchcluster/SearchCluster.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/searchcluster/SearchCluster.java
@@ -8,16 +8,18 @@ import com.yahoo.search.cluster.ClusterMonitor;
 import com.yahoo.search.cluster.NodeManager;
 import com.yahoo.yolean.UncheckedInterruptedException;
 
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
+import java.util.ListIterator;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.Executor;
 import java.util.logging.Logger;
 
 import static java.util.stream.Collectors.groupingBy;
+import static java.util.stream.Collectors.toMap;
 
 /**
  * A model of a search cluster we might want to dispatch queries to.
@@ -62,12 +64,14 @@ public class SearchCluster implements NodeManager<Node> {
 
     /** Sets the new nodes to monitor to be the new nodes, but keep any existing node instances which equal the new ones. */
     public ClusterMonitor<Node> updateNodes(Collection<Node> newNodes, double minActivedocsPercentage) {
-        Collection<Node> retainedNodes = groups.nodes();
-        Collection<Node> currentNodes = new HashSet<>(newNodes);
-        retainedNodes.retainAll(currentNodes);          // Throw away all old nodes which are not in the new set.
-        currentNodes.removeIf(retainedNodes::contains); // Throw away all new nodes for which we have more information in an old object.
-        Collection<Node> addedNodes = List.copyOf(currentNodes);
-        currentNodes.addAll(retainedNodes);             // Keep the old nodes that were replaced in the new set.
+        List<Node> currentNodes = new ArrayList<>(newNodes);
+        List<Node> addedNodes = new ArrayList<>();
+        Map<Node, Node> retainedNodes = groups.nodes().stream().collect(toMap(node -> node, node -> node));
+        for (int i = 0; i < currentNodes.size(); i++) {
+            Node retained = retainedNodes.get(currentNodes.get(i));
+            if (retained != null) currentNodes.set(i, retained);
+            else addedNodes.add(currentNodes.get(i));
+        }
         SearchGroupsImpl groups = toGroups(currentNodes, minActivedocsPercentage);
         ClusterMonitor<Node> monitor = new ClusterMonitor<>(this, false);
         for (Node node : groups.nodes()) monitor.add(node, true);
diff --git a/container-search/src/test/java/com/yahoo/search/dispatch/searchcluster/SearchClusterTest.java b/container-search/src/test/java/com/yahoo/search/dispatch/searchcluster/SearchClusterTest.java
index 65b0261b1f8..f6a1ca5cae3 100644
--- a/container-search/src/test/java/com/yahoo/search/dispatch/searchcluster/SearchClusterTest.java
+++ b/container-search/src/test/java/com/yahoo/search/dispatch/searchcluster/SearchClusterTest.java
@@ -403,14 +403,20 @@ public class SearchClusterTest {
             SearchGroups newGroups = state.searchCluster.groupList();
             assertEquals(Set.copyOf(updatedNodes), newGroups.nodes());
 
-            Map<Node, Node> oldNodesByIdentity = newGroups.nodes().stream().collect(toMap(identity(), identity()));
+            Map<Node, Node> oldNodesByIdentity = oldGroups.nodes().stream().collect(toMap(identity(), identity()));
             Map<Node, Node> newNodesByIdentity = newGroups.nodes().stream().collect(toMap(identity(), identity()));
             assertSame(updatedNodes.get(0), newNodesByIdentity.get(updatedNodes.get(0)));
             assertSame(updatedNodes.get(1), newNodesByIdentity.get(updatedNodes.get(1)));
             assertSame(updatedNodes.get(2), newNodesByIdentity.get(updatedNodes.get(2)));
-            assertSame(oldNodesByIdentity.get(referenceNodes.get(3)), newNodesByIdentity.get(updatedNodes.get(3)));
+            assertSame(oldNodesByIdentity.get(updatedNodes.get(3)), newNodesByIdentity.get(updatedNodes.get(3)));
             assertSame(updatedNodes.get(4), newNodesByIdentity.get(updatedNodes.get(4)));
             assertSame(updatedNodes.get(5), newNodesByIdentity.get(updatedNodes.get(5)));
+
+            // Also verify search-path index within group follows node order, as given by config.
+            int[] pathIndexWithinGroup = new int[3];
+            for (Node node : updatedNodes)
+                assertEquals(pathIndexWithinGroup[node.group()]++, newNodesByIdentity.get(node).pathIndex(),
+                             "search path index within group should follow updated node order for: " + node);
         }
     }