Provision load balancer for controller cluster

3 files changed, 25 insertions, 4 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 edf2932ad6e..2538619367f 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
@@ -31,7 +31,8 @@ public interface LoadBalancerService {
     /** Returns whether load balancers created by this service can forward traffic to given node and cluster type */
     default boolean canForwardTo(NodeType nodeType, ClusterSpec.Type clusterType) {
         return (nodeType == NodeType.tenant && clusterType.isContainer()) ||
-               (nodeType == NodeType.config && clusterType == ClusterSpec.Type.admin);
+               (nodeType == NodeType.config && clusterType == ClusterSpec.Type.admin) ||
+               (nodeType == NodeType.controller && clusterType.isContainer());
     }
 
     /** Load balancer protocols */
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 460e1e71e65..bb25e8371a2 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
@@ -51,12 +51,14 @@ public class LoadBalancerProvisioner {
     private final CuratorDatabaseClient db;
     private final LoadBalancerService service;
     private final BooleanFlag provisionConfigServerLoadBalancer;
+    private final BooleanFlag provisionControllerLoadBalancer;
 
     public LoadBalancerProvisioner(NodeRepository nodeRepository, LoadBalancerService service, FlagSource flagSource) {
         this.nodeRepository = nodeRepository;
         this.db = nodeRepository.database();
         this.service = service;
         this.provisionConfigServerLoadBalancer = Flags.CONFIGSERVER_PROVISION_LB.bindTo(flagSource);
+        this.provisionControllerLoadBalancer = Flags.CONTROLLER_PROVISION_LB.bindTo(flagSource);
         // Read and write all load balancers to make sure they are stored in the latest version of the serialization format
         for (var id : db.readLoadBalancerIds()) {
             try (var lock = db.lock(id.application())) {
@@ -147,11 +149,12 @@ public class LoadBalancerProvisioner {
         db.writeLoadBalancers(deactivatedLoadBalancers, transaction);
     }
 
-    // TODO(mpolden): Inline when feature flag is removed
+    // TODO(mpolden): Inline when feature flags are removed
     private boolean canForwardTo(NodeType type, ClusterSpec cluster) {
         boolean canForwardTo = service.canForwardTo(type, cluster.type());
-        if (canForwardTo && type == NodeType.config) {
-            return provisionConfigServerLoadBalancer.value();
+        if (canForwardTo) {
+            if (type == NodeType.config) return provisionConfigServerLoadBalancer.value();
+            if (type == NodeType.controller) return provisionControllerLoadBalancer.value();
         }
         return canForwardTo;
     }
@@ -206,6 +209,8 @@ public class LoadBalancerProvisioner {
                                  .owner(application);
         if (nodes.stream().anyMatch(node -> node.type() == NodeType.config)) {
             nodes = nodes.nodeType(NodeType.config).type(ClusterSpec.Type.admin);
+        } else if (nodes.stream().anyMatch(node -> node.type() == NodeType.controller)) {
+            nodes = nodes.nodeType(NodeType.controller).container();
         } else {
             nodes = nodes.nodeType(NodeType.tenant).container();
         }
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 f48127f650d..b4a1a4afe9b 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
@@ -230,6 +230,21 @@ public class LoadBalancerProvisionerTest {
         assertEquals(cluster, lbs.get().get(0).id().cluster());
     }
 
+    @Test
+    public void provision_load_balancer_controller_cluster() {
+        flagSource.withBooleanFlag(Flags.CONTROLLER_PROVISION_LB.id(), true);
+        ApplicationId controllerApp = ApplicationId.from("hosted-vespa", "controller", "default");
+        Supplier<List<LoadBalancer>> lbs = () -> tester.nodeRepository().loadBalancers(controllerApp).asList();
+        var cluster = ClusterSpec.Id.from("zone-config-servers");
+        var nodes = prepare(controllerApp, Capacity.fromRequiredNodeType(NodeType.controller), false,
+                            clusterRequest(ClusterSpec.Type.container, cluster));
+        assertEquals(1, lbs.get().size());
+        assertEquals("Prepare provisions load balancer with reserved nodes", 2, lbs.get().get(0).instance().reals().size());
+        tester.activate(controllerApp, nodes);
+        assertSame(LoadBalancer.State.active, lbs.get().get(0).state());
+        assertEquals(cluster, lbs.get().get(0).id().cluster());
+    }
+
     private void dirtyNodesOf(ApplicationId application) {
         tester.nodeRepository().setDirty(tester.nodeRepository().getNodes(application), Agent.system, this.getClass().getSimpleName());
     }