Revert "Remove redundant rate limit"

This reverts commit c02fa97f0949d2be1baa9318d08e43a57196fcd5.

3 files changed, 58 insertions, 3 deletions

diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/DnsMaintainer.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/DnsMaintainer.java
index c52306266a4..7e0032f03c5 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/DnsMaintainer.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/DnsMaintainer.java
@@ -2,6 +2,7 @@
 package com.yahoo.vespa.hosted.controller.maintenance;
 
 import com.yahoo.vespa.hosted.controller.Controller;
+import com.yahoo.vespa.hosted.controller.api.integration.dns.NameService;
 import com.yahoo.vespa.hosted.controller.api.integration.dns.Record;
 import com.yahoo.vespa.hosted.controller.api.integration.dns.RecordData;
 import com.yahoo.vespa.hosted.controller.dns.NameServiceQueue.Priority;
@@ -11,7 +12,12 @@ import com.yahoo.vespa.hosted.controller.rotation.RotationLock;
 import com.yahoo.vespa.hosted.controller.rotation.RotationRepository;
 
 import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
 import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicInteger;
 
 /**
  * Performs DNS maintenance tasks such as removing DNS aliases for unassigned rotations.
@@ -20,6 +26,8 @@ import java.util.Map;
  */
 public class DnsMaintainer extends Maintainer {
 
+    private final AtomicInteger rotationIndex = new AtomicInteger(0);
+
     public DnsMaintainer(Controller controller, Duration interval, JobControl jobControl) {
         super(controller, interval, jobControl);
     }
@@ -32,7 +40,7 @@ public class DnsMaintainer extends Maintainer {
     protected void maintain() {
         try (RotationLock lock = rotationRepository().lock()) {
             Map<RotationId, Rotation> unassignedRotations = rotationRepository().availableRotations(lock);
-            unassignedRotations.values().forEach(this::removeCname);
+            rotationToCheckOf(unassignedRotations.values()).ifPresent(this::removeCname);
         }
     }
 
@@ -42,4 +50,20 @@ public class DnsMaintainer extends Maintainer {
         controller().nameServiceForwarder().removeRecords(Record.Type.CNAME, RecordData.fqdn(rotation.name()), Priority.normal);
     }
 
+    /**
+     * Returns the rotation that should be checked in this run. We check only one rotation per run to avoid running into
+     * rate limits that may be imposed by the {@link NameService} implementation.
+     */
+    private Optional<Rotation> rotationToCheckOf(Collection<Rotation> rotations) {
+        if (rotations.isEmpty()) return Optional.empty();
+        List<Rotation> rotationList = new ArrayList<>(rotations);
+        int index = rotationIndex.getAndUpdate((i) -> {
+            if (i < rotationList.size() - 1) {
+                return ++i;
+            }
+            return 0;
+        });
+        return Optional.of(rotationList.get(index));
+    }
+
 }
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/ControllerTester.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/ControllerTester.java
index d01e2f2497f..2b0ee741e7e 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/ControllerTester.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/ControllerTester.java
@@ -74,7 +74,7 @@ import static org.junit.Assert.assertNotNull;
  */
 public final class ControllerTester {
 
-    private static final int availableRotations = 10;
+    public static final int availableRotations = 10;
 
     private final AthenzDbMock athenzDb;
     private final ManualClock clock;
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/DnsMaintainerTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/DnsMaintainerTest.java
index e40ea49bd8c..13218cc2442 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/DnsMaintainerTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/DnsMaintainerTest.java
@@ -6,12 +6,17 @@ import com.yahoo.config.provision.Environment;
 import com.yahoo.config.provision.RegionName;
 import com.yahoo.config.provision.zone.ZoneId;
 import com.yahoo.vespa.hosted.controller.Application;
+import com.yahoo.vespa.hosted.controller.ControllerTester;
 import com.yahoo.vespa.hosted.controller.api.integration.dns.Record;
+import com.yahoo.vespa.hosted.controller.api.integration.dns.RecordData;
 import com.yahoo.vespa.hosted.controller.api.integration.dns.RecordName;
 import com.yahoo.vespa.hosted.controller.application.ApplicationPackage;
+import com.yahoo.vespa.hosted.controller.application.Endpoint;
 import com.yahoo.vespa.hosted.controller.deployment.ApplicationPackageBuilder;
 import com.yahoo.vespa.hosted.controller.deployment.DeploymentTester;
 import com.yahoo.vespa.hosted.controller.persistence.MockCuratorDb;
+import com.yahoo.vespa.hosted.controller.rotation.Rotation;
+import com.yahoo.vespa.hosted.controller.rotation.RotationId;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -96,15 +101,41 @@ public class DnsMaintainerTest {
         tester.controllerTester().deleteApplication(application.id());
 
         // DnsMaintainer removes records
-        maintainer.maintain();
+        for (int i = 0; i < ControllerTester.availableRotations; i++) {
+            maintainer.maintain();
+        }
         tester.updateDns();
         assertFalse("DNS record removed", findCname.apply("app1--tenant1.global.vespa.yahooapis.com").isPresent());
         assertFalse("DNS record removed", findCname.apply("app1--tenant1.global.vespa.oath.cloud").isPresent());
         assertFalse("DNS record removed", findCname.apply("app1.tenant1.global.vespa.yahooapis.com").isPresent());
     }
 
+    @Test
+    public void rate_limit_record_removal() {
+        // Create stale records
+        int staleTotal = ControllerTester.availableRotations;
+        for (int i = 1; i <= staleTotal; i++) {
+            Rotation r = rotation(i);
+            tester.controllerTester().nameService().createCname(RecordName.from("stale-record-" + i + "." +
+                                                                                Endpoint.OATH_DNS_SUFFIX),
+                                                                RecordData.from(r.name() + "."));
+        }
+
+        // One record is removed per run
+        for (int i = 1; i <= staleTotal*2; i++) {
+            maintainer.run();
+            tester.updateDns();
+            assertEquals(Math.max(staleTotal - i, 0), records().size());
+        }
+    }
+
     private Set<Record> records() {
         return tester.controllerTester().nameService().records();
     }
 
+    private static Rotation rotation(int n) {
+        String id = String.format("%02d", n);
+        return new Rotation(new RotationId("rotation-id-" + id), "rotation-fqdn-" + id);
+    }
+
 }