Unit test for pre-provisioning of LBs

3 files changed, 43 insertions, 3 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancerSpec.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancerSpec.java
index 49a5c9c3c5e..bde736e7a28 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancerSpec.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancerSpec.java
@@ -19,8 +19,8 @@ public record LoadBalancerSpec(ApplicationId application, ClusterSpec.Id cluster
                                ZoneEndpoint settings, CloudAccount cloudAccount) {
 
     public static final ApplicationId preProvisionOwner = ApplicationId.from("hosted-vespa", "pre-provision", "default");
-    public static LoadBalancerSpec preProvisionSpec(ClusterSpec.Id slot) {
-        return new LoadBalancerSpec(preProvisionOwner, slot, Set.of(), ZoneEndpoint.defaultEndpoint, CloudAccount.empty);
+    public static LoadBalancerSpec preProvisionSpec(ClusterSpec.Id slot, CloudAccount account) {
+        return new LoadBalancerSpec(preProvisionOwner, slot, Set.of(), ZoneEndpoint.defaultEndpoint, account);
     }
 
     public LoadBalancerSpec(ApplicationId application, ClusterSpec.Id cluster, Set<Real> reals,
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/LoadBalancerProvisioner.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/LoadBalancerProvisioner.java
index 04813b476ca..24bca326f82 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/LoadBalancerProvisioner.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/LoadBalancerProvisioner.java
@@ -292,6 +292,7 @@ public class LoadBalancerProvisioner {
             LoadBalancer chosen = candidate.orElseThrow(() -> new IllegalStateException("could not find load balancer " + slot + " in pre-provisioned pool"));
             if (chosen.state() != State.active || chosen.instance().isEmpty())
                 throw new IllegalStateException("expected active load balancer in pre-provisioned pool, but got " + chosen);
+            log.log(Level.INFO, "Using " + chosen + " from pre-provisioned pool");
             service.reallocate(chosen.instance().get(), spec);
             db.removeLoadBalancer(chosen.id()); // Using a transaction to remove this, and write the instance, would be better, but much hassle.
             return chosen.instance(); // Should be immediately written again outside of this!
@@ -326,8 +327,9 @@ public class LoadBalancerProvisioner {
         // No need for lock while we provision, since we'll write atomically only after we're done, and the job lock ensures single writer.
         while (head - tail < size) {
             ClusterSpec.Id slot = slotId(head);
+            LoadBalancerSpec spec = preProvisionSpec(slot, nodeRepository.zone().cloud().account());
             db.writeLoadBalancer(new LoadBalancer(new LoadBalancerId(preProvisionOwner, slot),
-                                                  Optional.of(service.provision(preProvisionSpec(slot), Optional.of(slot.value()))),
+                                                  Optional.of(service.provision(spec, Optional.of(slot.value()))),
                                                   State.active, // Keep the expirer away.
                                                   nodeRepository.clock().instant()),
                                  null);
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/LoadBalancerProvisionerTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/LoadBalancerProvisionerTest.java
index 3aae3122a5c..fb59b3077f8 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/LoadBalancerProvisionerTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/LoadBalancerProvisionerTest.java
@@ -32,6 +32,7 @@ import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.NodeList;
 import com.yahoo.vespa.hosted.provision.lb.LoadBalancer;
 import com.yahoo.vespa.hosted.provision.lb.LoadBalancerList;
+import com.yahoo.vespa.hosted.provision.lb.LoadBalancerSpec;
 import com.yahoo.vespa.hosted.provision.lb.Real;
 import com.yahoo.vespa.hosted.provision.maintenance.LoadBalancerExpirer;
 import com.yahoo.vespa.hosted.provision.maintenance.TestMetric;
@@ -182,6 +183,43 @@ public class LoadBalancerProvisionerTest {
     }
 
     @Test
+    public void pre_provision_load_balancers() {
+        flagSource.withIntFlag(PermanentFlags.PRE_PROVISIONED_LB_COUNT.id(), 2);
+        LoadBalancerProvisioner provisioner = new LoadBalancerProvisioner(tester.nodeRepository(), tester.loadBalancerService());
+        LoadBalancerExpirer expirer = new LoadBalancerExpirer(tester.nodeRepository(), Duration.ofDays(1), tester.loadBalancerService(), new NullMetric());
+        provisioner.refreshPool();
+        expirer.run();
+        assertEquals(2, tester.nodeRepository().loadBalancers().list().size());
+        assertEquals(2, tester.nodeRepository().loadBalancers().list(LoadBalancerSpec.preProvisionOwner).size());
+
+        // Provision a load balancer when the pool has two entries.
+        ClusterSpec.Id containerCluster = ClusterSpec.Id.from("qrs");
+        prepare(app1, clusterRequest(ClusterSpec.Type.container, containerCluster));
+        List<LoadBalancer> loadBalancers = tester.nodeRepository().loadBalancers().list(app1).asList();
+        assertEquals(1, loadBalancers.size());
+        assertEquals(1, tester.nodeRepository().loadBalancers().list(LoadBalancerSpec.preProvisionOwner).asList().size());
+        assertEquals(Optional.of("1"), loadBalancers.get(0).instance().get().idSeed());
+
+        // Shrink pool to 0 entries.
+        flagSource.withIntFlag(PermanentFlags.PRE_PROVISIONED_LB_COUNT.id(), 0);
+        provisioner.refreshPool();
+        expirer.run();
+        assertEquals(loadBalancers.stream().map(LoadBalancer::id).toList(),
+                     tester.nodeRepository().loadBalancers().list().mapToList(LoadBalancer::id));
+
+        // Increase pool to 1 entry again. Creating an LB fails; the slot and idSeed are reused on retry.
+        tester.loadBalancerService().throwOnCreate(true);
+        flagSource.withIntFlag(PermanentFlags.PRE_PROVISIONED_LB_COUNT.id(), 1);
+        assertEquals("Did not expect a new load balancer to be created",
+                     assertThrows(IllegalStateException.class, provisioner::refreshPool).getMessage());
+        tester.loadBalancerService().throwOnCreate(false);
+        provisioner.refreshPool();
+        assertEquals(List.of(Optional.of("3")),
+                     tester.nodeRepository().loadBalancers().list(LoadBalancerSpec.preProvisionOwner)
+                           .mapToList(lb -> lb.instance().get().idSeed()));
+    }
+
+    @Test
     public void provision_load_balancers_with_dynamic_node_provisioning() {
         NodeResources resources = new NodeResources(1, 4, 10, 0.3);
         tester.makeReadyHosts(2, resources);