Upgrade nodes on lowest OS version first

A minor optimization of the following scenario: Nodes are on v1, upgrade to v2
is ongoing, and upgrade to v3 is started before v2 is complete.

3 files changed, 30 insertions, 6 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeList.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeList.java
index 37987f0512d..21bde5900eb 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeList.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeList.java
@@ -9,6 +9,7 @@ import com.yahoo.config.provision.NodeType;
 
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.EnumSet;
 import java.util.Iterator;
 import java.util.List;
@@ -79,6 +80,16 @@ public class NodeList implements Iterable<Node> {
         return filter(node -> node.status().osVersion().changing());
     }
 
+    /** Returns a copy of this sorted by current OS version (lowest to highest) */
+    public NodeList byIncreasingOsVersion() {
+        return nodes.stream()
+                    .sorted(Comparator.comparing(node -> node.status()
+                                                             .osVersion()
+                                                             .current()
+                                                             .orElse(Version.emptyVersion)))
+                    .collect(collectingAndThen(Collectors.toList(), NodeList::wrap));
+    }
+
     /** Returns the subset of nodes that are currently on the given OS version */
     public NodeList onOsVersion(Version version) {
         return filter(node -> node.status().osVersion().matches(version));
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersions.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersions.java
index 106595fbd47..9b21d2aceba 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersions.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersions.java
@@ -123,6 +123,7 @@ public class OsVersions {
         var numberToUpgrade = Math.max(0, maxActiveUpgrades - nodes.changingOsVersion().size());
         var nodesToUpgrade = nodes.not().changingOsVersion()
                                   .not().onOsVersion(version)
+                                  .byIncreasingOsVersion()
                                   .first(numberToUpgrade);
         if (nodesToUpgrade.size() == 0) return;
         log.info("Upgrading " + nodesToUpgrade.size() + " nodes of type " + type + " to OS version " + version);
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/os/OsVersionsTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/os/OsVersionsTest.java
index ebb64d650f1..477319a8bbb 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/os/OsVersionsTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/os/OsVersionsTest.java
@@ -12,6 +12,7 @@ import java.util.Comparator;
 import java.util.List;
 import java.util.Optional;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -75,8 +76,11 @@ public class OsVersionsTest {
         tester.makeReadyNodes(totalNodes, "default", NodeType.host);
         Supplier<NodeList> hostNodes = () -> tester.nodeRepository().list().nodeType(NodeType.host);
 
-        // Some nodes have reported current version
-        setCurrentVersion(hostNodes.get().asList().subList(0, 2), Version.fromString("7.0"));
+        // 5 nodes have no version. The other 15 are spread across different versions
+        var hostNodesList = hostNodes.get().asList();
+        for (int i = totalNodes - maxActiveUpgrades - 1; i >= 0; i--) {
+            setCurrentVersion(List.of(hostNodesList.get(i)), new Version(7, 0, i));
+        }
 
         // Set target
         var version1 = Version.fromString("7.1");
@@ -86,20 +90,28 @@ public class OsVersionsTest {
         // Activate target
         for (int i = 0; i < totalNodes; i += maxActiveUpgrades) {
             versions.setActive(NodeType.host, true);
-            var nodesUpgrading = hostNodes.get().changingOsVersion();
+            var nodes = hostNodes.get();
+            var nodesUpgrading = nodes.changingOsVersion();
             assertEquals("Target is changed for a subset of nodes", maxActiveUpgrades, nodesUpgrading.size());
             assertEquals("Wanted version is set for nodes upgrading", version1,
                          nodesUpgrading.stream()
                                        .map(node -> node.status().osVersion().wanted().get())
                                        .min(Comparator.naturalOrder()).get());
+            var nodesOnLowestVersion = nodes.asList().stream()
+                                            .sorted(Comparator.comparing(node -> node.status().osVersion().current().orElse(Version.emptyVersion)))
+                                            .collect(Collectors.toList())
+                                            .subList(0, maxActiveUpgrades);
+            assertEquals("Nodes on lowest version are told to upgrade",
+                         nodesUpgrading.asList(), nodesOnLowestVersion);
             completeUpgradeOf(nodesUpgrading.asList());
         }
 
         // Activating again after all nodes have upgraded does nothing
         versions.setActive(NodeType.host, true);
-        assertEquals(version1, hostNodes.get().stream()
-                                        .map(n -> n.status().osVersion().current().get())
-                                        .min(Comparator.naturalOrder()).get());
+        assertEquals("All nodes upgraded", version1,
+                     hostNodes.get().stream()
+                              .map(n -> n.status().osVersion().current().get())
+                              .min(Comparator.naturalOrder()).get());
     }
 
     private void setCurrentVersion(List<Node> nodes, Version currentVersion) {