Refactor assignment of single rotation

2 files changed, 49 insertions, 50 deletions

diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/rotation/RotationRepository.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/rotation/RotationRepository.java
index 53ef786d3f6..e003c362b95 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/rotation/RotationRepository.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/rotation/RotationRepository.java
@@ -1,10 +1,11 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.controller.rotation;
 
+import com.yahoo.config.application.api.DeploymentInstanceSpec;
 import com.yahoo.config.application.api.DeploymentSpec;
+import com.yahoo.config.application.api.Endpoint;
 import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.ClusterSpec;
-import com.yahoo.config.provision.Environment;
 import com.yahoo.vespa.hosted.controller.ApplicationController;
 import com.yahoo.vespa.hosted.controller.Instance;
 import com.yahoo.vespa.hosted.controller.application.AssignedRotation;
@@ -65,27 +66,30 @@ public class RotationRepository {
      * If a rotation is already assigned to the application, that rotation will be returned.
      * If no rotation is assigned, return an available rotation. The caller is responsible for assigning the rotation.
      *
-     * @param deploymentSpec the deployment spec for the application
+     * @param instanceSpec the instance deployment spec
      * @param instance the instance requesting a rotation
      * @param lock lock which must be acquired by the caller
      */
-    private Rotation getOrAssignRotation(DeploymentSpec deploymentSpec, Instance instance, RotationLock lock) {
-        if ( ! instance.rotations().isEmpty()) {
-            return allRotations.get(instance.rotations().get(0).rotationId());
+    private AssignedRotation assignRotationTo(String globalServiceId, DeploymentInstanceSpec instanceSpec,
+                                              Instance instance, RotationLock lock) {
+        RotationId rotation;
+        if (instance.rotations().isEmpty()) {
+            rotation = findAvailableRotation(instance.id(), lock).id();
+        } else {
+            rotation = instance.rotations().get(0).rotationId();
         }
-
-        if (deploymentSpec.requireInstance(instance.name()).globalServiceId().isEmpty()) {
-            throw new IllegalArgumentException("global-service-id is not set in deployment spec for instance '" +
-                                               instance.name() + "'");
-        }
-        long productionZones = deploymentSpec.requireInstance(instance.name()).zones().stream()
-                                                     .filter(zone -> zone.concerns(Environment.prod))
-                                                     .count();
-        if (productionZones < 2) {
+        var productionRegions = instanceSpec.zones().stream()
+                                            .filter(zone -> zone.environment().isProduction())
+                                            .flatMap(zone -> zone.region().stream())
+                                            .collect(Collectors.toSet());
+        if (productionRegions.size() < 2) {
             throw new IllegalArgumentException("global-service-id is set but less than 2 prod zones are defined " +
                                                "in instance '" + instance.name() + "'");
         }
-        return findAvailableRotation(instance.id(), lock);
+        return new AssignedRotation(new ClusterSpec.Id(globalServiceId),
+                                    EndpointId.defaultId(),
+                                    rotation,
+                                    productionRegions);
     }
 
     /**
@@ -107,40 +111,27 @@ public class RotationRepository {
         }
 
         // Only allow one kind of configuration syntax
-        if (     deploymentSpec.requireInstance(instance.name()).globalServiceId().isPresent()
-            && ! deploymentSpec.requireInstance(instance.name()).endpoints().isEmpty()) {
+        var instanceSpec = deploymentSpec.requireInstance(instance.name());
+        if (     instanceSpec.globalServiceId().isPresent()
+            && ! instanceSpec.endpoints().isEmpty()) {
             throw new IllegalArgumentException("Cannot provision rotations with both global-service-id and 'endpoints'");
         }
 
-        // Support the older case of setting global-service-id
-        if (deploymentSpec.requireInstance(instance.name()).globalServiceId().isPresent()) {
-            var regions = deploymentSpec.requireInstance(instance.name()).zones().stream()
-                                                .filter(zone -> zone.environment().isProduction())
-                                                .flatMap(zone -> zone.region().stream())
-                                                .collect(Collectors.toSet());
-
-            var rotation = getOrAssignRotation(deploymentSpec, instance, lock);
-
-            return List.of(
-                    new AssignedRotation(
-                            new ClusterSpec.Id(deploymentSpec.requireInstance(instance.name()).globalServiceId().get()),
-                            EndpointId.defaultId(),
-                            rotation.id(),
-                            regions
-                    )
-            );
+        // Support legacy global-service-id
+        if (instanceSpec.globalServiceId().isPresent()) {
+            return List.of(assignRotationTo(instanceSpec.globalServiceId().get(), instanceSpec, instance, lock));
         }
 
-        return assignRotations(deploymentSpec, instance, lock);
+        return assignRotationsTo(instanceSpec.endpoints(), instance, lock);
     }
 
-    private List<AssignedRotation> assignRotations(DeploymentSpec deploymentSpec, Instance instance, RotationLock lock) {
+    private List<AssignedRotation> assignRotationsTo(List<Endpoint> endpoints, Instance instance, RotationLock lock) {
         var availableRotations = new ArrayList<>(availableRotations(lock).values());
         var assignedRotationsByEndpointId = instance.rotations().stream()
                                                     .collect(Collectors.toMap(AssignedRotation::endpointId,
                                                                               Function.identity()));
         var assignments = new ArrayList<AssignedRotation>();
-        for (var endpoint : deploymentSpec.requireInstance(instance.name()).endpoints()) {
+        for (var endpoint : endpoints) {
             var endpointId = EndpointId.of(endpoint.endpointId());
             var assignedRotation = assignedRotationsByEndpointId.get(endpointId);
             RotationId rotationId;
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/rotation/RotationRepositoryTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/rotation/RotationRepositoryTest.java
index 468ecc424a9..136ed508a33 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/rotation/RotationRepositoryTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/rotation/RotationRepositoryTest.java
@@ -1,6 +1,7 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.controller.rotation;
 
+import com.yahoo.config.provision.RegionName;
 import com.yahoo.config.provision.SystemName;
 import com.yahoo.config.provision.zone.RoutingMethod;
 import com.yahoo.vespa.hosted.controller.ControllerTester;
@@ -11,13 +12,13 @@ import com.yahoo.vespa.hosted.controller.deployment.DeploymentContext;
 import com.yahoo.vespa.hosted.controller.deployment.DeploymentTester;
 import com.yahoo.vespa.hosted.controller.integration.ZoneApiMock;
 import com.yahoo.vespa.hosted.rotation.config.RotationsConfig;
-import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 
 import java.net.URI;
 import java.util.List;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 import static org.junit.Assert.assertEquals;
@@ -50,16 +51,9 @@ public class RotationRepositoryTest {
             .region("us-west-1")
             .build();
 
-    private DeploymentTester tester;
-    private RotationRepository repository;
-    private DeploymentContext application;
-
-    @Before
-    public void before() {
-        tester = new DeploymentTester(new ControllerTester(rotationsConfig));
-        repository = tester.controller().routing().rotations();
-        application = tester.newDeploymentContext("tenant1", "app1", "default");
-    }
+    private final DeploymentTester tester = new DeploymentTester(new ControllerTester(rotationsConfig));
+    private final RotationRepository repository = tester.controller().routing().rotations();
+    private final DeploymentContext application = tester.newDeploymentContext("tenant1", "app1", "default");
 
     @Test
     public void assigns_and_reuses_rotation() {
@@ -75,16 +69,30 @@ public class RotationRepositoryTest {
                                                                                application.instance(),
                                                                               lock);
             assertSingleRotation(expected, rotations, repository);
+            assertEquals(Set.of(RegionName.from("us-west-1"), RegionName.from("us-east-3")),
+                         application.instance().rotations().get(0).regions());
         }
 
         // Submitting once more assigns same rotation
-        application.submit(applicationPackage);
+        application.submit(applicationPackage).deploy();
         assertEquals(List.of(expected.id()), rotationIds(application.instance().rotations()));
+
+        // Adding region updates rotation
+        var applicationPackage = new ApplicationPackageBuilder()
+                .globalServiceId("foo")
+                .region("us-east-3")
+                .region("us-west-1")
+                .region("us-central-1")
+                .build();
+        application.submit(applicationPackage).deploy();
+        assertEquals(Set.of(RegionName.from("us-west-1"), RegionName.from("us-east-3"),
+                     RegionName.from("us-central-1")),
+                     application.instance().rotations().get(0).regions());
     }
     
     @Test
     public void strips_whitespace_in_rotation_fqdn() {
-        tester = new DeploymentTester(new ControllerTester(rotationsConfigWhitespaces));
+        var tester = new DeploymentTester(new ControllerTester(rotationsConfigWhitespaces));
         RotationRepository repository = tester.controller().routing().rotations();
         var application2 = tester.newDeploymentContext("tenant1", "app2", "default");