Add method to get allocation deficit

2 files changed, 50 insertions, 7 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeAllocation.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeAllocation.java
index 89412fe6350..dcc3c4a0ef8 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeAllocation.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeAllocation.java
@@ -4,6 +4,7 @@ package com.yahoo.vespa.hosted.provision.provisioning;
 import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.ClusterMembership;
 import com.yahoo.config.provision.ClusterSpec;
+import com.yahoo.config.provision.Flavor;
 import com.yahoo.config.provision.SystemName;
 import com.yahoo.config.provision.TenantName;
 import com.yahoo.config.provision.Zone;
@@ -273,6 +274,22 @@ class NodeAllocation {
     }
 
     /**
+     * Returns {@link FlavorCount} describing the docker node deficit for the given {@link NodeSpec}.
+     *
+     * @return empty if the requested spec is not count based or the requested flavor type is not docker or
+     * the request is already fulfilled. Otherwise returns {@link FlavorCount} containing the required flavor
+     * and node count to cover the deficit.
+     */
+    Optional<FlavorCount> getFulfilledDockerDeficit() {
+        return Optional.of(requestedNodes)
+                .filter(NodeSpec.CountNodeSpec.class::isInstance)
+                .map(NodeSpec.CountNodeSpec.class::cast)
+                .map(spec -> new FlavorCount(spec.getFlavor(), spec.fulfilledDeficitCount(acceptedOfRequestedFlavor)))
+                .filter(flavorCount -> flavorCount.getFlavor().getType() == Flavor.Type.DOCKER_CONTAINER)
+                .filter(flavorCount -> flavorCount.getCount() > 0);
+    }
+
+    /**
      * Make the number of <i>non-retired</i> nodes in the list equal to the requested number
      * of nodes, and retire the rest of the list. Only retire currently active nodes.
      * Prefer to retire nodes of the wrong flavor.
@@ -346,4 +363,21 @@ class NodeAllocation {
         return Comparator.comparing((PrioritizableNode n) -> n.node.allocation().get().membership().index());
     }
 
+    static class FlavorCount {
+        private final Flavor flavor;
+        private final int count;
+
+        private FlavorCount(Flavor flavor, int count) {
+            this.flavor = flavor;
+            this.count = count;
+        }
+
+        Flavor getFlavor() {
+            return flavor;
+        }
+
+        int getCount() {
+            return count;
+        }
+    }
 }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeSpec.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeSpec.java
index c0479a054c6..e033d994f24 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeSpec.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeSpec.java
@@ -38,7 +38,9 @@ public interface NodeSpec {
     boolean saturatedBy(int count);
 
     /** Returns whether the given node count is sufficient to fulfill this spec */
-    boolean fulfilledBy(int count);
+    default boolean fulfilledBy(int count) {
+        return fulfilledDeficitCount(count) == 0;
+    }
 
     /** Returns whether this should throw an exception if the requested nodes are not fully available */
     boolean canFail();
@@ -46,6 +48,9 @@ public interface NodeSpec {
     /** Returns the ideal number of nodes that should be retired to fulfill this spec */
     int idealRetiredCount(int acceptedCount, int currentRetiredCount);
 
+    /** Returns number of additional nodes needed for this spec to be fulfilled given the current node count */
+    int fulfilledDeficitCount(int count);
+
     /** Returns a specification of a fraction of all the nodes of this. It is assumed the argument is a valid divisor. */
     NodeSpec fraction(int divisor);
 
@@ -106,15 +111,17 @@ public interface NodeSpec {
         public boolean saturatedBy(int count) { return fulfilledBy(count); } // min=max for count specs
 
         @Override
-        public boolean fulfilledBy(int count) { return count >= this.count; }
-
-        @Override
         public boolean canFail() { return canFail; }
 
         @Override
         public int idealRetiredCount(int acceptedCount, int currentRetiredCount) { return acceptedCount - this.count; }
 
         @Override
+        public int fulfilledDeficitCount(int count) {
+            return Math.max(this.count - count, 0);
+        }
+
+        @Override
         public NodeSpec fraction(int divisor) {
             return new CountNodeSpec(count/divisor, requestedFlavor, exclusive, canFail);
         }
@@ -166,9 +173,6 @@ public interface NodeSpec {
         public boolean specifiesNonStockFlavor() { return false; }
 
         @Override
-        public boolean fulfilledBy(int count) { return true; }
-
-        @Override
         public boolean saturatedBy(int count) { return false; }
 
         @Override
@@ -185,6 +189,11 @@ public interface NodeSpec {
         }
 
         @Override
+        public int fulfilledDeficitCount(int count) {
+            return 0;
+        }
+
+        @Override
         public NodeSpec fraction(int divisor) { return this; }
 
         @Override