Support multiple OS upgrade implementations

4 files changed, 92 insertions, 45 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/DelegatingUpgrader.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/DelegatingUpgrader.java
new file mode 100644
index 00000000000..66692f0af4e
--- /dev/null
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/DelegatingUpgrader.java
@@ -0,0 +1,61 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.provision.os;
+
+import com.yahoo.component.Version;
+import com.yahoo.config.provision.NodeType;
+import com.yahoo.vespa.hosted.provision.Node;
+import com.yahoo.vespa.hosted.provision.NodeList;
+import com.yahoo.vespa.hosted.provision.NodeRepository;
+import com.yahoo.vespa.hosted.provision.node.filter.NodeListFilter;
+
+import java.util.Objects;
+import java.util.Optional;
+import java.util.logging.Logger;
+
+/**
+ * An upgrader that delegates the upgrade to the node itself, triggered by changing its wanted OS version. This
+ * implementation limits the number of parallel upgrades to avoid overloading the orchestrator with suspension requests.
+ *
+ * Used in clouds where nodes can upgrade themselves in-place, without data loss.
+ *
+ * @author mpolden
+ */
+public class DelegatingUpgrader implements Upgrader {
+
+    private static final Logger LOG = Logger.getLogger(DelegatingUpgrader.class.getName());
+
+    private final NodeRepository nodeRepository;
+
+    /** The maximum number of nodes, within a single node type, that can upgrade in parallel. */
+    private final int maxActiveUpgrades;
+
+    public DelegatingUpgrader(NodeRepository nodeRepository, int maxActiveUpgrades) {
+        this.nodeRepository = Objects.requireNonNull(nodeRepository);
+        this.maxActiveUpgrades = maxActiveUpgrades;
+    }
+
+    @Override
+    public void upgrade(NodeType type, Version version) {
+        NodeList activeNodes = nodeRepository.list().nodeType(type).state(Node.State.active);
+        int numberToUpgrade = Math.max(0, maxActiveUpgrades - activeNodes.changingOsVersionTo(version).size());
+        NodeList nodesToUpgrade = activeNodes.not().changingOsVersionTo(version)
+                                             .not().onOsVersion(version)
+                                             .byIncreasingOsVersion()
+                                             .first(numberToUpgrade);
+        if (nodesToUpgrade.size() == 0) return;
+        LOG.info("Upgrading " + nodesToUpgrade.size() + " nodes of type " + type + " to OS version " +
+                 version.toFullString());
+        nodeRepository.upgradeOs(NodeListFilter.from(nodesToUpgrade.asList()), Optional.of(version));
+    }
+
+    @Override
+    public void disableUpgrade(NodeType type) {
+        NodeList nodesUpgrading = nodeRepository.list()
+                                                .nodeType(type)
+                                                .changingOsVersion();
+        if (nodesUpgrading.size() == 0) return;
+        LOG.info("Disabling OS upgrade of all " + type + " nodes");
+        nodeRepository.upgradeOs(NodeListFilter.from(nodesUpgrading.asList()), Optional.empty());
+    }
+
+}
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 be474eddf97..f1c4e896310 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
@@ -1,16 +1,13 @@
-// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.provision.os;
 
 import com.yahoo.component.Version;
 import com.yahoo.config.provision.NodeType;
 import com.yahoo.vespa.curator.Lock;
-import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.NodeRepository;
-import com.yahoo.vespa.hosted.provision.node.filter.NodeListFilter;
 import com.yahoo.vespa.hosted.provision.persistence.CuratorDatabaseClient;
 
 import java.util.Map;
-import java.util.Objects;
 import java.util.Optional;
 import java.util.logging.Logger;
 
