Support partial infrastructure upgrades

3 files changed, 117 insertions, 18 deletions

diff --git a/config-provisioning/src/main/java/com/yahoo/config/provision/zone/NodeSlice.java b/config-provisioning/src/main/java/com/yahoo/config/provision/zone/NodeSlice.java
new file mode 100644
index 00000000000..f7f6bd10603
--- /dev/null
+++ b/config-provisioning/src/main/java/com/yahoo/config/provision/zone/NodeSlice.java
@@ -0,0 +1,46 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.config.provision.zone;
+
+import java.util.Objects;
+import java.util.OptionalDouble;
+import java.util.OptionalLong;
+
+/**
+ * A slice of nodes, satisfied by either a minimum count or a fraction.
+ *
+ * @author mpolden
+ */
+public record NodeSlice(OptionalDouble fraction, OptionalLong minCount) {
+
+    public static final NodeSlice ALL = minCount(Long.MAX_VALUE);
+
+    public NodeSlice {
+        Objects.requireNonNull(fraction);
+        Objects.requireNonNull(minCount);
+        if (fraction.isEmpty() == minCount.isEmpty()) {
+            throw new IllegalArgumentException("Exactly one of 'fraction' or 'minCount' must be set");
+        }
+        if (fraction.isPresent() && fraction.getAsDouble() > 1.0D) {
+            throw new IllegalArgumentException("Fraction must be <= 1.0, got " + fraction.getAsDouble());
+        }
+    }
+
+    /** Returns whether this slice is satisfied by given node count, out of totalCount */
+    public boolean satisfiedBy(long count, long totalCount) {
+        if (fraction.isPresent()) {
+            return count >= totalCount * fraction.getAsDouble();
+        }
+        return count >= Math.min(minCount.orElse(0), totalCount);
+    }
+
+    /** Returns a slice matching the given fraction of nodes */
+    public static NodeSlice fraction(double fraction) {
+        return new NodeSlice(OptionalDouble.of(fraction), OptionalLong.empty());
+    }
+
+    /** Returns a slice matching the given minimum number of nodes */
+    public static NodeSlice minCount(long count) {
+        return new NodeSlice(OptionalDouble.empty(), OptionalLong.of(count));
+    }
+
+}
diff --git a/config-provisioning/src/main/java/com/yahoo/config/provision/zone/UpgradePolicy.java b/config-provisioning/src/main/java/com/yahoo/config/provision/zone/UpgradePolicy.java
index 85cc384660d..1c5bfad4c47 100644
--- a/config-provisioning/src/main/java/com/yahoo/config/provision/zone/UpgradePolicy.java
+++ b/config-provisioning/src/main/java/com/yahoo/config/provision/zone/UpgradePolicy.java
@@ -4,6 +4,7 @@ package com.yahoo.config.provision.zone;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 import java.util.Set;
 
 /**
@@ -12,14 +13,13 @@ import java.util.Set;
  *
  * @author mpolden
  */
-public class UpgradePolicy {
+public record UpgradePolicy(List<Step> steps) {
 
-    private final List<Set<ZoneApi>> steps;
-
-    private UpgradePolicy(List<Set<ZoneApi>> steps) {
+    public UpgradePolicy(List<Step> steps) {
+        Objects.requireNonNull(steps);
         for (int i = 0; i < steps.size(); i++) {
             for (int j = 0; j < i; j++) {
-                if (!Collections.disjoint(steps.get(i), steps.get(j))) {
+                if (!Collections.disjoint(steps.get(i).zones(), steps.get(j).zones())) {
                     throw new IllegalArgumentException("One or more zones are declared in multiple steps");
                 }
             }
@@ -27,14 +27,9 @@ public class UpgradePolicy {
         this.steps = List.copyOf(steps);
     }
 
-    /** Returns the steps in this */
-    public List<Set<ZoneApi>> steps() {
-        return steps;
-    }
-
     /** Returns a copy of this with the step order inverted */
     public UpgradePolicy inverted() {
-        List<Set<ZoneApi>> copy = new ArrayList<>(steps);
+        List<Step> copy = new ArrayList<>(steps);
         Collections.reverse(copy);
         return new UpgradePolicy(copy);
     }
@@ -43,21 +38,23 @@ public class UpgradePolicy {
         return new UpgradePolicy.Builder();
     }
 
-    public static class Builder {
+    public record Builder(List<Step> steps) {
 
-        private final List<Set<ZoneApi>> steps = new ArrayList<>();
+        private Builder() {
+            this(new ArrayList<>());
+        }
 
-        private Builder() {}
+        public Builder upgrade(Step step) {
+            this.steps.add(step);
+            return this;
+        }
 
-        /** Upgrade given zone as the next step */
         public Builder upgrade(ZoneApi zone) {
             return upgradeInParallel(zone);
         }
 
-        /** Upgrade given zones in parallel as the next step */
         public Builder upgradeInParallel(ZoneApi... zone) {
-            this.steps.add(Set.of(zone));
-            return this;
+            return upgrade(Step.of(zone));
         }
 
         public UpgradePolicy build() {
@@ -66,4 +63,28 @@ public class UpgradePolicy {
 
     }
 
+    /**
+     * An upgrade step, consisting of one or more zones. If a step contains multiple zones, those will be upgraded in
+     * parallel.
+     */
+    public record Step(Set<ZoneApi> zones, NodeSlice nodeSlice) {
+
+        public Step(Set<ZoneApi> zones, NodeSlice nodeSlice) {
+            if (zones.isEmpty()) throw new IllegalArgumentException("A step must contain at least one zone");
+            this.zones = Set.copyOf(Objects.requireNonNull(zones));
+            this.nodeSlice = Objects.requireNonNull(nodeSlice);
+        }
+
+        /** Create a step for given zones, which requires all nodes to complete upgrade */
+        public static Step of(ZoneApi... zone) {
+            return new Step(Set.of(zone), NodeSlice.ALL);
+        }
+
+        /** Returns a copy of this step, requiring only the given slice of nodes for each zone in this step to upgrade */
+        public Step require(NodeSlice slice) {
+            return new Step(zones, slice);
+        }
+
+    }
+
 }
diff --git a/config-provisioning/src/test/java/com/yahoo/config/provision/zone/NodeSliceTest.java b/config-provisioning/src/test/java/com/yahoo/config/provision/zone/NodeSliceTest.java
new file mode 100644
index 00000000000..b13eba5fa32
--- /dev/null
+++ b/config-provisioning/src/test/java/com/yahoo/config/provision/zone/NodeSliceTest.java
@@ -0,0 +1,32 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.config.provision.zone;
+
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/**
+ * @author mpolden
+ */
+public class NodeSliceTest {
+
+    @Test
+    void node_slice() {
+        NodeSlice fraction = NodeSlice.fraction(0.6);
+        assertFalse(fraction.satisfiedBy(0, 4));
+        assertFalse(fraction.satisfiedBy(1, 4));
+        assertFalse(fraction.satisfiedBy(2, 4));
+        assertTrue(fraction.satisfiedBy(3, 4));
+        assertTrue(fraction.satisfiedBy(4, 4));
+        assertTrue(fraction.satisfiedBy(5, 4));
+
+        NodeSlice fixed = NodeSlice.minCount(5);
+        assertFalse(fixed.satisfiedBy(0, 5));
+        assertFalse(fixed.satisfiedBy(4, 5));
+        assertTrue(fixed.satisfiedBy(3, 3));
+        assertTrue(fixed.satisfiedBy(5, 5));
+        assertTrue(fixed.satisfiedBy(6, 5));
+    }
+
+}