Merge remote-tracking branch 'origin/master' into ogronnesby/gpu-billing

10 files changed, 157 insertions, 20 deletions

diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/application/v4/model/ProtonMetrics.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/application/v4/model/SearchNodeMetrics.java
index 8d7daa44687..729c2488e2b 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/application/v4/model/ProtonMetrics.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/application/v4/model/SearchNodeMetrics.java
@@ -10,11 +10,11 @@ import java.util.Map;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
-public class ProtonMetrics {
+public class SearchNodeMetrics {
 
     private static final ObjectMapper jsonMapper = new ObjectMapper();
 
-    private static final Logger logger = Logger.getLogger(ProtonMetrics.class.getName());
+    private static final Logger logger = Logger.getLogger(SearchNodeMetrics.class.getName());
 
     public static final String DOCUMENTS_ACTIVE_COUNT = "documentsActiveCount";
     public static final String DOCUMENTS_READY_COUNT = "documentsReadyCount";
@@ -26,7 +26,7 @@ public class ProtonMetrics {
     private final String clusterId;
     private final Map<String, Double> metrics;
 
-    public ProtonMetrics(String clusterId) {
+    public SearchNodeMetrics(String clusterId) {
         this.clusterId = clusterId;
         metrics = new HashMap<>();
     }
@@ -45,7 +45,7 @@ public class ProtonMetrics {
 
     public double resourceMemoryUsageAverage() { return metrics.get(RESOURCE_MEMORY_USAGE_AVERAGE); }
 
-    public ProtonMetrics addMetric(String name, double value) {
+    public SearchNodeMetrics addMetric(String name, double value) {
         metrics.put(name, value);
         return this;
     }
diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/athenz/ZmsClientMock.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/athenz/ZmsClientMock.java
index ed185f8af32..94ed28778d7 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/athenz/ZmsClientMock.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/athenz/ZmsClientMock.java
@@ -257,6 +257,11 @@ public class ZmsClientMock implements ZmsClient {
     }
 
     @Override
+    public void updateProviderEndpoint(AthenzService athenzService, String endpoint) {
+
+    }
+
+    @Override
     public void deleteService(AthenzService athenzService) {
         athenz.getOrCreateDomain(athenzService.getDomain()).services.remove(athenzService.getName());
     }
diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/BillingController.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/BillingController.java
index 8b2f4187f65..5c6e5c9542a 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/BillingController.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/BillingController.java
@@ -118,4 +118,11 @@ public interface BillingController {
     default CollectionResult setCollectionMethod(TenantName tenant, CollectionMethod method) {
         return CollectionResult.error("Method not implemented");
     }
+
+    /** Test if the number of tenants with the given plan is under the given limit */
+    default boolean tenantsWithPlanUnderLimit(Plan plan, int limit) {
+        return true;
+    }
+
+    default void updateCache(List<TenantName> tenants) {}
 }
\ No newline at end of file
diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/BillingDatabaseClient.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/BillingDatabaseClient.java
index a89989393d0..13229b650da 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/BillingDatabaseClient.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/BillingDatabaseClient.java
@@ -107,6 +107,13 @@ public interface BillingDatabaseClient {
     Map<TenantName, Optional<Plan>> getPlans(List<TenantName> tenants);
 
     /**
+     * Returns a map with the count of plan usage.  Plans that are not in use will not appear in this result.
+     */
+    default Map<Plan, Long> getPlanCount(List<TenantName> tenants, Plan defaultPlan) {
+        return Map.of();
+    }
+
+    /**
      * Set the current plan for the given tenant
      *
      * @param tenantName The tenant to set the plan for
diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/MockBillingController.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/MockBillingController.java
index aa06e282e1c..671739bacab 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/MockBillingController.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/MockBillingController.java
@@ -14,6 +14,7 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.stream.Stream;
 
 /**
  * @author olaa
@@ -21,6 +22,9 @@ import java.util.Optional;
 public class MockBillingController implements BillingController {
 
     private final Clock clock;
+
+    PlanId defaultPlan = PlanId.from("trial");
+    List<TenantName> tenants = new ArrayList<>();
     Map<TenantName, PlanId> plans = new HashMap<>();
     Map<TenantName, PaymentInstrument> activeInstruments = new HashMap<>();
     Map<TenantName, List<Bill>> committedBills = new HashMap<>();
@@ -177,6 +181,23 @@ public class MockBillingController implements BillingController {
         return CollectionResult.success();
     }
 
+    @Override
+    public boolean tenantsWithPlanUnderLimit(Plan plan, int limit) {
+        if (limit < 0) return true;
+
+        var count = Stream.concat(tenants.stream(), plans.keySet().stream())
+                .distinct()
+                .map(tenant -> plans.getOrDefault(tenant, defaultPlan))
+                .filter(p -> p.equals(plan.id()))
+                .count();
+
+        return count < limit;
+    }
+
+    public void setTenants(List<TenantName> tenants) {
+        this.tenants = tenants;
+    }
+
     private PaymentInstrument createInstrument(String id) {
         return new PaymentInstrument(id,
                 "name",
diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/PlanResult.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/PlanResult.java
index d84d6313fa4..b86d1199956 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/PlanResult.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/billing/PlanResult.java
@@ -1,6 +1,7 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.controller.api.integration.billing;
 
+import java.util.Objects;
 import java.util.Optional;
 
 /**
@@ -32,4 +33,23 @@ public class PlanResult {
         return errorMessage;
     }
 
+    @Override
+    public String toString() {
+        return "PlanResult{" +
+                "errorMessage=" + errorMessage +
+                '}';
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        PlanResult that = (PlanResult) o;
+        return Objects.equals(errorMessage, that.errorMessage);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(errorMessage);
+    }
 }
diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/configserver/ConfigServer.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/configserver/ConfigServer.java
index 93ac16c606d..19bfc84db7a 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/configserver/ConfigServer.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/configserver/ConfigServer.java
@@ -13,7 +13,7 @@ import com.yahoo.vespa.flags.json.FlagData;
 import com.yahoo.vespa.hosted.controller.api.application.v4.model.ClusterMetrics;
 import com.yahoo.vespa.hosted.controller.api.application.v4.model.DeploymentData;
 import com.yahoo.vespa.hosted.controller.api.application.v4.model.EndpointStatus;
-import com.yahoo.vespa.hosted.controller.api.application.v4.model.ProtonMetrics;
+import com.yahoo.vespa.hosted.controller.api.application.v4.model.SearchNodeMetrics;
 import com.yahoo.vespa.hosted.controller.api.identifiers.DeploymentId;
 import com.yahoo.vespa.hosted.controller.api.integration.LogEntry;
 import com.yahoo.vespa.hosted.controller.api.integration.deployment.TestReport;
@@ -71,7 +71,7 @@ public interface ConfigServer {
 
     /**
      * Gets the contents of a file inside the current application package for a given deployment. If the path is to
-     * a directly, a JSON list with URLs to contents is returned.
+     * a directory, a JSON list with URLs to contents is returned.
      *
      * @param deployment deployment to get application package content for
      * @param path path within package to get
@@ -81,7 +81,7 @@ public interface ConfigServer {
 
     List<ClusterMetrics> getDeploymentMetrics(DeploymentId deployment);
 
-    List<ProtonMetrics> getProtonMetrics(DeploymentId deployment);
+    List<SearchNodeMetrics> getSearchNodeMetrics(DeploymentId deployment);
 
     List<String> getContentClusters(DeploymentId deployment);
 
diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/ResourceDatabaseClientMock.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/ResourceDatabaseClientMock.java
index c0d006dcb53..570ca87e538 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/ResourceDatabaseClientMock.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/ResourceDatabaseClientMock.java
@@ -79,7 +79,7 @@ public class ResourceDatabaseClientMock implements ResourceDatabaseClient {
                             a.getZoneId(),
                             plan,
                             a.resources().architecture(),
-                            a.getVersion().getMajor(),
+                            a.getMajorVersion(),
                             BigDecimal.valueOf(a.resources().vcpu()).multiply(d),
                             BigDecimal.valueOf(a.resources().memoryGb()).multiply(d),
                             BigDecimal.valueOf(a.resources().diskGb()).multiply(d),
@@ -114,7 +114,7 @@ public class ResourceDatabaseClientMock implements ResourceDatabaseClient {
                 .filter(snapshot -> snapshot.getTimestamp().isBefore(Instant.ofEpochMilli(end)))
                 .filter(snapshot -> snapshot.getApplicationId().tenant().equals(tenantName))
                 .collect(Collectors.groupingBy(
-                        usage -> Objects.hash(usage.getApplicationId(), usage.getZoneId(), tenantPlan.id().value(), usage.resources().architecture(), usage.getVersion().getMajor())
+                        usage -> Objects.hash(usage.getApplicationId(), usage.getZoneId(), tenantPlan.id().value(), usage.resources().architecture(), usage.getMajorVersion())
                 ))
                 .values().stream()
                 .map(snapshots -> resourceUsageFromSnapshots(tenantPlan, snapshots))
diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/ResourceSnapshot.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/ResourceSnapshot.java
index 4a849b831b7..b3a91767465 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/ResourceSnapshot.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/resource/ResourceSnapshot.java
@@ -20,22 +20,29 @@ import java.util.stream.Collectors;
  */
 public class ResourceSnapshot {
 
+    private static final NodeResources zero = new NodeResources(
+            0, 0, 0, 0,
+            NodeResources.DiskSpeed.any,
+            NodeResources.StorageType.any,
+            NodeResources.Architecture.any,
+            NodeResources.GpuResources.zero());
+
     private final ApplicationId applicationId;
     private final NodeResources resources;
     private final Instant timestamp;
     private final ZoneId zoneId;
-    private final Version version;
+    private final int majorVersion;
 
-    public ResourceSnapshot(ApplicationId applicationId, NodeResources resources, Instant timestamp, ZoneId zoneId, Version version) {
+    public ResourceSnapshot(ApplicationId applicationId, NodeResources resources, Instant timestamp, ZoneId zoneId, int majorVersion) {
         this.applicationId = applicationId;
         this.resources = resources;
         this.timestamp = timestamp;
         this.zoneId = zoneId;
-        this.version = version;
+        this.majorVersion = majorVersion;
     }
 
     public static ResourceSnapshot from(ApplicationId applicationId, int nodes, NodeResources resources, Instant timestamp, ZoneId zoneId) {
-        return new ResourceSnapshot(applicationId, resources.multipliedBy(nodes), timestamp, zoneId, Version.emptyVersion);
+        return new ResourceSnapshot(applicationId, resources.multipliedBy(nodes), timestamp, zoneId, 0);
     }
 
     public static ResourceSnapshot from(List<Node> nodes, Instant timestamp, ZoneId zoneId) {
@@ -44,8 +51,8 @@ public class ResourceSnapshot {
                                                  .map(node -> node.owner().get())
                                                  .collect(Collectors.toSet());
 
-        Set<Version> versions = nodes.stream()
-                .map(Node::currentVersion)
+        Set<Integer> versions = nodes.stream()
+                .map(n -> n.wantedVersion().getMajor())
                 .collect(Collectors.toSet());
 
         if (applicationIds.size() != 1) throw new IllegalArgumentException("List of nodes can only represent one application");
@@ -53,7 +60,7 @@ public class ResourceSnapshot {
 
         var resources = nodes.stream()
                 .map(Node::resources)
-                .reduce(NodeResources.zero(), NodeResources::add);
+                .reduce(zero, ResourceSnapshot::addResources);
 
         return new ResourceSnapshot(applicationIds.iterator().next(), resources, timestamp, zoneId, versions.iterator().next());
     }
@@ -74,8 +81,8 @@ public class ResourceSnapshot {
         return zoneId;
     }
 
-    public Version getVersion() {
-        return version;
+    public int getMajorVersion() {
+        return majorVersion;
     }
 
     @Override
@@ -88,11 +95,31 @@ public class ResourceSnapshot {
                 this.resources.equals(other.resources) &&
                 this.timestamp.equals(other.timestamp) &&
                 this.zoneId.equals(other.zoneId) &&
-                this.version.equals(other.version);
+                this.majorVersion == other.majorVersion;
     }
 
     @Override
     public int hashCode(){
-        return Objects.hash(applicationId, resources, timestamp, zoneId, version);
+        return Objects.hash(applicationId, resources, timestamp, zoneId, majorVersion);
+    }
+
+    /* This function does pretty much the same thing as NodeResources::add, but it allows adding resources
+     * where some dimensions that are not relevant for billing (yet) are not the same.
+     *
+     * TODO: Make this code respect all dimensions.
+     */
+    private static NodeResources addResources(NodeResources a, NodeResources b) {
+        if (a.architecture() != b.architecture() && a.architecture() != NodeResources.Architecture.any && b.architecture() != NodeResources.Architecture.any) {
+            throw new IllegalArgumentException(a + " and " + b + " are not interchangeable for resource snapshots");
+        }
+        return new NodeResources(
+                a.vcpu() + b.vcpu(),
+                a.memoryGb() + b.memoryGb(),
+                a.diskGb() + b.diskGb(),
+                0,
+                NodeResources.DiskSpeed.any,
+                NodeResources.StorageType.any,
+                a.architecture() == NodeResources.Architecture.any ? b.architecture() : a.architecture(),
+                a.gpuResources().plus(b.gpuResources()));
     }
 }
diff --git a/controller-api/src/test/java/com/yahoo/vespa/hosted/controller/api/resource/ResourceSnapshotTest.java b/controller-api/src/test/java/com/yahoo/vespa/hosted/controller/api/resource/ResourceSnapshotTest.java
new file mode 100644
index 00000000000..e5d615f8803
--- /dev/null
+++ b/controller-api/src/test/java/com/yahoo/vespa/hosted/controller/api/resource/ResourceSnapshotTest.java
@@ -0,0 +1,50 @@
+package com.yahoo.vespa.hosted.controller.api.resource;
+
+import com.yahoo.config.provision.ApplicationId;
+import com.yahoo.config.provision.NodeResources;
+import com.yahoo.config.provision.zone.ZoneId;
+import com.yahoo.vespa.hosted.controller.api.integration.configserver.Node;
+import com.yahoo.vespa.hosted.controller.api.integration.resource.ResourceSnapshot;
+import org.junit.jupiter.api.Test;
+
+import java.time.Instant;
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.fail;
+
+public class ResourceSnapshotTest {
+    @Test
+    void test_adding_resources_collapse_dimensions() {
+        var nodes = List.of(
+                nodeWithResources(NodeResources.zero().with(NodeResources.DiskSpeed.fast)),
+                nodeWithResources(NodeResources.zero().with(NodeResources.DiskSpeed.slow)));
+
+        // This should be OK and not throw exception
+        var snapshot = ResourceSnapshot.from(nodes, Instant.EPOCH, ZoneId.defaultId());
+
+        assertEquals(NodeResources.DiskSpeed.any, snapshot.resources().diskSpeed());
+    }
+
+    @Test
+    void test_adding_resources_fail() {
+        var nodes = List.of(
+                nodeWithResources(NodeResources.zero().with(NodeResources.Architecture.x86_64)),
+                nodeWithResources(NodeResources.zero().with(NodeResources.Architecture.arm64)));
+
+        try {
+            ResourceSnapshot.from(nodes, Instant.EPOCH, ZoneId.defaultId());
+            fail("Should throw an exception");
+        } catch (IllegalArgumentException e) {
+            // this should happen
+        }
+    }
+
+    private Node nodeWithResources(NodeResources resources) {
+        return Node.builder()
+                .hostname("a")
+                .owner(ApplicationId.defaultId())
+                .resources(resources)
+                .build();
+    }
+}