@@ -28,24 +25,16 @@ public class OsVersions {
 
     private static final Logger log = Logger.getLogger(OsVersions.class.getName());
 
-    /**
-     * The maximum number of nodes, within a single node type, that can upgrade in parallel. We limit the number of
-     * concurrent upgrades to avoid overloading the orchestrator.
-     */
-    private static final int MAX_ACTIVE_UPGRADES = 30;
-
-    private final NodeRepository nodeRepository;
     private final CuratorDatabaseClient db;
-    private final int maxActiveUpgrades;
+    private final Upgrader upgrader;
 
     public OsVersions(NodeRepository nodeRepository) {
-        this(nodeRepository, MAX_ACTIVE_UPGRADES);
+        this(nodeRepository, new DelegatingUpgrader(nodeRepository, 30));
     }
 
-    OsVersions(NodeRepository nodeRepository, int maxActiveUpgrades) {
-        this.nodeRepository = Objects.requireNonNull(nodeRepository, "nodeRepository must be non-null");
+    OsVersions(NodeRepository nodeRepository, Upgrader upgrader) {
         this.db = nodeRepository.database();
-        this.maxActiveUpgrades = maxActiveUpgrades;
+        this.upgrader = upgrader;
 
         // Read and write all versions to make sure they are stored in the latest version of the serialized format
         try (var lock = db.lockOsVersions()) {
@@ -72,7 +61,7 @@ public class OsVersions {
         try (Lock lock = db.lockOsVersions()) {
             var osVersions = db.readOsVersions();
             osVersions.remove(nodeType);
-            disableUpgrade(nodeType);
+            upgrader.disableUpgrade(nodeType);
             db.writeOsVersions(osVersions);
         }
     }
@@ -111,36 +100,13 @@ public class OsVersions {
             var currentVersion = osVersions.get(nodeType);
             if (currentVersion == null) return; // No target version set for this type
             if (active) {
-                upgrade(nodeType, currentVersion);
+                upgrader.upgrade(nodeType, currentVersion);
             } else {
-                disableUpgrade(nodeType);
+                upgrader.disableUpgrade(nodeType);
             }
         }
     }
 
-    /** Trigger upgrade of nodes of given type*/
-    private void upgrade(NodeType type, Version version) {
-        var activeNodes = nodeRepository.list().nodeType(type).state(Node.State.active);
-        var numberToUpgrade = Math.max(0, maxActiveUpgrades - activeNodes.changingOsVersionTo(version).size());
-        var nodesToUpgrade = activeNodes.not().changingOsVersionTo(version)
-                                        .not().onOsVersion(version)
-                                        .byIncreasingOsVersion()
-                                        .first(numberToUpgrade);
-        if (nodesToUpgrade.size() == 0) return;
-        log.info("Upgrading " + nodesToUpgrade.size() + " nodes of type " + type + " to OS version " + version.toFullString());
-        nodeRepository.upgradeOs(NodeListFilter.from(nodesToUpgrade.asList()), Optional.of(version));
-    }
-
-    /** Disable OS upgrade for all nodes of given type */
-    private void disableUpgrade(NodeType type) {
-        var nodesUpgrading = nodeRepository.list()
-                                           .nodeType(type)
-                                           .changingOsVersion();
-        if (nodesUpgrading.size() == 0) return;
-        log.info("Disabling OS upgrade of all " + type + " nodes");
-        nodeRepository.upgradeOs(NodeListFilter.from(nodesUpgrading.asList()), Optional.empty());
-    }
-
     private static void require(NodeType nodeType) {
         if (!nodeType.isDockerHost()) {
             throw new IllegalArgumentException("Node type '" + nodeType + "' does not support OS upgrades");
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/Upgrader.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/Upgrader.java
new file mode 100644
index 00000000000..9352871c7a6
--- /dev/null
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/Upgrader.java
@@ -0,0 +1,20 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.provision.os;
+
+import com.yahoo.component.Version;
+import com.yahoo.config.provision.NodeType;
+
+/**
+ * Interface for an OS upgrader.
+ *
+ * @author mpolden
+ */
+public interface Upgrader {
+
+    /** Trigger upgrade of nodes of given type */
+    void upgrade(NodeType type, Version version);
+
+    /** Disable OS upgrade for all nodes of given type */
+    void disableUpgrade(NodeType type);
+
+}
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 5f4bde85c88..fcc5b06b48d 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
@@ -35,7 +35,7 @@ public class OsVersionsTest {
 
     @Test
     public void test_versions() {
-        var versions = new OsVersions(tester.nodeRepository(), Integer.MAX_VALUE);
+        var versions = new OsVersions(tester.nodeRepository(), new DelegatingUpgrader(tester.nodeRepository(), Integer.MAX_VALUE));
         tester.makeReadyNodes(10, "default", NodeType.host);
         tester.prepareAndActivateInfraApplication(infraApplication, NodeType.host);
         Supplier<List<Node>> hostNodes = () -> tester.nodeRepository().getNodes(NodeType.host);
@@ -80,7 +80,7 @@ public class OsVersionsTest {
     public void test_max_active_upgrades() {
         int totalNodes = 20;
         int maxActiveUpgrades = 5;
-        var versions = new OsVersions(tester.nodeRepository(), maxActiveUpgrades);
+        var versions = new OsVersions(tester.nodeRepository(), new DelegatingUpgrader(tester.nodeRepository(), maxActiveUpgrades));
         tester.makeReadyNodes(totalNodes, "default", NodeType.host);
         Supplier<NodeList> hostNodes = () -> tester.nodeRepository().list().state(Node.State.active).nodeType(NodeType.host);
         tester.prepareAndActivateInfraApplication(infraApplication, NodeType.host);
@@ -127,7 +127,7 @@ public class OsVersionsTest {
 
     @Test
     public void test_newer_upgrade_aborts_upgrade_to_stale_version() {
-        var versions = new OsVersions(tester.nodeRepository(), Integer.MAX_VALUE);
+        var versions = new OsVersions(tester.nodeRepository(), new DelegatingUpgrader(tester.nodeRepository(), Integer.MAX_VALUE));
         tester.makeReadyNodes(10, "default", NodeType.host);
         tester.prepareAndActivateInfraApplication(infraApplication, NodeType.host);
         Supplier<NodeList> hostNodes = () -> tester.nodeRepository().list().nodeType(NodeType.host);