Apply capacity to quota check

15 files changed, 88 insertions, 22 deletions

diff --git a/config-model-api/abi-spec.json b/config-model-api/abi-spec.json
index 42e7e23dfcc..c0139c87349 100644
--- a/config-model-api/abi-spec.json
+++ b/config-model-api/abi-spec.json
@@ -1624,6 +1624,9 @@
     ],
     "methods" : [
       "public void <init>()",
+      "public void add(com.yahoo.config.provision.ClusterSpec, com.yahoo.config.provision.Capacity)",
+      "public java.util.Map clusters()",
+      "public java.util.Map capacities()",
       "public void add(com.yahoo.config.provision.ClusterSpec$Id, com.yahoo.config.provision.Capacity)",
       "public java.util.Map all()"
     ],
diff --git a/config-model-api/src/main/java/com/yahoo/config/model/api/Provisioned.java b/config-model-api/src/main/java/com/yahoo/config/model/api/Provisioned.java
index db0822c4c10..81e0cb58c7d 100644
--- a/config-model-api/src/main/java/com/yahoo/config/model/api/Provisioned.java
+++ b/config-model-api/src/main/java/com/yahoo/config/model/api/Provisioned.java
@@ -16,13 +16,28 @@ import java.util.Map;
  */
 public class Provisioned {
 
-    private final Map<ClusterSpec.Id, Capacity> provisioned = new HashMap<>();
+    private final Map<ClusterSpec.Id, ClusterSpec> clusters = new HashMap<>();
 
+    private final Map<ClusterSpec.Id, Capacity> capacities = new HashMap<>();
+
+    public void add(ClusterSpec cluster, Capacity capacity) {
+        clusters.put(cluster.id(), cluster);
+        capacities.put(cluster.id(), capacity);
+    }
+
+    /** Returns an unmodifiable map of all the cluster requests recorded during build of the model this belongs to */
+    public Map<ClusterSpec.Id, ClusterSpec> clusters() { return Collections.unmodifiableMap(clusters); }
+
+    /** Returns an unmodifiable map of all the capacity provision requests recorded during build of the model this belongs to */
+    public Map<ClusterSpec.Id, Capacity> capacities() { return Collections.unmodifiableMap(capacities); }
+
+    // TODO: Remove after June 2024
     public void add(ClusterSpec.Id id, Capacity capacity) {
-        provisioned.put(id, capacity);
+        capacities.put(id, capacity);
     }
 
     /** Returns an unmodifiable map of all the provision requests recorded during build of the model this belongs to */
-    public Map<ClusterSpec.Id, Capacity> all() { return Collections.unmodifiableMap(provisioned); }
+    // TODO: Remove after June 2024
+    public Map<ClusterSpec.Id, Capacity> all() { return Collections.unmodifiableMap(capacities); }
 
 }
diff --git a/config-model/pom.xml b/config-model/pom.xml
index a0bda542d5f..751d3de63ef 100644
--- a/config-model/pom.xml
+++ b/config-model/pom.xml
@@ -26,6 +26,12 @@
     </dependency>
     <dependency>
       <groupId>com.yahoo.vespa</groupId>
+      <artifactId>flags</artifactId>
+      <version>${project.version}</version>
+      <scope>provided</scope>
+    </dependency>
+    <dependency>
+      <groupId>com.yahoo.vespa</groupId>
       <artifactId>provided-dependencies</artifactId>
       <version>${project.version}</version>
       <scope>provided</scope>
diff --git a/config-model/src/main/java/com/yahoo/config/model/deploy/DeployState.java b/config-model/src/main/java/com/yahoo/config/model/deploy/DeployState.java
index 57a75bd8a38..0e4e296f5a1 100644
--- a/config-model/src/main/java/com/yahoo/config/model/deploy/DeployState.java
+++ b/config-model/src/main/java/com/yahoo/config/model/deploy/DeployState.java
@@ -39,6 +39,8 @@ import com.yahoo.vespa.config.ConfigDefinition;
 import com.yahoo.vespa.config.ConfigDefinitionBuilder;
 import com.yahoo.vespa.config.ConfigDefinitionKey;
 import com.yahoo.vespa.documentmodel.DocumentModel;
