Add token endpoint

5 files changed, 146 insertions, 8 deletions

diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/RoutingController.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/RoutingController.java
index c50cc6051e1..81362018939 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/RoutingController.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/RoutingController.java
@@ -13,6 +13,9 @@ import com.yahoo.config.provision.InstanceName;
 import com.yahoo.config.provision.SystemName;
 import com.yahoo.config.provision.zone.RoutingMethod;
 import com.yahoo.config.provision.zone.ZoneId;
+import com.yahoo.vespa.flags.BooleanFlag;
+import com.yahoo.vespa.flags.FetchVector;
+import com.yahoo.vespa.flags.Flags;
 import com.yahoo.vespa.hosted.controller.api.identifiers.DeploymentId;
 import com.yahoo.vespa.hosted.controller.api.integration.configserver.ContainerEndpoint;
 import com.yahoo.vespa.hosted.controller.api.integration.dns.Record;
@@ -71,6 +74,7 @@ public class RoutingController {
     private final Controller controller;
     private final RoutingPolicies routingPolicies;
     private final RotationRepository rotationRepository;
+    private final BooleanFlag createTokenEndpoint;
 
     public RoutingController(Controller controller, RotationsConfig rotationsConfig) {
         this.controller = Objects.requireNonNull(controller, "controller must be non-null");
@@ -78,6 +82,7 @@ public class RoutingController {
         this.rotationRepository = new RotationRepository(Objects.requireNonNull(rotationsConfig, "rotationsConfig must be non-null"),
                                                          controller.applications(),
                                                          controller.curator());
+        this.createTokenEndpoint = Flags.ENABLE_DATAPLANE_PROXY.bindTo(controller.flagSource());
     }
 
     /** Create a routing context for given deployment */
@@ -109,11 +114,12 @@ public class RoutingController {
 
     /** Read and return zone-scoped endpoints for given deployment */
     public EndpointList readEndpointsOf(DeploymentId deployment) {
+        boolean addTokenEndpoint = createTokenEndpoint.with(FetchVector.Dimension.APPLICATION_ID, deployment.applicationId().serializedForm()).value();
         Set<Endpoint> endpoints = new LinkedHashSet<>();
         // To discover the cluster name for a zone-scoped endpoint, we need to read routing policies
         for (var policy : routingPolicies.read(deployment)) {
             RoutingMethod routingMethod = controller.zoneRegistry().routingMethod(policy.id().zone());
-            endpoints.addAll(policy.zoneEndpointsIn(controller.system(), routingMethod));
+            endpoints.addAll(policy.zoneEndpointsIn(controller.system(), routingMethod, addTokenEndpoint));
             endpoints.add(policy.regionEndpointIn(controller.system(), routingMethod));
         }
         return EndpointList.copyOf(endpoints);
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/Endpoint.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/Endpoint.java
index 80b52e0c7a4..7b2e8d0f4ed 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/Endpoint.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/Endpoint.java
@@ -45,10 +45,11 @@ public class Endpoint {
     private final Scope scope;
     private final boolean legacy;
     private final RoutingMethod routingMethod;
+    private boolean tokenEndpoint;
 
     private Endpoint(TenantAndApplicationId application, Optional<InstanceName> instanceName, EndpointId id,
                      ClusterSpec.Id cluster, URI url, URI legacyRegionalUrl, List<Target> targets, Scope scope, Port port, boolean legacy,
-                     RoutingMethod routingMethod, boolean certificateName) {
+                     RoutingMethod routingMethod, boolean certificateName, boolean tokenEndpoint) {
         Objects.requireNonNull(application, "application must be non-null");
         Objects.requireNonNull(instanceName, "instanceName must be non-null");
         Objects.requireNonNull(cluster, "cluster must be non-null");
@@ -66,6 +67,7 @@ public class Endpoint {
         this.scope = requireScope(scope, routingMethod);
         this.legacy = legacy;
         this.routingMethod = routingMethod;
+        this.tokenEndpoint = tokenEndpoint;
     }
 
     /**
@@ -353,6 +355,10 @@ public class Endpoint {
         return targets;
     }
 
+    public boolean isTokenEndpoint() {
+        return tokenEndpoint;
+    }
+
     /** An endpoint's scope */
     public enum Scope {
 
@@ -477,6 +483,7 @@ public class Endpoint {
         private RoutingMethod routingMethod = RoutingMethod.sharedLayer4;
         private boolean legacy = false;
         private boolean certificateName = false;
+        private boolean tokenEndpoint = false;
 
         private EndpointBuilder(TenantAndApplicationId application, Optional<InstanceName> instance) {
             this.application = Objects.requireNonNull(application);
@@ -544,6 +551,11 @@ public class Endpoint {
             return this;
         }
 
+        public EndpointBuilder tokenEndpoint() {
+            this.tokenEndpoint = true;
+            return this;
+        }
+
         /** Sets the port of this */
         public EndpointBuilder on(Port port) {
             this.port = port;
@@ -576,7 +588,8 @@ public class Endpoint {
             if (routingMethod.isDirect() && !port.isDefault()) {
                 throw new IllegalArgumentException("Routing method " + routingMethod + " can only use default port");
             }
-            URI url = createUrl(endpointOrClusterAsString(endpointId, cluster),
+            String prefix = tokenEndpoint ? "token-" : "";
+            URI url = createUrl(prefix + endpointOrClusterAsString(endpointId, cluster),
                                 Objects.requireNonNull(application, "application must be non-null"),
                                 Objects.requireNonNull(instance, "instance must be non-null"),
                                 Objects.requireNonNull(targets, "targets must be non-null"),
@@ -604,7 +617,8 @@ public class Endpoint {
                                 port,
                                 legacy,
                                 routingMethod,
-                                certificateName);
+                                certificateName,
+                                tokenEndpoint);
         }
 
         private Scope requireUnset(Scope scope) {
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/routing/RoutingPolicies.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/routing/RoutingPolicies.java
index e0259b21384..cd632917842 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/routing/RoutingPolicies.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/routing/RoutingPolicies.java
@@ -7,6 +7,9 @@ import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.zone.RoutingMethod;
 import com.yahoo.config.provision.zone.ZoneId;
 import com.yahoo.transaction.Mutex;
+import com.yahoo.vespa.flags.BooleanFlag;
+import com.yahoo.vespa.flags.FetchVector;
+import com.yahoo.vespa.flags.Flags;
 import com.yahoo.vespa.hosted.controller.Application;
 import com.yahoo.vespa.hosted.controller.Controller;
 import com.yahoo.vespa.hosted.controller.api.identifiers.ClusterId;
@@ -59,10 +62,12 @@ public class RoutingPolicies {
 
     private final Controller controller;
     private final CuratorDb db;
+    private final BooleanFlag createTokenEndpoint;
 
     public RoutingPolicies(Controller controller) {
         this.controller = Objects.requireNonNull(controller, "controller must be non-null");
         this.db = controller.curator();
+        this.createTokenEndpoint = Flags.ENABLE_DATAPLANE_PROXY.bindTo(controller.flagSource());
         try (var lock = db.lockRoutingPolicies()) { // Update serialized format
             for (var policy : db.readRoutingPolicies().entrySet()) {
                 db.writeRoutingPolicies(policy.getKey(), policy.getValue());
@@ -394,7 +399,8 @@ public class RoutingPolicies {
 
     /** Update zone DNS record for given policy */
     private void updateZoneDnsOf(RoutingPolicy policy, LoadBalancer loadBalancer, DeploymentId deploymentId) {
-        for (var endpoint : policy.zoneEndpointsIn(controller.system(), RoutingMethod.exclusive)) {
+        boolean addTokenEndpoint = createTokenEndpoint.with(FetchVector.Dimension.APPLICATION_ID, deploymentId.applicationId().serializedForm()).value();
+        for (var endpoint : policy.zoneEndpointsIn(controller.system(), RoutingMethod.exclusive, addTokenEndpoint)) {
             var name = RecordName.from(endpoint.dnsName());
             var record = policy.canonicalName().isPresent() ?
                     new Record(Record.Type.CNAME, name, RecordData.fqdn(policy.canonicalName().get().value())) :
@@ -406,6 +412,7 @@ public class RoutingPolicies {
 
     private void setPrivateDns(Endpoint endpoint, LoadBalancer loadBalancer, DeploymentId deploymentId) {
         if (loadBalancer.service().isEmpty()) return;
+        if (endpoint.isTokenEndpoint()) return;
         controller.serviceRegistry().vpcEndpointService()
                   .setPrivateDns(DomainName.of(endpoint.dnsName()),
                                  new ClusterId(deploymentId, endpoint.cluster()),
@@ -462,12 +469,13 @@ public class RoutingPolicies {
      * @return the updated policies
      */
     private RoutingPolicyList removePoliciesUnreferencedBy(LoadBalancerAllocation allocation, RoutingPolicyList instancePolicies, @SuppressWarnings("unused") Mutex lock) {
+        boolean addTokenEndpoint = createTokenEndpoint.with(FetchVector.Dimension.APPLICATION_ID, allocation.deployment.applicationId().serializedForm()).value();
         Map<RoutingPolicyId, RoutingPolicy> newPolicies = new LinkedHashMap<>(instancePolicies.asMap());
         Set<RoutingPolicyId> activeIds = allocation.asPolicyIds();
         RoutingPolicyList removable = instancePolicies.deployment(allocation.deployment)
                                                       .not().matching(policy -> activeIds.contains(policy.id()));
         for (var policy : removable) {
-            for (var endpoint : policy.zoneEndpointsIn(controller.system(), RoutingMethod.exclusive)) {
+            for (var endpoint : policy.zoneEndpointsIn(controller.system(), RoutingMethod.exclusive, addTokenEndpoint)) {
                 nameServiceForwarderIn(allocation.deployment.zoneId()).removeRecords(Record.Type.CNAME,
                                                                                      RecordName.from(endpoint.dnsName()),
                                                                                      Priority.normal,
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/routing/RoutingPolicy.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/routing/RoutingPolicy.java
index b4d83b7ded6..fb8f5e8e129 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/routing/RoutingPolicy.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/routing/RoutingPolicy.java
@@ -104,9 +104,15 @@ public record RoutingPolicy(RoutingPolicyId id,
     }
 
     /** Returns the zone endpoints of this */
-    public List<Endpoint> zoneEndpointsIn(SystemName system, RoutingMethod routingMethod) {
+    public List<Endpoint> zoneEndpointsIn(SystemName system, RoutingMethod routingMethod, boolean includeTokenEndpoint) {
         DeploymentId deployment = new DeploymentId(id.owner(), id.zone());
-        return List.of(endpoint(routingMethod).target(id.cluster(), deployment).in(system));
+        Endpoint zoneEndpoint = endpoint(routingMethod).target(id.cluster(), deployment).in(system);
+        if (includeTokenEndpoint) {
+            Endpoint tokenEndpoint = endpoint(routingMethod).target(id.cluster(), deployment).tokenEndpoint().in(system);
+            return List.of(zoneEndpoint, tokenEndpoint);
+        } else {
+            return List.of(zoneEndpoint);
+        }
     }
 
     /** Returns the region endpoint of this */
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/routing/RoutingPoliciesTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/routing/RoutingPoliciesTest.java
index 6ef4da8b60b..6294f9879b4 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/routing/RoutingPoliciesTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/routing/RoutingPoliciesTest.java
@@ -16,6 +16,7 @@ import com.yahoo.config.provision.RegionName;
 import com.yahoo.config.provision.SystemName;
 import com.yahoo.config.provision.zone.RoutingMethod;
 import com.yahoo.config.provision.zone.ZoneId;
+import com.yahoo.vespa.flags.Flags;
 import com.yahoo.vespa.hosted.controller.ControllerTester;
 import com.yahoo.vespa.hosted.controller.Instance;
 import com.yahoo.vespa.hosted.controller.api.identifiers.DeploymentId;
@@ -52,7 +53,9 @@ import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.stream.Collector;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertSame;
@@ -336,6 +339,103 @@ public class RoutingPoliciesTest {
     }
 
     @Test
+    void zone_token_endpoints() {
+        var tester = new RoutingPoliciesTester();
+        tester.enableTokenEndpoint(true);
+
+        var context1 = tester.newDeploymentContext("tenant1", "app1", "default");
+        var context2 = tester.newDeploymentContext("tenant1", "app2", "default");
+
+        // Deploy application
+        int clustersPerZone = 2;
+        tester.provisionLoadBalancers(clustersPerZone, context1.instanceId(), false, zone1, zone2);
+        context1.submit(applicationPackage).deferLoadBalancerProvisioningIn(Environment.prod).deploy();
+
+        // Deployment creates records and policies for all clusters in all zones
+        Set<String> expectedRecords = includeTokenEndpoints(Set.of(
+                "c0.app1.tenant1.us-west-1.vespa.oath.cloud",
+                "c1.app1.tenant1.us-west-1.vespa.oath.cloud",
+                "c0.app1.tenant1.us-central-1.vespa.oath.cloud",
+                "c1.app1.tenant1.us-central-1.vespa.oath.cloud"
+        ));
+        assertEquals(expectedRecords, tester.recordNames());
+        assertEquals(4, tester.policiesOf(context1.instanceId()).size());
+
+        // Next deploy does nothing
+        context1.submit(applicationPackage).deferLoadBalancerProvisioningIn(Environment.prod).deploy();
+        assertEquals(expectedRecords, tester.recordNames());
+        assertEquals(4, tester.policiesOf(context1.instanceId()).size());
+
+        // Add 1 cluster in each zone and deploy
+        tester.provisionLoadBalancers(clustersPerZone + 1, context1.instanceId(), false, zone1, zone2);
+        context1.submit(applicationPackage).deferLoadBalancerProvisioningIn(Environment.prod).deploy();
+        expectedRecords = includeTokenEndpoints(Set.of(
+                "c0.app1.tenant1.us-west-1.vespa.oath.cloud",
+                "c1.app1.tenant1.us-west-1.vespa.oath.cloud",
+                "c2.app1.tenant1.us-west-1.vespa.oath.cloud",
+                "c0.app1.tenant1.us-central-1.vespa.oath.cloud",
+                "c1.app1.tenant1.us-central-1.vespa.oath.cloud",
+                "c2.app1.tenant1.us-central-1.vespa.oath.cloud"
+        ));
+        assertEquals(expectedRecords, tester.recordNames());
+        assertEquals(6, tester.policiesOf(context1.instanceId()).size());
+
+        // Deploy another application
+        tester.provisionLoadBalancers(clustersPerZone, context2.instanceId(), false, zone1, zone2);
+        context2.submit(applicationPackage).deferLoadBalancerProvisioningIn(Environment.prod).deploy();
+        expectedRecords = includeTokenEndpoints(Set.of(
+                "c0.app1.tenant1.us-west-1.vespa.oath.cloud",
+                "c1.app1.tenant1.us-west-1.vespa.oath.cloud",
+                "c2.app1.tenant1.us-west-1.vespa.oath.cloud",
+                "c0.app1.tenant1.us-central-1.vespa.oath.cloud",
+                "c1.app1.tenant1.us-central-1.vespa.oath.cloud",
+                "c2.app1.tenant1.us-central-1.vespa.oath.cloud",
+                "c0.app2.tenant1.us-central-1.vespa.oath.cloud",
+                "c1.app2.tenant1.us-central-1.vespa.oath.cloud",
+                "c0.app2.tenant1.us-west-1.vespa.oath.cloud",
+                "c1.app2.tenant1.us-west-1.vespa.oath.cloud"
+        ));
+        assertEquals(expectedRecords.stream().sorted().toList(), tester.recordNames().stream().sorted().toList());
+        assertEquals(4, tester.policiesOf(context2.instanceId()).size());
+
+        // Deploy removes cluster from app1
+        tester.provisionLoadBalancers(clustersPerZone, context1.instanceId(), false, zone1, zone2);
+        context1.submit(applicationPackage).deferLoadBalancerProvisioningIn(Environment.prod).deploy();
+        expectedRecords = includeTokenEndpoints(Set.of(
+                "c0.app1.tenant1.us-west-1.vespa.oath.cloud",
+                "c1.app1.tenant1.us-west-1.vespa.oath.cloud",
+                "c0.app1.tenant1.us-central-1.vespa.oath.cloud",
+                "c1.app1.tenant1.us-central-1.vespa.oath.cloud",
+                "c0.app2.tenant1.us-central-1.vespa.oath.cloud",
+                "c1.app2.tenant1.us-central-1.vespa.oath.cloud",
+                "c0.app2.tenant1.us-west-1.vespa.oath.cloud",
+                "c1.app2.tenant1.us-west-1.vespa.oath.cloud"
+        ));
+        assertEquals(expectedRecords, tester.recordNames());
+
+        // Remove app2 completely
+        tester.controllerTester().controller().applications().requireInstance(context2.instanceId()).deployments().keySet()
+                .forEach(zone -> tester.controllerTester().controller().applications().deactivate(context2.instanceId(), zone));
+        context2.flushDnsUpdates();
+        expectedRecords = includeTokenEndpoints(Set.of(
+                "c0.app1.tenant1.us-west-1.vespa.oath.cloud",
+                "c1.app1.tenant1.us-west-1.vespa.oath.cloud",
+                "c0.app1.tenant1.us-central-1.vespa.oath.cloud",
+                "c1.app1.tenant1.us-central-1.vespa.oath.cloud"
+        ));
+        assertEquals(expectedRecords, tester.recordNames());
+        assertTrue(tester.routingPolicies().read(context2.instanceId()).isEmpty(), "Removes stale routing policies " + context2.application());
+        assertEquals(4, tester.routingPolicies().read(context1.instanceId()).size(), "Keeps routing policies for " + context1.application());
+    }
+
+    private Set<String> includeTokenEndpoints(Set<String> records) {
+        return Stream.concat(
+                        records.stream(),
+                        records.stream().map(v -> "token-" + v))
+                .collect(Collectors.toSet());
+    }
+
+    @Test
     void zone_routing_policies_without_dns_update() {
         var tester = new RoutingPoliciesTester(new DeploymentTester(), false);
         var context = tester.newDeploymentContext("tenant1", "app1", "default");
@@ -1109,6 +1209,10 @@ public class RoutingPoliciesTest {
                          .toList();
         }
 
+        void enableTokenEndpoint(boolean enabled) {
+            tester.controllerTester().flagSource().withBooleanFlag(Flags.ENABLE_DATAPLANE_PROXY.id(), enabled);
+        }
+
         /** Assert that an application endpoint points to given targets and weights */
         private void assertTargets(TenantAndApplicationId application, EndpointId endpointId, ClusterSpec.Id cluster,
                                    int loadBalancerId, Map<DeploymentId, Integer> deploymentWeights) {