Ensure that combined load balancer is not expired too early

5 files changed, 26 insertions, 24 deletions

diff --git a/node-repository/pom.xml b/node-repository/pom.xml
index c4eeaa1d2d0..c63ce62f7c1 100644
--- a/node-repository/pom.xml
+++ b/node-repository/pom.xml
@@ -125,11 +125,6 @@
             <version>${project.version}</version>
             <scope>test</scope>
         </dependency>
-        <dependency>
-            <groupId>org.junit.jupiter</groupId>
-            <artifactId>junit-jupiter-params</artifactId>
-            <scope>test</scope>
-        </dependency>
     </dependencies>
     <build>
         <plugins>
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancer.java
index f8de6e1db16..0135f89c47e 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/lb/LoadBalancer.java
@@ -1,6 +1,7 @@
 // Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.provision.lb;
 
+import com.yahoo.config.provision.ClusterSpec;
 import com.yahoo.vespa.applicationmodel.InfrastructureApplication;
 import com.yahoo.vespa.hosted.provision.maintenance.LoadBalancerExpirer;
 
@@ -79,6 +80,11 @@ public class LoadBalancer {
         return new LoadBalancer(id, idSeed, Optional.of(instance), state, changedAt);
     }
 
+    /** Returns the effective container ID of given cluster. For combined clusters this returns the ID of the container cluster */
+    public static ClusterSpec.Id containerId(ClusterSpec cluster) {
+        return cluster.combinedId().orElse(cluster.id());
+    }
+
     public enum State {
 
         /** The load balancer has been provisioned and reserved for an application */
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/LoadBalancerExpirer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/LoadBalancerExpirer.java
index 105928c503e..3eb26c83a44 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/LoadBalancerExpirer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/LoadBalancerExpirer.java
@@ -160,10 +160,14 @@ public class LoadBalancerExpirer extends NodeRepositoryMaintainer {
 
     private List<Node> allocatedNodes(LoadBalancerId loadBalancer) {
         return nodeRepository().nodes()
-                .list(Node.State.active, Node.State.inactive, Node.State.reserved)
-                .owner(loadBalancer.application())
-                .cluster(loadBalancer.cluster())
-                .asList();
+                               .list(Node.State.active, Node.State.inactive, Node.State.reserved)
+                               .owner(loadBalancer.application())
+                               // Always match the cluster by the effective container ID
+                               // TODO(mpolden): Remove this and use NodeList::cluster once combined disappears in Vespa 9
+                               .matching((node) -> node.allocation().isPresent() &&
+                                                   LoadBalancer.containerId(node.allocation().get().membership().cluster())
+                                                               .equals(loadBalancer.cluster()))
+                               .asList();
     }
 
 }
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 239b962360b..022c00693c2 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
@@ -101,7 +101,7 @@ public class LoadBalancerProvisioner {
     public void prepare(ApplicationId application, ClusterSpec cluster, NodeSpec requested) {
         if (!shouldProvision(application, requested.type(), cluster.type())) return;
         try (var lock = db.lock(application)) {
-            ClusterSpec.Id clusterId = effectiveId(cluster);
+            ClusterSpec.Id clusterId = LoadBalancer.containerId(cluster);
             LoadBalancerId loadBalancerId = requireNonClashing(new LoadBalancerId(application, clusterId));
             prepare(loadBalancerId, cluster.zoneEndpoint(), requested);
         }
@@ -121,7 +121,7 @@ public class LoadBalancerProvisioner {
         Map<ClusterSpec.Id, ZoneEndpoint> activatingClusters = clusters.stream()
                                                                        // .collect(Collectors.toMap(ClusterSpec::id, ClusterSpec::zoneEndpoint));
                                                                        // TODO: this dies with combined clusters
-                                                                       .collect(groupingBy(LoadBalancerProvisioner::effectiveId,
+                                                                       .collect(groupingBy(LoadBalancer::containerId,
                                                                                           reducing(ZoneEndpoint.defaultEndpoint,
                                                                                                    ClusterSpec::zoneEndpoint,
                                                                                                    (o, n) -> o.isDefault() ? n : o)));
@@ -391,7 +391,7 @@ public class LoadBalancerProvisioner {
         } else {
             nodes = nodes.nodeType(NodeType.tenant).container();
         }
-        return nodes.groupingBy(node -> effectiveId(node.allocation().get().membership().cluster()));
+        return nodes.groupingBy(node -> LoadBalancer.containerId(node.allocation().get().membership().cluster()));
     }
 
     /** Returns real servers for given nodes */
@@ -436,8 +436,4 @@ public class LoadBalancerProvisioner {
         return reachable;
     }
 
-    private static ClusterSpec.Id effectiveId(ClusterSpec cluster) {
-        return cluster.combinedId().orElse(cluster.id());
-    }
-
 }
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 e3b5f25f104..dddf4c48ae6 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
@@ -38,8 +38,6 @@ import com.yahoo.vespa.hosted.provision.maintenance.TestMetric;
 import com.yahoo.vespa.hosted.provision.node.Agent;
 import com.yahoo.vespa.hosted.provision.node.IP;
 import org.junit.Test;
-import org.junit.jupiter.params.ParameterizedTest;
-import org.junit.jupiter.params.provider.MethodSource;
 
 import java.time.Duration;
 import java.util.Collections;
@@ -477,14 +475,17 @@ public class LoadBalancerProvisionerTest {
         assertEquals(cloudAccount1, loadBalancers.first().get().instance().get().cloudAccount());
     }
 
-    // TODO: 'combined' does not work
-    public static Object[] data() {
-        return new Object[]{container /*, combined */};
+    @Test
+    public void container_load_balancer_with_existing_in_removable_state() {
+        load_balancer_with_existing_in_removable_state(container);
+    }
+
+    @Test
+    public void combined_load_balancer_with_existing_in_removable_state() {
+        load_balancer_with_existing_in_removable_state(combined);
     }
 
-    @MethodSource("data")
-    @ParameterizedTest(name = "clusterType={0}")
-    public void load_balancer_with_existing_in_removable_state(ClusterSpec.Type clusterType) {
+    private void load_balancer_with_existing_in_removable_state(ClusterSpec.Type clusterType) {
         ClusterResources resources = new ClusterResources(3, 1, nodeResources);
         Capacity capacity = Capacity.from(resources, resources, IntRange.empty(), false, true, Optional.of(CloudAccount.empty), ClusterInfo.empty());
         ClusterSpec clusterSpec = clusterRequest(