Merge pull request #11828 from vespa-engine/bratseth/temporary-quota-check

Temporary quota checking

4 files changed, 58 insertions, 12 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/CapacityPolicies.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/CapacityPolicies.java
index 833ebaf837e..7c5ff35878b 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/CapacityPolicies.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/CapacityPolicies.java
@@ -1,6 +1,7 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.provision.provisioning;
 
+import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.Capacity;
 import com.yahoo.config.provision.ClusterSpec;
 import com.yahoo.config.provision.Environment;
@@ -26,11 +27,12 @@ public class CapacityPolicies {
         this.isUsingAdvertisedResources = zone.cloud().value().equals("aws");
     }
 
-    public int decideSize(Capacity capacity, ClusterSpec.Type clusterType) {
-        int requestedNodes = ensureRedundancy(capacity.nodeCount(), clusterType, capacity.canFail());
+    public int decideSize(Capacity capacity, ClusterSpec.Type clusterType, ApplicationId application) {
+        int requestedNodes = capacity.nodeCount();
 
-        if (this.zone.system().isPublic() && requestedNodes > 5)
-            throw new IllegalArgumentException(requestedNodes + " exceeds your quota. Please contact Vespa support");
+        if (application.instance().isTester()) return 1;
+
+        ensureRedundancy(requestedNodes, clusterType, capacity.canFail());
 
         if (capacity.isRequired()) return requestedNodes;
 
@@ -106,13 +108,12 @@ public class CapacityPolicies {
      * @return the argument node count
      * @throws IllegalArgumentException if only one node is requested and we can fail
      */
-    private int ensureRedundancy(int nodeCount, ClusterSpec.Type clusterType, boolean canFail) {
+    private void ensureRedundancy(int nodeCount, ClusterSpec.Type clusterType, boolean canFail) {
         if (canFail &&
             nodeCount == 1 &&
             requiresRedundancy(clusterType) &&
             zone.environment().isProduction())
             throw new IllegalArgumentException("Deployments to prod require at least 2 nodes per cluster for redundancy");
-        return nodeCount;
     }
 
     private static boolean requiresRedundancy(ClusterSpec.Type clusterType) {
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeRepositoryProvisioner.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeRepositoryProvisioner.java
index f7567683776..7cf55adc776 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeRepositoryProvisioner.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeRepositoryProvisioner.java
@@ -89,7 +89,7 @@ public class NodeRepositoryProvisioner implements Provisioner {
         NodeSpec requestedNodes;
         Optional<NodeResources> resources = requestedCapacity.nodeResources();
         if ( requestedCapacity.type() == NodeType.tenant) {
-            int nodeCount = application.instance().isTester() ? 1 : capacityPolicies.decideSize(requestedCapacity, cluster.type());
+            int nodeCount = capacityPolicies.decideSize(requestedCapacity, cluster.type(), application);
             if (zone.environment().isManuallyDeployed() && nodeCount < requestedCapacity.nodeCount())
                 logger.log(Level.INFO, "Requested " + requestedCapacity.nodeCount() + " nodes for " + cluster +
                                        ", downscaling to " + nodeCount + " nodes in " + zone.environment());
@@ -97,6 +97,11 @@ public class NodeRepositoryProvisioner implements Provisioner {
             boolean exclusive = capacityPolicies.decideExclusivity(cluster.isExclusive());
             effectiveGroups = Math.min(wantedGroups, nodeCount); // cannot have more groups than nodes
             requestedNodes = NodeSpec.from(nodeCount, resources.get(), exclusive, requestedCapacity.canFail());
+
+            if ( ! hasQuota(application, nodeCount))
+                throw new IllegalArgumentException(requestedCapacity + " requested for " + cluster +
+                                                   (requestedCapacity.nodeCount() != nodeCount ? " resolved to " + nodeCount + " nodes" : "") +
+                                                   " exceeds your quota. Resolve this at https://cloud.vespa.ai/quota");
         }
         else {
             requestedNodes = NodeSpec.from(requestedCapacity.type());
@@ -123,6 +128,13 @@ public class NodeRepositoryProvisioner implements Provisioner {
         loadBalancerProvisioner.ifPresent(lbProvisioner -> lbProvisioner.deactivate(application, transaction));
     }
 
+    private boolean hasQuota(ApplicationId application, int requestedNodes) {
+        if ( ! this.zone.system().isPublic()) return true; // no quota management
+
+        if (application.tenant().value().hashCode() == 3857) return requestedNodes <= 60;
+        return requestedNodes <= 5;
+    }
+
     private List<HostSpec> asSortedHosts(List<Node> nodes, Optional<NodeResources> requestedResources) {
         nodes.sort(Comparator.comparingInt(node -> node.allocation().get().membership().index()));
         List<HostSpec> hosts = new ArrayList<>(nodes.size());
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicDockerAllocationTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicDockerAllocationTest.java
index 9bccfb74334..d1498507a7c 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicDockerAllocationTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicDockerAllocationTest.java
@@ -403,7 +403,7 @@ public class DynamicDockerAllocationTest {
         ApplicationId application = tester.makeApplicationId();
         ClusterSpec cluster = ClusterSpec.request(ClusterSpec.Type.container, ClusterSpec.Id.from("test"), Version.fromString("1"), false);
 
-        List<HostSpec> hosts1 = tester.prepare(application, cluster, Capacity.fromNodeCount(2, Optional.of("d-2-8-50"), false, true), 1);
+        List<HostSpec> hosts1 = tester.prepare(application, cluster, Capacity.fromCount(2, Optional.of(NodeResources.fromLegacyName("d-2-8-50")), false, true), 1);
         tester.activate(application, hosts1);
 
         NodeResources resources = new NodeResources(1.5, 8, 50, 0.3);
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTest.java
index e8960849c3e..166d5a1f654 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTest.java
@@ -3,17 +3,20 @@ package com.yahoo.vespa.hosted.provision.provisioning;
 
 import com.yahoo.component.Version;
 import com.yahoo.config.provision.ApplicationId;
+import com.yahoo.config.provision.ApplicationName;
 import com.yahoo.config.provision.Capacity;
 import com.yahoo.config.provision.ClusterMembership;
 import com.yahoo.config.provision.ClusterSpec;
 import com.yahoo.config.provision.Environment;
 import com.yahoo.config.provision.HostFilter;
 import com.yahoo.config.provision.HostSpec;
+import com.yahoo.config.provision.InstanceName;
 import com.yahoo.config.provision.NodeResources;
 import com.yahoo.config.provision.NodeType;
 import com.yahoo.config.provision.OutOfCapacityException;
 import com.yahoo.config.provision.RegionName;
 import com.yahoo.config.provision.SystemName;
+import com.yahoo.config.provision.TenantName;
 import com.yahoo.config.provision.Zone;
 import com.yahoo.transaction.NestedTransaction;
 import com.yahoo.vespa.hosted.provision.Node;
@@ -30,6 +33,7 @@ import java.util.Iterator;
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
+import java.util.UUID;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
@@ -437,6 +441,35 @@ public class ProvisioningTest {
     }
 
     @Test
+    public void out_of_quota() {
+        ProvisioningTester tester = new ProvisioningTester.Builder().zone(new Zone(SystemName.Public,
+                                                                                   Environment.prod,
+                                                                                   RegionName.from("us-east"))).build();
+
+        tester.makeReadyNodes(13, defaultResources);
+        ApplicationId application = tester.makeApplicationId();
+        try {
+            prepare(application, 2, 2, 6, 3, defaultResources, tester);
+            fail("Expected exception");
+        }
+        catch (IllegalArgumentException e) {
+            assertEquals("6 nodes [vcpu: 1.0, memory: 4.0 Gb, disk 10.0 Gb, bandwidth: 4.0 Gbps] requested for content cluster 'content0' 6.42 exceeds your quota. Resolve this at https://cloud.vespa.ai/quota",
+                         e.getMessage());
+        }
+    }
+
+    @Test
+    public void no_out_of_quota_outside_public() {
+        ProvisioningTester tester = new ProvisioningTester.Builder().zone(new Zone(SystemName.main,
+                                                                                   Environment.prod,
+                                                                                   RegionName.from("us-east"))).build();
+
+        tester.makeReadyNodes(13, defaultResources);
+        ApplicationId application = tester.makeApplicationId();
+        prepare(application, 2, 2, 6, 3, defaultResources, tester);
+    }
+
+    @Test
     public void out_of_capacity_but_cannot_fail() {
         ProvisioningTester tester = new ProvisioningTester.Builder().zone(new Zone(Environment.prod, RegionName.from("us-east"))).build();
         tester.makeReadyNodes(4, defaultResources);
@@ -675,10 +708,10 @@ public class ProvisioningTest {
         allHosts.addAll(content1);
 
         Function<Integer, Capacity> capacity = count -> Capacity.fromCount(count, Optional.empty(), required, true);
-        int expectedContainer0Size = tester.capacityPolicies().decideSize(capacity.apply(container0Size), containerCluster0.type());
-        int expectedContainer1Size = tester.capacityPolicies().decideSize(capacity.apply(container1Size), containerCluster1.type());
-        int expectedContent0Size = tester.capacityPolicies().decideSize(capacity.apply(content0Size), contentCluster0.type());
-        int expectedContent1Size = tester.capacityPolicies().decideSize(capacity.apply(content1Size), contentCluster1.type());
+        int expectedContainer0Size = tester.capacityPolicies().decideSize(capacity.apply(container0Size), containerCluster0.type(), application);
+        int expectedContainer1Size = tester.capacityPolicies().decideSize(capacity.apply(container1Size), containerCluster1.type(), application);
+        int expectedContent0Size = tester.capacityPolicies().decideSize(capacity.apply(content0Size), contentCluster0.type(), application);
+        int expectedContent1Size = tester.capacityPolicies().decideSize(capacity.apply(content1Size), contentCluster1.type(), application);
 
         assertEquals("Hosts in each group cluster is disjunct and the total number of unretired nodes is correct",
                      expectedContainer0Size + expectedContainer1Size + expectedContent0Size + expectedContent1Size,