Maintain cache of plans in use

4 files changed, 55 insertions, 0 deletions

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);
+    }
 }