Add SQL-based DNS refresh processing mechanism (#1971)

core/src/test/java/google/registry/dns/DnsUtilsTest.java

@@ -14,6 +14,7 @@
 
 package google.registry.dns;
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.createTld;
@@ -26,6 +27,7 @@ import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Ordering;
@@ -36,6 +38,7 @@ import google.registry.model.common.DnsRefreshRequest;
 import google.registry.persistence.transaction.JpaTestExtensions;
 import google.registry.persistence.transaction.JpaTestExtensions.JpaIntegrationTestExtension;
 import google.registry.testing.FakeClock;
+import java.util.Comparator;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
 import org.junit.jupiter.api.Assertions;
@@ -53,8 +56,7 @@ public class DnsUtilsTest {
 
   private final DnsQueue dnsQueue = mock(DnsQueue.class);
   private final DnsUtils dnsUtils = new DnsUtils(dnsQueue);
-
-  FakeClock clock = new FakeClock(DateTime.parse("2020-02-02T01:23:45Z"));
+  private final FakeClock clock = new FakeClock(DateTime.parse("2020-02-02T01:23:45Z"));
 
   @RegisterExtension
   JpaIntegrationTestExtension jpa =
@@ -138,12 +140,133 @@ public class DnsUtilsTest {
     assertRequest(request, TargetType.DOMAIN, domainName, tld, clock.nowUtc().plusMinutes(3));
   }
 
+  @Test
+  void testSuccess_ProcessRequests() {
+    ImmutableList<DnsRefreshRequest> requests = processRequests();
+    DateTime processtime = clock.nowUtc();
+    assertThat(requests.size()).isEqualTo(4);
+    assertRequest(
+        requests.get(0),
+        TargetType.DOMAIN,
+        "test2.tld",
+        "tld",
+        clock.nowUtc().minusMinutes(4),
+        processtime);
+    assertRequest(
+        requests.get(1),
+        TargetType.DOMAIN,
+        "test1.tld",
+        "tld",
+        clock.nowUtc().minusMinutes(3),
+        processtime);
+    assertRequest(
+        requests.get(2),
+        TargetType.HOST,
+        "ns1.test2.tld",
+        "tld",
+        clock.nowUtc().minusMinutes(1),
+        processtime);
+    assertRequest(
+        requests.get(3),
+        TargetType.DOMAIN,
+        "test5.tld",
+        "tld",
+        clock.nowUtc().minusMinutes(1),
+        processtime);
+    requests = loadAllOf(DnsRefreshRequest.class);
+    assertThat(requests.size()).isEqualTo(7);
+    // The four processed records should have updated process time in SQL as well.
+    assertThat(requests.stream().filter(e -> e.getLastProcessTime().equals(processtime)).count())
+        .isEqualTo(4);
+    clock.advanceOneMilli();
+
+    // Requests within cooldown period not included.
+    requests =
+        dnsUtils.readAndUpdateRequestsWithLatestProcessTime("tld", Duration.standardMinutes(1), 4);
+    assertThat(requests.size()).isEqualTo(1);
+    assertRequest(
+        requests.get(0),
+        TargetType.DOMAIN,
+        "test6.tld",
+        "tld",
+        clock.nowUtc().minusMinutes(1).minusMillis(1),
+        clock.nowUtc());
+  }
+
+  @Test
+  void testSuccess_deleteRequests() {
+    dnsUtils.deleteRequests(processRequests());
+    ImmutableList<DnsRefreshRequest> remainingRequests =
+        loadAllOf(DnsRefreshRequest.class).stream()
+            .sorted(Comparator.comparing(DnsRefreshRequest::getRequestTime))
+            .collect(toImmutableList());
+    assertThat(remainingRequests.size()).isEqualTo(3);
+    assertRequest(
+        remainingRequests.get(0),
+        TargetType.DOMAIN,
+        "something.example",
+        "example",
+        clock.nowUtc().minusMinutes(2));
+    assertRequest(
+        remainingRequests.get(1),
+        TargetType.DOMAIN,
+        "test6.tld",
+        "tld",
+        clock.nowUtc().minusMinutes(1));
+    assertRequest(
+        remainingRequests.get(2),
+        TargetType.DOMAIN,
+        "test4.tld",
+        "tld",
+        clock.nowUtc().plusMinutes(1));
+    tm().transact(() -> tm().delete(remainingRequests.get(2)));
+    assertThat(loadAllOf(DnsRefreshRequest.class).size()).isEqualTo(2);
+    // Should not throw even though one of the request is already deleted.
+    dnsUtils.deleteRequests(remainingRequests);
+    assertThat(loadAllOf(DnsRefreshRequest.class).size()).isEqualTo(0);
+  }
+
+  private ImmutableList<DnsRefreshRequest> processRequests() {
+    useDnsSql();
+    createTld("example");
+    // Domain Included.
+    dnsUtils.requestDomainDnsRefresh("test1.tld", Duration.standardMinutes(1));
+    // This one should be returned before test1.tld, even though it's added later, because of
+    // the delay specified in test1.tld.
+    dnsUtils.requestDomainDnsRefresh("test2.tld");
+    // Not included because the TLD is not under management.
+    dnsUtils.requestDomainDnsRefresh("something.example", Duration.standardMinutes(2));
+    clock.advanceBy(Duration.standardMinutes(3));
+    // Host included.
+    dnsUtils.requestHostDnsRefresh("ns1.test2.tld");
+    // Not included because the request time is in the future
+    dnsUtils.requestDomainDnsRefresh("test4.tld", Duration.standardMinutes(2));
+    // Included after the previous one. Same request time, order by insertion order (i.e. ID);
+    dnsUtils.requestDomainDnsRefresh("test5.tld");
+    // Not included because batch size is exceeded;
+    dnsUtils.requestDomainDnsRefresh("test6.tld");
+    clock.advanceBy(Duration.standardMinutes(1));
+    return dnsUtils.readAndUpdateRequestsWithLatestProcessTime(
+        "tld", Duration.standardMinutes(1), 4);
+  }
+
   private static void assertRequest(
       DnsRefreshRequest request, TargetType type, String name, String tld, DateTime requestTime) {
+    assertRequest(request, type, name, tld, requestTime, START_OF_TIME);
+  }
+
+  private static void assertRequest(
+      DnsRefreshRequest request,
+      TargetType type,
+      String name,
+      String tld,
+      DateTime requestTime,
+      DateTime processTime) {
     assertThat(request.getType()).isEqualTo(type);
     assertThat(request.getName()).isEqualTo(name);
     assertThat(request.getTld()).isEqualTo(tld);
     assertThat(request.getRequestTime()).isEqualTo(requestTime);
+    assertThat(request.getLastProcessTime()).isEqualTo(processTime);
   }
 
   private void useDnsSql() {

core/src/test/java/google/registry/dns/PublishDnsUpdatesActionTest.java

@@ -22,7 +22,7 @@ import static google.registry.dns.DnsModule.PARAM_HOSTS;
 import static google.registry.dns.DnsModule.PARAM_LOCK_INDEX;
 import static google.registry.dns.DnsModule.PARAM_NUM_PUBLISH_LOCKS;
 import static google.registry.dns.DnsModule.PARAM_PUBLISH_TASK_ENQUEUED;
-import static google.registry.dns.DnsModule.PARAM_REFRESH_REQUEST_CREATED;
+import static google.registry.dns.DnsModule.PARAM_REFRESH_REQUEST_TIME;
 import static google.registry.dns.PublishDnsUpdatesAction.RETRIES_BEFORE_PERMANENT_FAILURE;
 import static google.registry.request.RequestParameters.PARAM_TLD;
 import static google.registry.testing.DatabaseHelper.createTld;
@@ -302,7 +302,7 @@ public class PublishDnsUpdatesActionTest {
             .param(PARAM_LOCK_INDEX, "1")
             .param(PARAM_NUM_PUBLISH_LOCKS, "1")
             .param(PARAM_PUBLISH_TASK_ENQUEUED, clock.nowUtc().toString())
-            .param(PARAM_REFRESH_REQUEST_CREATED, clock.nowUtc().minusHours(2).toString())
+            .param(PARAM_REFRESH_REQUEST_TIME, clock.nowUtc().minusHours(2).toString())
             .param(PARAM_DOMAINS, "example1.xn--q9jyb4c,example2.xn--q9jyb4c")
             .param(PARAM_HOSTS, "")
             .header("content-type", "application/x-www-form-urlencoded"),
@@ -313,7 +313,7 @@ public class PublishDnsUpdatesActionTest {
             .param(PARAM_LOCK_INDEX, "1")
             .param(PARAM_NUM_PUBLISH_LOCKS, "1")
             .param(PARAM_PUBLISH_TASK_ENQUEUED, clock.nowUtc().toString())
-            .param(PARAM_REFRESH_REQUEST_CREATED, clock.nowUtc().minusHours(2).toString())
+            .param(PARAM_REFRESH_REQUEST_TIME, clock.nowUtc().minusHours(2).toString())
             .param(PARAM_DOMAINS, "example3.xn--q9jyb4c,example4.xn--q9jyb4c")
             .param(PARAM_HOSTS, "ns1.example.xn--q9jyb4c")
             .header("content-type", "application/x-www-form-urlencoded"));
@@ -341,7 +341,7 @@ public class PublishDnsUpdatesActionTest {
             .param(PARAM_LOCK_INDEX, "1")
             .param(PARAM_NUM_PUBLISH_LOCKS, "1")
             .param(PARAM_PUBLISH_TASK_ENQUEUED, clock.nowUtc().toString())
-            .param(PARAM_REFRESH_REQUEST_CREATED, clock.nowUtc().minusHours(2).toString())
+            .param(PARAM_REFRESH_REQUEST_TIME, clock.nowUtc().minusHours(2).toString())
             .param(PARAM_DOMAINS, "example1.xn--q9jyb4c,example2.xn--q9jyb4c")
             .param(PARAM_HOSTS, "")
             .header("content-type", "application/x-www-form-urlencoded"),
@@ -352,7 +352,7 @@ public class PublishDnsUpdatesActionTest {
             .param(PARAM_LOCK_INDEX, "1")
             .param(PARAM_NUM_PUBLISH_LOCKS, "1")
             .param(PARAM_PUBLISH_TASK_ENQUEUED, clock.nowUtc().toString())
-            .param(PARAM_REFRESH_REQUEST_CREATED, clock.nowUtc().minusHours(2).toString())
+            .param(PARAM_REFRESH_REQUEST_TIME, clock.nowUtc().minusHours(2).toString())
             .param(PARAM_DOMAINS, "example3.xn--q9jyb4c,example4.xn--q9jyb4c,example5.xn--q9jyb4c")
             .param(PARAM_HOSTS, "ns1.example.xn--q9jyb4c")
             .header("content-type", "application/x-www-form-urlencoded"));
@@ -378,7 +378,7 @@ public class PublishDnsUpdatesActionTest {
             .param(PARAM_LOCK_INDEX, "1")
             .param(PARAM_NUM_PUBLISH_LOCKS, "1")
             .param(PARAM_PUBLISH_TASK_ENQUEUED, clock.nowUtc().toString())
-            .param(PARAM_REFRESH_REQUEST_CREATED, clock.nowUtc().minusHours(2).toString())
+            .param(PARAM_REFRESH_REQUEST_TIME, clock.nowUtc().minusHours(2).toString())
             .param(PARAM_DOMAINS, "example1.xn--q9jyb4c")
             .param(PARAM_HOSTS, "")
             .header("content-type", "application/x-www-form-urlencoded"),
@@ -389,7 +389,7 @@ public class PublishDnsUpdatesActionTest {
             .param(PARAM_LOCK_INDEX, "1")
             .param(PARAM_NUM_PUBLISH_LOCKS, "1")
             .param(PARAM_PUBLISH_TASK_ENQUEUED, clock.nowUtc().toString())
-            .param(PARAM_REFRESH_REQUEST_CREATED, clock.nowUtc().minusHours(2).toString())
+            .param(PARAM_REFRESH_REQUEST_TIME, clock.nowUtc().minusHours(2).toString())
             .param(PARAM_DOMAINS, "")
             .param(PARAM_HOSTS, "ns1.example.xn--q9jyb4c")
             .header("content-type", "application/x-www-form-urlencoded"));

core/src/test/java/google/registry/dns/ReadDnsQueueActionTest.java

@@ -144,7 +144,7 @@ public class ReadDnsQueueActionTest {
                     .url(PublishDnsUpdatesAction.PATH)
                     .param("tld", tldToDnsWriter.getKey())
                     .param("dnsWriter", tldToDnsWriter.getValue())
-                    .param("itemsCreated", "3000-01-01T00:00:00.000Z")
+                    .param("requestTime", "3000-01-01T00:00:00.000Z")
                     .param("enqueued", "3000-01-01T01:00:00.000Z")
                     // Single-lock TLDs should use lock 1 of 1 by default
                     .param("lockIndex", "1")
@@ -245,7 +245,7 @@ public class ReadDnsQueueActionTest {
         DNS_PUBLISH_PUSH_QUEUE_NAME,
         new TaskMatcher()
             .param("enqueued", "3000-02-05T01:00:00.000Z")
-            .param("itemsCreated", "3000-02-03T00:00:00.000Z")
+            .param("requestTime", "3000-02-03T00:00:00.000Z")
             .param("tld", "com")
             .param("dnsWriter", "comWriter")
             .param("domains", "domain1.com,domain2.com,domain3.com")
@@ -484,7 +484,7 @@ public class ReadDnsQueueActionTest {
             .url(PublishDnsUpdatesAction.PATH)
             .param("tld", "multilock.uk")
             .param("dnsWriter", "multilockWriter")
-            .param("itemsCreated", "3000-01-01T00:00:00.000Z")
+            .param("requestTime", "3000-01-01T00:00:00.000Z")
             .param("enqueued", "3000-01-01T01:00:00.000Z")
             .param("domains", "hello.multilock.uk")
             .param("hosts", "ns1.abc.hello.multilock.uk,ns2.hello.multilock.uk")
@@ -493,7 +493,7 @@ public class ReadDnsQueueActionTest {
             .url(PublishDnsUpdatesAction.PATH)
             .param("tld", "multilock.uk")
             .param("dnsWriter", "multilockWriter")
-            .param("itemsCreated", "3000-01-01T00:00:00.000Z")
+            .param("requestTime", "3000-01-01T00:00:00.000Z")
             .param("enqueued", "3000-01-01T01:00:00.000Z")
             .param("domains", "another.multilock.uk")
             .param("hosts", "ns3.def.another.multilock.uk,ns4.another.multilock.uk")

core/src/test/java/google/registry/dns/ReadDnsRefreshRequestsActionTest.java

@@ -0,0 +1,296 @@
+// Copyright 2017 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.dns;
+
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.common.DatabaseMigrationStateSchedule.MigrationState;
+import static google.registry.model.common.DatabaseMigrationStateSchedule.set;
+import static google.registry.model.common.DatabaseMigrationStateSchedule.useUncachedForTest;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.testing.DatabaseHelper.createTld;
+import static google.registry.testing.DatabaseHelper.loadAllOf;
+import static google.registry.testing.DatabaseHelper.persistResource;
+import static google.registry.testing.DatabaseHelper.persistResources;
+import static google.registry.util.DateTimeUtils.START_OF_TIME;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.anyCollection;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.eq;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSortedMap;
+import com.google.common.collect.Ordering;
+import google.registry.dns.DnsConstants.TargetType;
+import google.registry.model.common.DnsRefreshRequest;
+import google.registry.model.common.DnsRefreshRequestTest;
+import google.registry.model.tld.Registry;
+import google.registry.persistence.transaction.JpaTestExtensions;
+import google.registry.persistence.transaction.JpaTestExtensions.JpaIntegrationTestExtension;
+import google.registry.testing.CloudTasksHelper;
+import google.registry.testing.CloudTasksHelper.TaskMatcher;
+import google.registry.testing.FakeClock;
+import java.util.Collection;
+import java.util.Optional;
+import org.joda.time.DateTime;
+import org.joda.time.DateTimeZone;
+import org.joda.time.Duration;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+import org.mockito.ArgumentCaptor;
+
+/** Unit tests for {@link DnsRefreshRequestTest}. */
+public class ReadDnsRefreshRequestsActionTest {
+
+  private final FakeClock clock = new FakeClock(DateTime.parse("2020-02-02T01:23:45Z"));
+  private final CloudTasksHelper cloudTasksHelper = new CloudTasksHelper(clock);
+  private final DnsUtils dnsUtils = new DnsUtils(null);
+  private final Optional<Integer> jitterSeconds = Optional.of(5);
+
+  @RegisterExtension
+  JpaIntegrationTestExtension jpa =
+      new JpaTestExtensions.Builder().withClock(clock).buildIntegrationTestExtension();
+
+  private final ReadDnsRefreshRequestsAction action =
+      spy(
+          new ReadDnsRefreshRequestsAction(
+              2,
+              Duration.standardSeconds(10),
+              jitterSeconds,
+              "tld",
+              clock,
+              dnsUtils,
+              null,
+              cloudTasksHelper.getTestCloudTasksUtils()));
+
+  private ImmutableList<DnsRefreshRequest> requests;
+
+  @BeforeAll
+  static void beforeAll() {
+    useUncachedForTest();
+  }
+
+  @BeforeEach
+  void beforeEach() {
+    useDnsSql();
+    persistResource(
+        createTld("tld")
+            .asBuilder()
+            .setDnsWriters(ImmutableSet.of("FooWriter", "BarWriter"))
+            .setNumDnsPublishLocks(2)
+            .build());
+    requests =
+        new ImmutableList.Builder<DnsRefreshRequest>()
+            .add(new DnsRefreshRequest(TargetType.DOMAIN, "domain.tld", "tld", clock.nowUtc()))
+            .add(
+                new DnsRefreshRequest(
+                    TargetType.HOST, "ns1.domain.tld", "tld", clock.nowUtc().minusMinutes(1)))
+            .add(
+                new DnsRefreshRequest(
+                    TargetType.DOMAIN, "future.tld", "tld", clock.nowUtc().plusMinutes(1)))
+            .build();
+    clock.advanceBy(Duration.standardMinutes(5));
+    persistResources(requests);
+    requests = loadAllOf(DnsRefreshRequest.class);
+  }
+
+  @Test
+  void testSuccess_runAction_pausedTld() {
+    persistResource(createTld("tld").asBuilder().setDnsPaused(true).build());
+    action.run();
+    verify(action, never()).enqueueUpdates(anyInt(), anyInt(), anyCollection());
+    verify(action, never()).processRequests(anyCollection());
+  }
+
+  @Test
+  void testSuccess_runAction_requestTimeInTheFuture() {
+    clock.setTo(DateTime.parse("2000-01-01T00:00:00Z"));
+    action.run();
+    verify(action, never()).enqueueUpdates(anyInt(), anyInt(), anyCollection());
+    verify(action, never()).processRequests(anyCollection());
+  }
+
+  @Test
+  void testSuccess_runAction_oneBatch() {
+    // The read batch size is 2 * 2 = 4. All requests will be in the same bucket, even though the
+    // bucket size should be roughly 2. But this is as expected because getLockIndex() only evenly
+    // distributes tasks in a statistical sense.
+    doReturn(2).when(action).getLockIndex(anyInt(), any(DnsRefreshRequest.class));
+    action.run();
+    @SuppressWarnings("unchecked")
+    ArgumentCaptor<Collection<DnsRefreshRequest>> captor =
+        ArgumentCaptor.forClass(Collection.class);
+    verify(action, times(1)).enqueueUpdates(eq(2), eq(2), captor.capture());
+    assertThat(captor.getValue().size()).isEqualTo(3);
+    verify(action, times(1)).processRequests(captor.capture());
+    assertThat(captor.getValue().size()).isEqualTo(3);
+    assertThat(loadAllOf(DnsRefreshRequest.class).isEmpty()).isTrue();
+  }
+
+  @Test
+  void testSuccess_runAction_twoBatches() {
+    // Make the read batch size 2 * 1 = 2.
+    persistResource(Registry.get("tld").asBuilder().setNumDnsPublishLocks(1).build());
+    doReturn(1).when(action).getLockIndex(anyInt(), any(DnsRefreshRequest.class));
+    doAnswer(
+            invocation -> {
+              @SuppressWarnings("unchecked")
+              ImmutableList<DnsRefreshRequest> ans =
+                  (ImmutableList<DnsRefreshRequest>) invocation.callRealMethod();
+              // The next read should not time out as we only increment by one millisecond, whereas
+              // the timeout is set to 10 seconds.
+              clock.advanceOneMilli();
+              return ans;
+            })
+        .when(action)
+        .processRequests(anyCollection());
+    action.run();
+    @SuppressWarnings("unchecked")
+    ArgumentCaptor<Collection<DnsRefreshRequest>> captor =
+        ArgumentCaptor.forClass(Collection.class);
+    verify(action, times(2)).enqueueUpdates(eq(1), eq(1), captor.capture());
+    assertThat(captor.getAllValues().get(0).size()).isEqualTo(2);
+    assertThat(captor.getAllValues().get(1).size()).isEqualTo(1);
+    verify(action, times(2)).processRequests(captor.capture());
+    assertThat(captor.getAllValues().get(0).size()).isEqualTo(2);
+    assertThat(captor.getAllValues().get(1).size()).isEqualTo(1);
+    assertThat(loadAllOf(DnsRefreshRequest.class).isEmpty()).isTrue();
+  }
+
+  @Test
+  void testSuccess_runAction_timeOutAfterFirstRead() {
+    // Make the process batch size 2 * 1 = 2.
+    persistResource(Registry.get("tld").asBuilder().setNumDnsPublishLocks(1).build());
+    // Both requests in the first batch will be bucketed to the same bucket.
+    doReturn(1).when(action).getLockIndex(anyInt(), any(DnsRefreshRequest.class));
+    doAnswer(
+            invocation -> {
+              @SuppressWarnings("unchecked")
+              ImmutableList<DnsRefreshRequest> ans =
+                  (ImmutableList<DnsRefreshRequest>) invocation.callRealMethod();
+              // After this function is called once, the loop in run() should top when it checks
+              // if the current time is before the request end time.
+              clock.advanceBy(Duration.standardHours(1));
+              return ans;
+            })
+        .when(action)
+        .processRequests(anyCollection());
+    action.run();
+    verify(action, times(1)).enqueueUpdates(anyInt(), anyInt(), anyCollection());
+    verify(action, times(1)).processRequests(anyCollection());
+    // The third request is left untouched because it is not read;
+    ImmutableList<DnsRefreshRequest> remainingRequests = loadAllOf(DnsRefreshRequest.class);
+    assertThat(remainingRequests.size()).isEqualTo(1);
+    assertThat(remainingRequests.get(0).getLastProcessTime()).isEqualTo(START_OF_TIME);
+  }
+
+  @Test
+  void testSuccess_processTasks() {
+    doReturn(2)
+        .doReturn(1)
+        .doReturn(2)
+        .when(action)
+        .getLockIndex(eq(2), any(DnsRefreshRequest.class));
+    action.processRequests(requests);
+    verify(action).enqueueUpdates(2, 2, ImmutableSet.of(requests.get(0), requests.get(2)));
+    verify(action).enqueueUpdates(1, 2, ImmutableSet.of(requests.get(1)));
+    assertThat(loadAllOf(DnsRefreshRequest.class)).isEmpty();
+  }
+
+  @Test
+  void testSuccess_processTasks_enqueueFailed_tasksNotDeleted() {
+    doReturn(2)
+        .doReturn(1)
+        .doReturn(2)
+        .when(action)
+        .getLockIndex(eq(2), any(DnsRefreshRequest.class));
+    doThrow(new RuntimeException("Something went wrong!"))
+        .when(action)
+        .enqueueUpdates(eq(2), eq(2), anyCollection());
+    action.processRequests(requests);
+    verify(action).enqueueUpdates(2, 2, ImmutableSet.of(requests.get(0), requests.get(2)));
+    verify(action).enqueueUpdates(1, 2, ImmutableSet.of(requests.get(1)));
+    assertThat(loadAllOf(DnsRefreshRequest.class).size()).isEqualTo(2);
+  }
+
+  @Test
+  void testSuccess_enqueueTasks() {
+    action.enqueueUpdates(2, 3, requests);
+    cloudTasksHelper.assertTasksEnqueued(
+        "dns-publish",
+        new TaskMatcher()
+            .url("/_dr/task/publishDnsUpdates")
+            .service("BACKEND")
+            .param("tld", "tld")
+            .param("dnsWriter", "FooWriter")
+            .param("lockIndex", "2")
+            .param("numPublishLocks", "3")
+            .param("enqueued", clock.nowUtc().toString())
+            .param("requestTime", clock.nowUtc().minusMinutes(6).toString())
+            .param("domains", "domain.tld,future.tld")
+            .param("hosts", "ns1.domain.tld"),
+        new TaskMatcher()
+            .url("/_dr/task/publishDnsUpdates")
+            .service("BACKEND")
+            .param("tld", "tld")
+            .param("dnsWriter", "BarWriter")
+            .param("lockIndex", "2")
+            .param("numPublishLocks", "3")
+            .param("enqueued", clock.nowUtc().toString())
+            .param("requestTime", clock.nowUtc().minusMinutes(6).toString())
+            .param("domains", "domain.tld,future.tld")
+            .param("hosts", "ns1.domain.tld"));
+    cloudTasksHelper
+        .getTestTasksFor("dns-publish")
+        .forEach(
+            task -> {
+              DateTime scheduledTime =
+                  new DateTime(task.getScheduleTime().getSeconds() * 1000, DateTimeZone.UTC);
+              assertThat(new Duration(clock.nowUtc(), scheduledTime))
+                  .isAtMost(Duration.standardSeconds(jitterSeconds.get()));
+            });
+  }
+
+  private void useDnsSql() {
+    DateTime currentTime = clock.nowUtc();
+    clock.setTo(START_OF_TIME);
+    tm().transact(
+            () ->
+                set(
+                    new ImmutableSortedMap.Builder<DateTime, MigrationState>(Ordering.natural())
+                        .put(START_OF_TIME, MigrationState.DATASTORE_ONLY)
+                        .put(START_OF_TIME.plusMillis(1), MigrationState.DATASTORE_PRIMARY)
+                        .put(START_OF_TIME.plusMillis(2), MigrationState.DATASTORE_PRIMARY_NO_ASYNC)
+                        .put(
+                            START_OF_TIME.plusMillis(3), MigrationState.DATASTORE_PRIMARY_READ_ONLY)
+                        .put(START_OF_TIME.plusMillis(4), MigrationState.SQL_PRIMARY_READ_ONLY)
+                        .put(START_OF_TIME.plusMillis(5), MigrationState.SQL_PRIMARY)
+                        .put(START_OF_TIME.plusMillis(6), MigrationState.SQL_ONLY)
+                        .put(START_OF_TIME.plusMillis(7), MigrationState.SEQUENCE_BASED_ALLOCATE_ID)
+                        .put(START_OF_TIME.plusMillis(8), MigrationState.NORDN_SQL)
+                        .put(START_OF_TIME.plusMillis(9), MigrationState.DNS_SQL)
+                        .build()));
+    clock.setTo(currentTime);
+  }
+}