Support reallocating load balancers from preprovisioned pool

4 files changed, 48 insertions, 39 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancerService.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancerService.java
index 40323a4f2d3..6649b18ed02 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancerService.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancerService.java
@@ -36,6 +36,8 @@ public interface LoadBalancerService {
      */
     LoadBalancerInstance configure(LoadBalancerInstance instance, LoadBalancerSpec spec, boolean force);
 
+    void reallocate(LoadBalancerInstance provisioned, LoadBalancerSpec spec);
+
     /** Permanently remove given load balancer */
     void remove(LoadBalancer loadBalancer);
 
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancerServiceMock.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancerServiceMock.java
index 3b444312b14..554ec26d5f4 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancerServiceMock.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancerServiceMock.java
@@ -3,6 +3,7 @@ package com.yahoo.vespa.hosted.provision.lb;
 
 import ai.vespa.http.DomainName;
 import com.google.common.collect.ImmutableSet;
+import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.ClusterSpec;
 import com.yahoo.config.provision.EndpointsChecker.Availability;
 import com.yahoo.config.provision.EndpointsChecker.Endpoint;
@@ -17,18 +18,32 @@ import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.atomic.AtomicBoolean;
 
+import static java.util.Objects.requireNonNull;
+import static java.util.stream.Collectors.toMap;
+
 /**
  * @author mpolden
  */
 public class LoadBalancerServiceMock implements LoadBalancerService {
 
-    private final Map<LoadBalancerId, LoadBalancerInstance> instances = new HashMap<>();
+    private record Key(ApplicationId application, ClusterSpec.Id cluster, UUID idSeed) {
+        @Override public int hashCode() { return idSeed == null ? Objects.hash(application, cluster) : Objects.hash(idSeed); }
+        @Override public boolean equals(Object o) {
+            if (o == this) return true;
+            if ( ! (o instanceof Key key)) return false;
+            if (idSeed != null) return Objects.equals(idSeed, key.idSeed);
+            return Objects.equals(application, key.application) &&
+                   Objects.equals(cluster, key.cluster);
+        }
+    }
+    private final Map<Key, LoadBalancerInstance> instances = new HashMap<>();
     private boolean throwOnCreate = false;
     private boolean supportsProvisioning = true;
     private final AtomicBoolean uuid = new AtomicBoolean(true);
 
     public Map<LoadBalancerId, LoadBalancerInstance> instances() {
-        return Collections.unmodifiableMap(instances);
+        return instances.entrySet().stream().collect(toMap(e -> new LoadBalancerId(e.getKey().application, e.getKey().cluster),
+                                                           Map.Entry::getValue));
     }
 
     public LoadBalancerServiceMock throwOnCreate(boolean throwOnCreate) {
@@ -57,7 +72,6 @@ public class LoadBalancerServiceMock implements LoadBalancerService {
     public LoadBalancerInstance provision(LoadBalancerSpec spec) {
         if (throwOnCreate) throw new IllegalStateException("Did not expect a new load balancer to be created");
         Optional<UUID> idSeed = uuid.getAndSet(false) ? Optional.of(UUID.fromString("c11272ab-d20e-4c86-b808-ffedaa00c480")) : Optional.empty();
-        var id = new LoadBalancerId(spec.application(), spec.cluster());
         var instance = new LoadBalancerInstance(
                 idSeed,
                 Optional.of(DomainName.of("lb-" + spec.application().toShortString() + "-" + spec.cluster().value())),
@@ -70,14 +84,14 @@ public class LoadBalancerServiceMock implements LoadBalancerService {
                 spec.settings(),
                 spec.settings().isPrivateEndpoint() ? List.of(PrivateServiceId.of("service")) : List.of(),
                 spec.cloudAccount());
-        instances.put(id, instance);
+        instances.put(new Key(spec.application(), spec.cluster(), idSeed.orElse(null)), instance);
         return instance;
     }
 
     @Override
     public LoadBalancerInstance configure(LoadBalancerInstance instance, LoadBalancerSpec spec, boolean force) {
-        var id = new LoadBalancerId(spec.application(), spec.cluster());
-        var oldInstance = Objects.requireNonNull(instances.get(id), "expected existing load balancer " + id);
+        var id = new Key(spec.application(), spec.cluster(), instance.idSeed().orElse(null));
+        var oldInstance = requireNonNull(instances.get(id), "expected existing load balancer " + id);
         if (!force && !oldInstance.reals().isEmpty() && spec.reals().isEmpty()) {
             throw new IllegalArgumentException("Refusing to remove all reals from load balancer " + id);
         }
@@ -89,8 +103,17 @@ public class LoadBalancerServiceMock implements LoadBalancerService {
     }
 
     @Override
+    public void reallocate(LoadBalancerInstance provisioned, LoadBalancerSpec spec) {
+        instances.put(new Key(spec.application(), spec.cluster(), null),
+                      requireNonNull(instances.remove(new Key(null, null, provisioned.idSeed().get())))); // ᕙ༼◕_◕༽ᕤ
+    }
+
+    @Override
     public void remove(LoadBalancer loadBalancer) {
-        instances.remove(loadBalancer.id());
+        requireNonNull(instances.remove(new Key(loadBalancer.id().application(),
+                                                loadBalancer.id().cluster(),
+                                                loadBalancer.instance().get().idSeed().orElse(null))),
+                       "expected load balancer to exist: " + loadBalancer.id());
     }
 
     @Override
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 e49e5aed739..199578d66f2 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
@@ -16,13 +16,17 @@ import java.util.Set;
  *
  * @author mpolden
  */
-public class LoadBalancerSpec {
-
-    private final ApplicationId application;
-    private final ClusterSpec.Id cluster;
-    private final Set<Real> reals;
-    private final ZoneEndpoint settings;
-    private final CloudAccount cloudAccount;
+public record LoadBalancerSpec(ApplicationId application, ClusterSpec.Id cluster, Set<Real> reals,
+                               ZoneEndpoint settings, CloudAccount cloudAccount) {
+
+    public static final ApplicationId preProvisionOwner = ApplicationId.from("hosted-vespa", "pre-provision", "default");
+    public static LoadBalancerSpec preProvisionSpec(int slot) {
+        return new LoadBalancerSpec(preProvisionOwner,
+                                    ClusterSpec.Id.from("slot-" + slot),
+                                    Set.of(),
+                                    ZoneEndpoint.defaultEndpoint,
+                                    CloudAccount.empty);
+    }
 
     public LoadBalancerSpec(ApplicationId application, ClusterSpec.Id cluster, Set<Real> reals,
                             ZoneEndpoint settings, CloudAccount cloudAccount) {
@@ -33,29 +37,4 @@ public class LoadBalancerSpec {
         this.cloudAccount = Objects.requireNonNull(cloudAccount);
     }
 
-    /** Owner of the load balancer */
-    public ApplicationId application() {
-        return application;
-    }
-
-    /** The target cluster of this load balancer */
-    public ClusterSpec.Id cluster() {
-        return cluster;
-    }
-
-    /** Real servers to attach to this load balancer */
-    public Set<Real> reals() {
-        return reals;
-    }
-
-    /** Static user-configured settings for this load balancer. */
-    public ZoneEndpoint settings() {
-        return settings;
-    }
-
-    /** Cloud account to use when satisfying this */
-    public CloudAccount cloudAccount() {
-        return cloudAccount;
-    }
-
 }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/SharedLoadBalancerService.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/SharedLoadBalancerService.java
index eb99e5647df..9124b012dfe 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/SharedLoadBalancerService.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/SharedLoadBalancerService.java
@@ -40,6 +40,11 @@ public class SharedLoadBalancerService implements LoadBalancerService {
         return instance.with(spec.reals(), spec.settings(), Optional.empty());
     }
 
+    @Override
+    public void reallocate(LoadBalancerInstance provisioned, LoadBalancerSpec spec) {
+        throw new UnsupportedOperationException("reallocate is not supported with " + getClass());
+    }
+
     private LoadBalancerInstance create(LoadBalancerSpec spec) {
         if ( ! spec.settings().isPublicEndpoint())
             throw new IllegalArgumentException("non-public endpoints is not supported with " + getClass());