Maintain cache of plans in use

10 files changed, 92 insertions, 4 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);
+    }
 }
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/CloudDatabaseMaintainer.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/CloudDatabaseMaintainer.java
new file mode 100644
index 00000000000..914707aa318
--- /dev/null
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/CloudDatabaseMaintainer.java
@@ -0,0 +1,26 @@
+package com.yahoo.vespa.hosted.controller.maintenance;
+
+import com.yahoo.vespa.hosted.controller.Controller;
+import com.yahoo.vespa.hosted.controller.tenant.Tenant;
+import com.yahoo.yolean.Exceptions;
+
+import java.time.Duration;
+
+public class CloudDatabaseMaintainer extends ControllerMaintainer {
+
+    public CloudDatabaseMaintainer(Controller controller, Duration interval) {
+        super(controller, interval);
+    }
+
+    @Override
+    protected double maintain() {
+        try {
+            var tenants = controller().tenants().asList().stream().map(Tenant::name).toList();
+            controller().serviceRegistry().billingController().updateCache(tenants);
+        } catch (Exception e) {
+            log.warning("Could not update cloud database cache: " + Exceptions.toMessageString(e));
+            return 0.0;
+        }
+        return 1.0;
+    }
+}
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/ControllerMaintenance.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/ControllerMaintenance.java
index bcc8296da2a..05159b38ec6 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/ControllerMaintenance.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/ControllerMaintenance.java
@@ -75,6 +75,7 @@ public class ControllerMaintenance extends AbstractComponent {
         maintainers.add(new TenantRoleCleanupMaintainer(controller, intervals.tenantRoleMaintainer));
         maintainers.add(new ChangeRequestMaintainer(controller, intervals.changeRequestMaintainer));
         maintainers.add(new VcmrMaintainer(controller, intervals.vcmrMaintainer, metric));
+        maintainers.add(new CloudDatabaseMaintainer(controller, intervals.defaultInterval));
         maintainers.add(new CloudTrialExpirer(controller, intervals.defaultInterval));
         maintainers.add(new RetriggerMaintainer(controller, intervals.retriggerMaintainer));
         maintainers.add(new UserManagementMaintainer(controller, intervals.userManagementMaintainer, controller.serviceRegistry().roleMaintainer()));
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/user/UserApiHandler.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/user/UserApiHandler.java
index a9787e28113..3811ec22555 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/user/UserApiHandler.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/user/UserApiHandler.java
@@ -368,9 +368,8 @@ public class UserApiHandler extends ThreadedHttpRequestHandler {
 
     private boolean hasTrialCapacity() {
         if (! controller.system().isPublic()) return true;
-        var existing = controller.tenants().asList().stream().map(Tenant::name).toList();
-        var trialTenants = controller.serviceRegistry().billingController().tenantsWithPlan(existing, PlanId.from("trial"));
-        return maxTrialTenants.value() < 0 || trialTenants.size() < maxTrialTenants.value();
+        var plan = controller.serviceRegistry().planRegistry().plan("trial");
+        return controller.serviceRegistry().billingController().tenantsWithPlanUnderLimit(plan.get(), maxTrialTenants.value());
     }
 
     private static Inspector bodyInspector(HttpRequest request) {
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ServiceRegistryMock.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ServiceRegistryMock.java
index 998b371bbf1..d80c55fd00a 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ServiceRegistryMock.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ServiceRegistryMock.java
@@ -85,7 +85,7 @@ public class ServiceRegistryMock extends AbstractComponent implements ServiceReg
     private final MockEnclaveAccessService mockAMIService = new MockEnclaveAccessService();
     private final MockResourceTagger mockResourceTagger = new MockResourceTagger();
     private final MockRoleService roleService = new MockRoleService();
-    private final BillingController billingController = new MockBillingController(clock);
+    private final MockBillingController billingController = new MockBillingController(clock);
     private final ArtifactRegistryMock containerRegistry = new ArtifactRegistryMock();
     private final NoopTenantSecretService tenantSecretService = new NoopTenantSecretService();
     private final NoopEndpointSecretManager secretManager = new NoopEndpointSecretManager();
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/controller/responses/maintenance.json b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/controller/responses/maintenance.json
index 113ad1af6d1..34a3520dd1d 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/controller/responses/maintenance.json
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/controller/responses/maintenance.json
@@ -28,6 +28,9 @@
       "name": "ChangeRequestMaintainer"
     },
     {
+      "name": "CloudDatabaseMaintainer"
+    },
+    {
       "name": "CloudTrialExpirer"
     },
     {
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/user/UserApiTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/user/UserApiTest.java
index 075e001655f..b3c992cdac7 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/user/UserApiTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/user/UserApiTest.java
@@ -8,6 +8,7 @@ import com.yahoo.vespa.flags.Flags;
 import com.yahoo.vespa.flags.InMemoryFlagSource;
 import com.yahoo.vespa.flags.PermanentFlags;
 import com.yahoo.vespa.hosted.controller.ControllerTester;
+import com.yahoo.vespa.hosted.controller.api.integration.billing.MockBillingController;
 import com.yahoo.vespa.hosted.controller.api.integration.billing.PlanId;
 import com.yahoo.jdisc.http.filter.security.misc.User;
 import com.yahoo.vespa.hosted.controller.api.integration.user.UserId;
@@ -19,6 +20,7 @@ import com.yahoo.vespa.hosted.controller.tenant.Tenant;
 import org.junit.jupiter.api.Test;
 
 import java.io.File;
+import java.util.List;
 import java.util.Set;
 
 import static com.yahoo.application.container.handler.Request.Method.DELETE;
@@ -262,6 +264,7 @@ public class UserApiTest extends ControllerContainerCloudTest {
                     .withBooleanFlag(PermanentFlags.ENABLE_PUBLIC_SIGNUP_FLOW.id(), true);
             Set<Role> operator = Set.of(Role.hostedOperator(), Role.hostedSupporter(), Role.hostedAccountant());
             User user = new User("dev@domail", "Joe Developer", "dev", null);
+            tester.controller().serviceRegistry().billingController().updateCache(List.of());
 
             Role developer = Role.developer(TenantName.from("scoober"));
             tester.userManagement().createRole(developer);
@@ -292,6 +295,7 @@ public class UserApiTest extends ControllerContainerCloudTest {
             User user = new User("dev@domail", "Joe Developer", "dev", null);
 
             controller.createTenant("tenant1", Tenant.Type.cloud);
+            ((MockBillingController) controller.serviceRegistry().billingController()).setTenants(controller.controller().tenants().asList().stream().map(Tenant::name).toList());
 
             tester.assertResponse(
                     request("/user/v1/user").user(user),