+import com.yahoo.vespa.flags.FlagSource;
+import com.yahoo.vespa.flags.InMemoryFlagSource;
 import com.yahoo.vespa.model.container.search.QueryProfiles;
 import com.yahoo.vespa.model.container.search.QueryProfilesBuilder;
 import com.yahoo.vespa.model.container.search.SemanticRules;
@@ -79,6 +81,7 @@ public class DeployState implements ConfigDefinitionStore {
     private final Version vespaVersion;
     private final Set<ContainerEndpoint> endpoints;
     private final Zone zone; // TODO: Zone is set separately both here and in properties
+    private final FlagSource flagSource;
     private final QueryProfiles queryProfiles;
     private final SemanticRules semanticRules;
     private final ImportedMlModels importedModels;
@@ -118,6 +121,7 @@ public class DeployState implements ConfigDefinitionStore {
                         Set<ContainerEndpoint> endpoints,
                         Collection<MlModelImporter> modelImporters,
                         Zone zone,
+                        FlagSource flagSource,
                         QueryProfiles queryProfiles,
                         SemanticRules semanticRules,
                         Instant now,
@@ -143,6 +147,7 @@ public class DeployState implements ConfigDefinitionStore {
         this.configDefinitionRepo = configDefinitionRepo;
         this.endpoints = Set.copyOf(endpoints);
         this.zone = zone;
+        this.flagSource = flagSource;
         this.queryProfiles = queryProfiles; // TODO: Remove this by seeing how pagetemplates are propagated
         this.semanticRules = semanticRules; // TODO: Remove this by seeing how pagetemplates are propagated
         this.importedModels = importMlModels(applicationPackage, modelImporters, executor);
@@ -273,6 +278,8 @@ public class DeployState implements ConfigDefinitionStore {
     /** Returns the zone in which this is currently running */
     public Zone zone() { return zone; }
 
+    public FlagSource flagSource() { return flagSource; }
+
     public QueryProfiles getQueryProfiles() { return queryProfiles; }
 
     public SemanticRules getSemanticRules() { return semanticRules; }
@@ -330,6 +337,7 @@ public class DeployState implements ConfigDefinitionStore {
         private Set<ContainerEndpoint> endpoints = Set.of();
         private Collection<MlModelImporter> modelImporters = List.of();
         private Zone zone = Zone.defaultZone();
+        private FlagSource flagSource = new InMemoryFlagSource();
         private Instant now = Instant.now();
         private Version wantedNodeVespaVersion = Vtag.currentVersion;
         private boolean accessLoggingEnabledByDefault = true;
@@ -407,6 +415,11 @@ public class DeployState implements ConfigDefinitionStore {
             return this;
         }
 
+        public Builder flagSource(FlagSource flagSource) {
+            this.flagSource = flagSource;
+            return this;
+        }
+
         public Builder now(Instant now) {
             this.now = now;
             return this;
@@ -483,6 +496,7 @@ public class DeployState implements ConfigDefinitionStore {
                                    endpoints,
                                    modelImporters,
                                    zone,
+                                   flagSource,
                                    queryProfiles,
                                    semanticRules,
                                    now,
diff --git a/config-model/src/main/java/com/yahoo/config/model/provision/InMemoryProvisioner.java b/config-model/src/main/java/com/yahoo/config/model/provision/InMemoryProvisioner.java
index befe57a97e4..da0fd265724 100644
--- a/config-model/src/main/java/com/yahoo/config/model/provision/InMemoryProvisioner.java
+++ b/config-model/src/main/java/com/yahoo/config/model/provision/InMemoryProvisioner.java
@@ -158,7 +158,7 @@ public class InMemoryProvisioner implements HostProvisioner {
 
     @Override
     public List<HostSpec> prepare(ClusterSpec cluster, Capacity requested, ProvisionLogger logger) {
-        provisioned.add(cluster.id(), requested);
+        provisioned.add(cluster, requested);
         clusters.add(cluster);
         if (environment == Environment.dev && ! requested.isRequired()) {
             requested = requested.withLimits(requested.minResources().withNodes(1),
diff --git a/config-model/src/main/java/com/yahoo/vespa/model/VespaModelFactory.java b/config-model/src/main/java/com/yahoo/vespa/model/VespaModelFactory.java
index c876976917b..efe83fb4e91 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/VespaModelFactory.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/VespaModelFactory.java
@@ -28,6 +28,8 @@ import com.yahoo.config.provision.QuotaExceededException;
 import com.yahoo.config.provision.TransientException;
 import com.yahoo.config.provision.Zone;
 import com.yahoo.vespa.config.VespaVersion;
+import com.yahoo.vespa.flags.FlagSource;
+import com.yahoo.vespa.flags.InMemoryFlagSource;
 import com.yahoo.vespa.model.application.validation.Validation;
 import com.yahoo.vespa.model.application.validation.Validator;
 import org.xml.sax.SAXException;
@@ -52,6 +54,7 @@ public class VespaModelFactory implements ModelFactory {
     private final ConfigModelRegistry configModelRegistry;
     private final Collection<MlModelImporter> modelImporters;
     private final Zone zone;
+    private final FlagSource flagSource;
     private final Clock clock;
     private final Version version;
     private final List<Validator> additionalValidators;
@@ -60,7 +63,7 @@ public class VespaModelFactory implements ModelFactory {
     @Inject
     public VespaModelFactory(ComponentRegistry<ConfigModelPlugin> pluginRegistry,
                              ComponentRegistry<Validator> additionalValidators,
-                             Zone zone) {
+                             Zone zone, FlagSource flagSource) {
         this.version = new Version(VespaVersion.major, VespaVersion.minor, VespaVersion.micro);
         List<ConfigModelBuilder<?>> modelBuilders = new ArrayList<>();
         for (ConfigModelPlugin plugin : pluginRegistry.allComponents()) {
@@ -76,6 +79,7 @@ public class VespaModelFactory implements ModelFactory {
                 new XGBoostImporter(),
                 new LightGBMImporter());
         this.zone = zone;
+        this.flagSource = flagSource;
         this.additionalValidators = List.copyOf(additionalValidators.allComponents());
 
         this.clock = Clock.systemUTC();
@@ -84,7 +88,7 @@ public class VespaModelFactory implements ModelFactory {
     // For testing only
     protected VespaModelFactory(ConfigModelRegistry configModelRegistry) {
         this(new Version(VespaVersion.major, VespaVersion.minor, VespaVersion.micro), configModelRegistry,
-                Clock.systemUTC(), Zone.defaultZone());
+             Clock.systemUTC(), Zone.defaultZone());
     }
 
     private VespaModelFactory(Version version, ConfigModelRegistry configModelRegistry, Clock clock, Zone zone) {
@@ -98,6 +102,7 @@ public class VespaModelFactory implements ModelFactory {
         this.modelImporters = List.of();
         this.additionalValidators = List.of();
         this.zone = zone;
+        this.flagSource = new InMemoryFlagSource();
         this.clock = clock;
     }
 
@@ -192,6 +197,7 @@ public class VespaModelFactory implements ModelFactory {
             .endpoints(modelContext.properties().endpoints())
             .modelImporters(modelImporters)
             .zone(zone)
+            .flagSource(flagSource)
             .now(clock.instant())
             .wantedNodeVespaVersion(modelContext.wantedNodeVespaVersion())
             .wantedDockerImageRepo(modelContext.wantedDockerImageRepo())
diff --git a/config-model/src/main/java/com/yahoo/vespa/model/application/validation/QuotaValidator.java b/config-model/src/main/java/com/yahoo/vespa/model/application/validation/QuotaValidator.java
index 4d9386b5f19..0984770ef49 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/application/validation/QuotaValidator.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/application/validation/QuotaValidator.java
@@ -1,9 +1,12 @@
 // Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.model.application.validation;
 
+import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.Capacity;
+import com.yahoo.config.provision.CapacityPolicies;
 import com.yahoo.config.provision.ClusterResources;
 import com.yahoo.config.provision.ClusterSpec;
+import com.yahoo.config.provision.Exclusivity;
 import com.yahoo.config.provision.NodeResources;
 import com.yahoo.config.provision.QuotaExceededException;
 import com.yahoo.config.provision.SystemName;
@@ -31,25 +34,34 @@ public class QuotaValidator implements Validator {
 
     @Override
     public void validate(Context context) {
+        var zone = context.deployState().zone();
+        var flagSource = context.deployState().flagSource();
+        var capacityPolicies = new CapacityPolicies(zone, new Exclusivity(zone, flagSource), flagSource);
         var quota = context.deployState().getProperties().quota();
         quota.maxClusterSize().ifPresent(maxClusterSize -> validateMaxClusterSize(maxClusterSize, context.model()));
-        quota.budgetAsDecimal().ifPresent(budget -> validateBudget(budget, context.model(), context.deployState().getProperties().zone()));
+        quota.budgetAsDecimal().ifPresent(budget -> validateBudget(budget, context, capacityPolicies));
     }
 
-    private void validateBudget(BigDecimal budget, VespaModel model, Zone zone) {
-        var maxSpend = model.allClusters().stream()
-                .filter(id -> !adminClusterIds(model).contains(id))
-                .map(id -> model.provisioned().all().getOrDefault(id, zeroCapacity))
-                .mapToDouble(c -> c.maxResources().cost()) // TODO: This may be unspecified -> 0
-                .sum();
+    private void validateBudget(BigDecimal budget, Context context,
+                                CapacityPolicies capacityPolicies) {
+        var zone = context.deployState().getProperties().zone();
+        var application = context.model().applicationPackage().getApplicationId();
+
+        var maxSpend = 0.0;
+        for (var id : context.model().allClusters()) {
+            if (adminClusterIds(context.model()).contains(id)) continue;
+            var cluster = context.model().provisioned().clusters().get(id);
+            var capacity = context.model().provisioned().capacities().getOrDefault(id, zeroCapacity);
+            maxSpend += capacityPolicies.applyOn(capacity, application, cluster.isExclusive()).maxResources().cost();
+        }
 
-        var actualSpend = model.allocatedHosts().getHosts().stream()
+        var actualSpend = context.model().allocatedHosts().getHosts().stream()
                          .filter(hostSpec -> hostSpec.membership().get().cluster().type() != ClusterSpec.Type.admin)
                          .mapToDouble(hostSpec -> hostSpec.advertisedResources().cost())
                          .sum();
 
         if (Math.abs(actualSpend) < 0.01) {
-            log.warning("Deploying application " + model.applicationPackage().getApplicationId() + " with zero budget use.  This is suspicious, but not blocked");
+            log.warning("Deploying application " + application + " with zero budget use.  This is suspicious, but not blocked");
             return;
         }
 
@@ -69,7 +81,7 @@ public class QuotaValidator implements Validator {
 
     /** Check that all clusters in the application do not exceed the quota max cluster size. */
     private void validateMaxClusterSize(int maxClusterSize, VespaModel model) {
-        var invalidClusters = model.provisioned().all().entrySet().stream()
+        var invalidClusters = model.provisioned().capacities().entrySet().stream()
                 .filter(entry -> entry.getValue() != null)
                 .filter(entry -> {
                     var cluster = entry.getValue();
diff --git a/config-model/src/main/java/com/yahoo/vespa/model/application/validation/change/ResourcesReductionValidator.java b/config-model/src/main/java/com/yahoo/vespa/model/application/validation/change/ResourcesReductionValidator.java
index 5d7a8779005..42410dc3acf 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/application/validation/change/ResourcesReductionValidator.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/application/validation/change/ResourcesReductionValidator.java
@@ -60,9 +60,9 @@ public class ResourcesReductionValidator implements ChangeValidator {
      * This will always yield specified node resources on hosted instances and never on self-hosted instances.
      */
     private ClusterResources clusterResources(ClusterSpec.Id id, VespaModel model) {
-        if ( ! model.provisioned().all().containsKey(id)) return null;
+        if ( ! model.provisioned().capacities().containsKey(id)) return null;
 
-        ClusterResources resources = model.provisioned().all().get(id).maxResources();
+        ClusterResources resources = model.provisioned().capacities().get(id).maxResources();
         if ( ! resources.nodeResources().isUnspecified()) return resources;
 
         var containerCluster = model.getContainerClusters().get(id.value());
diff --git a/config-model/src/test/java/com/yahoo/vespa/model/ClusterInfoTest.java b/config-model/src/test/java/com/yahoo/vespa/model/ClusterInfoTest.java
index 4df9f261dfe..7aa6eb76995 100644
--- a/config-model/src/test/java/com/yahoo/vespa/model/ClusterInfoTest.java
+++ b/config-model/src/test/java/com/yahoo/vespa/model/ClusterInfoTest.java
@@ -258,7 +258,7 @@ public class ClusterInfoTest {
                                   .provisioned(provisioner.provisioned())
                                   .build();
         new VespaModel(new NullConfigModelRegistry(), deployState);
-        return deployState.provisioned().all();
+        return deployState.provisioned().capacities();
     }
 
 }
diff --git a/config-model/src/test/java/com/yahoo/vespa/model/application/validation/QuotaValidatorTest.java b/config-model/src/test/java/com/yahoo/vespa/model/application/validation/QuotaValidatorTest.java
index 89f81dfdaef..590433757c3 100644
--- a/config-model/src/test/java/com/yahoo/vespa/model/application/validation/QuotaValidatorTest.java
+++ b/config-model/src/test/java/com/yahoo/vespa/model/application/validation/QuotaValidatorTest.java
@@ -54,6 +54,12 @@ public class QuotaValidatorTest {
     }
 
     @Test
+    void test_deploy_within_quota_budget_because_in_dev() {
+        var tester = new ValidationTester(13, false, new TestProperties().setHostedVespa(true).setQuota(quota).setZone(devZone));
+        tester.deploy(null, getServices(10), Environment.dev, null, CONTAINER_CLUSTER);
+    }
+
+    @Test
     void test_deploy_above_quota_budget_in_publiccd() {
         var tester = new ValidationTester(13, false, new TestProperties().setHostedVespa(true).setQuota(quota.withBudget(BigDecimal.ONE)).setZone(publicCdZone));
         try {
diff --git a/config-model/src/test/java/com/yahoo/vespa/model/container/xml/ContainerModelBuilderTest.java b/config-model/src/test/java/com/yahoo/vespa/model/container/xml/ContainerModelBuilderTest.java
index e704da08d18..0e8ce4748b4 100644
--- a/config-model/src/test/java/com/yahoo/vespa/model/container/xml/ContainerModelBuilderTest.java
+++ b/config-model/src/test/java/com/yahoo/vespa/model/container/xml/ContainerModelBuilderTest.java
@@ -598,7 +598,7 @@ public class ContainerModelBuilderTest extends ContainerModelBuilderTestBase {
                                                 .setCloudAccount(cloudAccount))
                 .build());
         assertEquals(2, model.hostSystem().getHosts().size());
-        assertEquals(List.of(cloudAccount), model.provisioned().all().values()
+        assertEquals(List.of(cloudAccount), model.provisioned().capacities().values()
                                                  .stream()
                                                  .map(capacity -> capacity.cloudAccount().get())
                                                  .toList());
diff --git a/configserver-flags/src/main/java/com/yahoo/vespa/configserver/flags/ConfigServerFlagSource.java b/configserver-flags/src/main/java/com/yahoo/vespa/configserver/flags/ConfigServerFlagSource.java
index f30d329cf09..7a0b3695e7c 100644
--- a/configserver-flags/src/main/java/com/yahoo/vespa/configserver/flags/ConfigServerFlagSource.java
+++ b/configserver-flags/src/main/java/com/yahoo/vespa/configserver/flags/ConfigServerFlagSource.java
@@ -13,6 +13,7 @@ import java.nio.file.FileSystems;
  * @author hakonhall
  */
 public class ConfigServerFlagSource extends OrderedFlagSource {
+
     @Inject
     public ConfigServerFlagSource(FlagsDb flagsDb) {
         this(FileSystems.getDefault(), flagsDb);
@@ -21,4 +22,5 @@ public class ConfigServerFlagSource extends OrderedFlagSource {
     ConfigServerFlagSource(FileSystem fileSystem, FlagsDb flagsDb) {
         super(new BootstrapFlagSource(fileSystem), new ZooKeeperFlagSource(flagsDb));
     }
+
 }
diff --git a/configserver-flags/src/main/java/com/yahoo/vespa/configserver/flags/db/ZooKeeperFlagSource.java b/configserver-flags/src/main/java/com/yahoo/vespa/configserver/flags/db/ZooKeeperFlagSource.java
index f68c4657d8a..e57088416a0 100644
--- a/configserver-flags/src/main/java/com/yahoo/vespa/configserver/flags/db/ZooKeeperFlagSource.java
+++ b/configserver-flags/src/main/java/com/yahoo/vespa/configserver/flags/db/ZooKeeperFlagSource.java
@@ -13,6 +13,7 @@ import java.util.Optional;
  * @author hakonhall
  */
 public class ZooKeeperFlagSource implements FlagSource {
+
     private final FlagsDb flagsDb;
 
     public ZooKeeperFlagSource(FlagsDb flagsDb) {
@@ -23,4 +24,5 @@ public class ZooKeeperFlagSource implements FlagSource {
     public Optional<RawFlag> fetch(FlagId id, FetchVector vector) {
         return flagsDb.getValue(id).flatMap(data -> data.resolve(vector));
     }
+
 }
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/ApplicationRepository.java b/configserver/src/main/java/com/yahoo/vespa/config/server/ApplicationRepository.java
index 07a8179732f..22199dfc5b4 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/ApplicationRepository.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/ApplicationRepository.java
@@ -1136,7 +1136,7 @@ public class ApplicationRepository implements com.yahoo.config.provision.Deploye
 
     public double getQuotaUsageRate(ApplicationId applicationId) {
         var application = getApplication(applicationId);
-        return application.getModel().provisioned().all().values().stream()
+        return application.getModel().provisioned().capacities().values().stream()
                 .map(Capacity::maxResources)// TODO: This may be unspecified -> 0
                 .mapToDouble(resources -> resources.nodes() * resources.nodeResources().cost())
                 .sum();
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/provision/ProvisionerAdapter.java b/configserver/src/main/java/com/yahoo/vespa/config/server/provision/ProvisionerAdapter.java
index 956573ecbd4..27805b64e79 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/provision/ProvisionerAdapter.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/provision/ProvisionerAdapter.java
@@ -44,7 +44,7 @@ public class ProvisionerAdapter implements HostProvisioner {
 
     @Override
     public List<HostSpec> prepare(ClusterSpec cluster, Capacity capacity, ProvisionLogger logger) {
-        provisioned.add(cluster.id(), capacity);
+        provisioned.add(cluster, capacity);
         return provisioner.prepare(applicationId, cluster, capacity, logger);
     }