Split failing dns update batches and kill after 10 retries (#1664)

* Split failing dns update batches and kill after 10 retries

* format fixes

* Add another test

* Switch to CloudTasks

* Change back to app engine header

* Change to immutableList and other changes

* Change to optional header

* Add bug ID to todo

* Switch to constructor injection

* Remove old queue

* Set response status

* Change to Optional<Integer>

* Rename action status

* Switched to use CLoudTaskHelper

* Remove spy in test

core/src/main/java/google/registry/dns/DnsMetrics.java

@@ -40,7 +40,14 @@ public class DnsMetrics {
   public enum CommitStatus { SUCCESS, FAILURE }
 
   /** Disposition of the publish action. */
-  public enum ActionStatus { SUCCESS, COMMIT_FAILURE, LOCK_FAILURE, BAD_WRITER, BAD_LOCK_INDEX }
+  public enum ActionStatus {
+    SUCCESS,
+    COMMIT_FAILURE,
+    LOCK_FAILURE,
+    BAD_WRITER,
+    BAD_LOCK_INDEX,
+    MAX_RETRIES_EXCEEDED
+  }
 
   private static final ImmutableSet<LabelDescriptor> LABEL_DESCRIPTORS_FOR_PUBLISH_REQUESTS =
       ImmutableSet.of(

core/src/main/java/google/registry/dns/PublishDnsUpdatesAction.java

@@ -14,6 +14,7 @@
 
 package google.registry.dns;
 
+import static google.registry.dns.DnsConstants.DNS_PUBLISH_PUSH_QUEUE_NAME;
 import static google.registry.dns.DnsModule.PARAM_DNS_WRITER;
 import static google.registry.dns.DnsModule.PARAM_DOMAINS;
 import static google.registry.dns.DnsModule.PARAM_HOSTS;
@@ -24,7 +25,11 @@ import static google.registry.dns.DnsModule.PARAM_REFRESH_REQUEST_CREATED;
 import static google.registry.request.Action.Method.POST;
 import static google.registry.request.RequestParameters.PARAM_TLD;
 import static google.registry.util.CollectionUtils.nullToEmpty;
+import static javax.servlet.http.HttpServletResponse.SC_ACCEPTED;
 
+import com.google.common.base.Joiner;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMultimap;
 import com.google.common.flogger.FluentLogger;
 import com.google.common.net.InternetDomainName;
 import google.registry.config.RegistryConfig.Config;
@@ -34,12 +39,17 @@ import google.registry.dns.DnsMetrics.PublishStatus;
 import google.registry.dns.writer.DnsWriter;
 import google.registry.model.tld.Registry;
 import google.registry.request.Action;
+import google.registry.request.Action.Service;
+import google.registry.request.Header;
 import google.registry.request.HttpException.ServiceUnavailableException;
 import google.registry.request.Parameter;
+import google.registry.request.Response;
 import google.registry.request.auth.Auth;
 import google.registry.request.lock.LockHandler;
 import google.registry.util.Clock;
+import google.registry.util.CloudTasksUtils;
 import google.registry.util.DomainNameUtils;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.Callable;
 import javax.inject.Inject;
@@ -57,13 +67,19 @@ public final class PublishDnsUpdatesAction implements Runnable, Callable<Void> {
 
   public static final String PATH = "/_dr/task/publishDnsUpdates";
   public static final String LOCK_NAME = "DNS updates";
+  // TODO(b/236726584): Remove App Engine header once CloudTasksUtils is refactored to create HTTP
+  // tasks.
+  public static final String APP_ENGINE_RETRY_HEADER = "X-AppEngine-TaskRetryCount";
+  public static final String CLOUD_TASKS_RETRY_HEADER = "X-CloudTasks-TaskRetryCount";
+  public static final int RETRIES_BEFORE_PERMANENT_FAILURE = 10;
 
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
-  @Inject DnsQueue dnsQueue;
+  private final DnsQueue dnsQueue;
-  @Inject DnsWriterProxy dnsWriterProxy;
+  private final DnsWriterProxy dnsWriterProxy;
-  @Inject DnsMetrics dnsMetrics;
+  private final DnsMetrics dnsMetrics;
-  @Inject @Config("publishDnsUpdatesLockDuration") Duration timeout;
+  private final Duration timeout;
+  private final int retryCount;
 
   /**
    * The DNS writer to use for this batch.
@@ -73,24 +89,61 @@ public final class PublishDnsUpdatesAction implements Runnable, Callable<Void> {
    * writers configured in {@link Registry#getDnsWriters()}, as of the time the batch was written
    * out (and not necessarily currently).
    */
-  @Inject @Parameter(PARAM_DNS_WRITER) String dnsWriter;
+  private final String dnsWriter;
+
+  private final DateTime enqueuedTime;
+  private final DateTime itemsCreateTime;
+  private final int lockIndex;
+  private final int numPublishLocks;
+  private final Set<String> domains;
+  private final Set<String> hosts;
+  private final String tld;
+  private final LockHandler lockHandler;
+  private final Clock clock;
+  private final CloudTasksUtils cloudTasksUtils;
+  private final Response response;
 
   @Inject
-  @Parameter(PARAM_PUBLISH_TASK_ENQUEUED)
+  public PublishDnsUpdatesAction(
-  DateTime enqueuedTime;
+      @Parameter(PARAM_DNS_WRITER) String dnsWriter,
-
+      @Parameter(PARAM_PUBLISH_TASK_ENQUEUED) DateTime enqueuedTime,
-  @Inject
+      @Parameter(PARAM_REFRESH_REQUEST_CREATED) DateTime itemsCreateTime,
-  @Parameter(PARAM_REFRESH_REQUEST_CREATED)
+      @Parameter(PARAM_LOCK_INDEX) int lockIndex,
-  DateTime itemsCreateTime;
+      @Parameter(PARAM_NUM_PUBLISH_LOCKS) int numPublishLocks,
-
+      @Parameter(PARAM_DOMAINS) Set<String> domains,
-  @Inject @Parameter(PARAM_LOCK_INDEX) int lockIndex;
+      @Parameter(PARAM_HOSTS) Set<String> hosts,
-  @Inject @Parameter(PARAM_NUM_PUBLISH_LOCKS) int numPublishLocks;
+      @Parameter(PARAM_TLD) String tld,
-  @Inject @Parameter(PARAM_DOMAINS) Set<String> domains;
+      @Config("publishDnsUpdatesLockDuration") Duration timeout,
-  @Inject @Parameter(PARAM_HOSTS) Set<String> hosts;
+      @Header(APP_ENGINE_RETRY_HEADER) Optional<Integer> appEngineRetryCount,
-  @Inject @Parameter(PARAM_TLD) String tld;
+      @Header(CLOUD_TASKS_RETRY_HEADER) Optional<Integer> cloudTasksRetryCount,
-  @Inject LockHandler lockHandler;
+      DnsQueue dnsQueue,
-  @Inject Clock clock;
+      DnsWriterProxy dnsWriterProxy,
-  @Inject PublishDnsUpdatesAction() {}
+      DnsMetrics dnsMetrics,
+      LockHandler lockHandler,
+      Clock clock,
+      CloudTasksUtils cloudTasksUtils,
+      Response response) {
+    this.dnsQueue = dnsQueue;
+    this.dnsWriterProxy = dnsWriterProxy;
+    this.dnsMetrics = dnsMetrics;
+    this.timeout = timeout;
+    this.retryCount =
+        cloudTasksRetryCount.orElse(
+            appEngineRetryCount.orElseThrow(
+                () -> new IllegalStateException("Missing a valid retry count header")));
+    this.dnsWriter = dnsWriter;
+    this.enqueuedTime = enqueuedTime;
+    this.itemsCreateTime = itemsCreateTime;
+    this.lockIndex = lockIndex;
+    this.numPublishLocks = numPublishLocks;
+    this.domains = domains;
+    this.hosts = hosts;
+    this.tld = tld;
+    this.lockHandler = lockHandler;
+    this.clock = clock;
+    this.cloudTasksUtils = cloudTasksUtils;
+    this.response = response;
+  }
 
   private void recordActionResult(ActionStatus status) {
     DateTime now = clock.nowUtc();
@@ -122,9 +175,9 @@ public final class PublishDnsUpdatesAction implements Runnable, Callable<Void> {
       return;
     }
     // If executeWithLocks fails to get the lock, it does not throw an exception, simply returns
-    // false. We need to make sure to take note of this error; otherwise, a failed lock might result
+    // false. We need to make sure to take note of this error; otherwise, a failed lock might
-    // in the update task being dequeued and dropped. A message will already have been logged
+    // result in the update task being dequeued and dropped. A message will already have been
-    // to indicate the problem.
+    // logged to indicate the problem.
     if (!lockHandler.executeWithLocks(
         this,
         tld,
@@ -138,10 +191,64 @@ public final class PublishDnsUpdatesAction implements Runnable, Callable<Void> {
   /** Runs the task, with the lock. */
   @Override
   public Void call() {
-    processBatch();
+    try {
+      processBatch();
+    } catch (Throwable e) {
+      // Retry the batch 3 times
+      if (retryCount < 3) {
+        throw e;
+      }
+      // After 3 retries, split the batch
+      if (domains.size() > 1 || hosts.size() > 1) {
+        // split batch and requeue
+        splitBatch();
+      } else if (domains.size() == 1 && hosts.size() == 1) {
+        // Enqueue the single domain and single host separately
+        enqueue(ImmutableList.copyOf(domains), ImmutableList.of());
+        enqueue(ImmutableList.of(), ImmutableList.copyOf(hosts));
+      }
+      // If the batch only contains 1 name, allow it more retries
+      else if (retryCount < RETRIES_BEFORE_PERMANENT_FAILURE) {
+        throw e;
+      }
+      // If we get here, we should terminate this task as it is likely a perpetually failing task.
+      // TODO(b/237302821): Send an email notifying partner the dns update failed
+      recordActionResult(ActionStatus.MAX_RETRIES_EXCEEDED);
+      response.setStatus(SC_ACCEPTED);
+      logger.atSevere().withCause(e).log("Terminated task after too many retries");
+    }
     return null;
   }
 
+  /** Splits the domains and hosts in a batch into smaller batches and adds them to the queue. */
+  private void splitBatch() {
+    ImmutableList<String> domainList = ImmutableList.copyOf(domains);
+    ImmutableList<String> hostList = ImmutableList.copyOf(hosts);
+
+    enqueue(domainList.subList(0, domains.size() / 2), hostList.subList(0, hosts.size() / 2));
+    enqueue(
+        domainList.subList(domains.size() / 2, domains.size()),
+        hostList.subList(hosts.size() / 2, hosts.size()));
+  }
+
+  private void enqueue(ImmutableList<String> domains, ImmutableList<String> hosts) {
+    cloudTasksUtils.enqueue(
+        DNS_PUBLISH_PUSH_QUEUE_NAME,
+        cloudTasksUtils.createPostTask(
+            PATH,
+            Service.BACKEND.toString(),
+            ImmutableMultimap.<String, String>builder()
+                .put(PARAM_TLD, tld)
+                .put(PARAM_DNS_WRITER, dnsWriter)
+                .put(PARAM_LOCK_INDEX, Integer.toString(lockIndex))
+                .put(PARAM_NUM_PUBLISH_LOCKS, Integer.toString(numPublishLocks))
+                .put(PARAM_PUBLISH_TASK_ENQUEUED, clock.nowUtc().toString())
+                .put(PARAM_REFRESH_REQUEST_CREATED, itemsCreateTime.toString())
+                .put(PARAM_DOMAINS, Joiner.on(",").join(domains))
+                .put(PARAM_HOSTS, Joiner.on(",").join(hosts))
+                .build()));
+  }
+
   /** Adds all the domains and hosts in the batch back to the queue to be processed later. */
   private void requeueBatch() {
     logger.atInfo().log("Requeueing batch for retry.");

core/src/main/java/google/registry/request/RequestModule.java

@@ -15,8 +15,11 @@
 package google.registry.request;
 
 import static com.google.common.net.MediaType.JSON_UTF_8;
+import static google.registry.dns.PublishDnsUpdatesAction.APP_ENGINE_RETRY_HEADER;
+import static google.registry.dns.PublishDnsUpdatesAction.CLOUD_TASKS_RETRY_HEADER;
 import static google.registry.model.tld.Registries.assertTldExists;
 import static google.registry.model.tld.Registries.assertTldsExist;
+import static google.registry.request.RequestParameters.extractOptionalHeader;
 import static google.registry.request.RequestParameters.extractRequiredParameter;
 import static google.registry.request.RequestParameters.extractSetOfParameters;
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -40,6 +43,7 @@ import google.registry.util.RequestStatusChecker;
 import google.registry.util.RequestStatusCheckerImpl;
 import java.io.IOException;
 import java.util.Map;
+import java.util.Optional;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 import javax.servlet.http.HttpSession;
@@ -245,4 +249,16 @@ public final class RequestModule {
     }
     return params.build();
   }
+
+  @Provides
+  @Header(APP_ENGINE_RETRY_HEADER)
+  static Optional<Integer> provideAppEngineRetryCount(HttpServletRequest req) {
+    return extractOptionalHeader(req, APP_ENGINE_RETRY_HEADER).map(Integer::parseInt);
+  }
+
+  @Provides
+  @Header(CLOUD_TASKS_RETRY_HEADER)
+  static Optional<Integer> provideCloudTasksRetryCount(HttpServletRequest req) {
+    return extractOptionalHeader(req, CLOUD_TASKS_RETRY_HEADER).map(Integer::parseInt);
+  }
 }

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

@@ -15,10 +15,21 @@
 package google.registry.dns;
 
 import static com.google.common.truth.Truth.assertThat;
+import static google.registry.dns.DnsConstants.DNS_PUBLISH_PUSH_QUEUE_NAME;
+import static google.registry.dns.DnsModule.PARAM_DNS_WRITER;
+import static google.registry.dns.DnsModule.PARAM_DOMAINS;
+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.request.RequestParameters.PARAM_TLD;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.persistActiveDomain;
 import static google.registry.testing.DatabaseHelper.persistActiveSubordinateHost;
 import static google.registry.testing.DatabaseHelper.persistResource;
+import static javax.servlet.http.HttpServletResponse.SC_ACCEPTED;
+import static javax.servlet.http.HttpServletResponse.SC_OK;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doThrow;
@@ -39,9 +50,14 @@ import google.registry.model.tld.Registry;
 import google.registry.request.HttpException.ServiceUnavailableException;
 import google.registry.request.lock.LockHandler;
 import google.registry.testing.AppEngineExtension;
+import google.registry.testing.CloudTasksHelper;
+import google.registry.testing.CloudTasksHelper.TaskMatcher;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeLockHandler;
+import google.registry.testing.FakeResponse;
 import google.registry.testing.InjectExtension;
+import java.util.Optional;
+import java.util.Set;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
 import org.junit.jupiter.api.BeforeEach;
@@ -57,10 +73,12 @@ public class PublishDnsUpdatesActionTest {
 
   @RegisterExtension public final InjectExtension inject = new InjectExtension();
   private final FakeClock clock = new FakeClock(DateTime.parse("1971-01-01TZ"));
+  private final FakeResponse response = new FakeResponse();
   private final FakeLockHandler lockHandler = new FakeLockHandler(true);
   private final DnsWriter dnsWriter = mock(DnsWriter.class);
   private final DnsMetrics dnsMetrics = mock(DnsMetrics.class);
   private final DnsQueue dnsQueue = mock(DnsQueue.class);
+  private final CloudTasksHelper cloudTasksHelper = new CloudTasksHelper();
   private PublishDnsUpdatesAction action;
 
   @BeforeEach
@@ -80,30 +98,65 @@ public class PublishDnsUpdatesActionTest {
     clock.advanceOneMilli();
   }
 
-  private PublishDnsUpdatesAction createAction(String tld) {
+  private PublishDnsUpdatesAction createAction(String tld, Set<String> domains, Set<String> hosts) {
-    PublishDnsUpdatesAction action = new PublishDnsUpdatesAction();
+    return createAction(tld, domains, hosts, 0, "correctWriter", 1, 1, lockHandler);
-    action.timeout = Duration.standardSeconds(10);
+  }
-    action.tld = tld;
+
-    action.hosts = ImmutableSet.of();
+  private PublishDnsUpdatesAction createAction(
-    action.domains = ImmutableSet.of();
+      String tld, Set<String> domains, Set<String> hosts, Integer retryCount) {
-    action.itemsCreateTime = clock.nowUtc().minusHours(2);
+    return createAction(tld, domains, hosts, retryCount, "correctWriter", 1, 1, lockHandler);
-    action.enqueuedTime = clock.nowUtc().minusHours(1);
+  }
-    action.dnsWriter = "correctWriter";
+
-    action.dnsWriterProxy = new DnsWriterProxy(ImmutableMap.of("correctWriter", dnsWriter));
+  private PublishDnsUpdatesAction createActionBadDnsWriter(
-    action.dnsMetrics = dnsMetrics;
+      String tld, Set<String> domains, Set<String> hosts) {
-    action.dnsQueue = dnsQueue;
+    return createAction(tld, domains, hosts, 0, "wrongWriter", 1, 1, lockHandler);
-    action.lockIndex = 1;
+  }
-    action.numPublishLocks = 1;
+
-    action.lockHandler = lockHandler;
+  private PublishDnsUpdatesAction createActionWithCustomLocks(
-    action.clock = clock;
+      String tld,
-    return action;
+      Set<String> domains,
+      Set<String> hosts,
+      int lockIndex,
+      int numPublishLocks,
+      LockHandler lockHandler) {
+    return createAction(
+        tld, domains, hosts, 0, "correctWriter", lockIndex, numPublishLocks, lockHandler);
+  }
+
+  private PublishDnsUpdatesAction createAction(
+      String tld,
+      Set<String> domains,
+      Set<String> hosts,
+      Integer retryCount,
+      String dnsWriterString,
+      int lockIndex,
+      int numPublishLocks,
+      LockHandler lockHandler) {
+    return new PublishDnsUpdatesAction(
+        dnsWriterString,
+        clock.nowUtc().minusHours(1),
+        clock.nowUtc().minusHours(2),
+        lockIndex,
+        numPublishLocks,
+        domains,
+        hosts,
+        tld,
+        Duration.standardSeconds(10),
+        Optional.ofNullable(retryCount),
+        Optional.empty(),
+        dnsQueue,
+        new DnsWriterProxy(ImmutableMap.of("correctWriter", dnsWriter)),
+        dnsMetrics,
+        lockHandler,
+        clock,
+        cloudTasksHelper.getTestCloudTasksUtils(),
+        response);
   }
 
   @Test
   void testHost_published() {
-    action = createAction("xn--q9jyb4c");
+    action =
-    action.hosts = ImmutableSet.of("ns1.example.xn--q9jyb4c");
+        createAction("xn--q9jyb4c", ImmutableSet.of(), ImmutableSet.of("ns1.example.xn--q9jyb4c"));
-
     action.run();
 
     verify(dnsWriter).publishHost("ns1.example.xn--q9jyb4c");
@@ -125,13 +178,12 @@ public class PublishDnsUpdatesActionTest {
             Duration.standardHours(1));
     verifyNoMoreInteractions(dnsMetrics);
     verifyNoMoreInteractions(dnsQueue);
+    assertThat(response.getStatus()).isEqualTo(SC_OK);
   }
 
   @Test
   void testDomain_published() {
-    action = createAction("xn--q9jyb4c");
+    action = createAction("xn--q9jyb4c", ImmutableSet.of("example.xn--q9jyb4c"), ImmutableSet.of());
-    action.domains = ImmutableSet.of("example.xn--q9jyb4c");
-
     action.run();
 
     verify(dnsWriter).publishDomain("example.xn--q9jyb4c");
@@ -153,18 +205,22 @@ public class PublishDnsUpdatesActionTest {
             Duration.standardHours(1));
     verifyNoMoreInteractions(dnsMetrics);
     verifyNoMoreInteractions(dnsQueue);
+    assertThat(response.getStatus()).isEqualTo(SC_OK);
   }
 
   @Test
   void testAction_acquiresCorrectLock() {
     persistResource(Registry.get("xn--q9jyb4c").asBuilder().setNumDnsPublishLocks(4).build());
-    action = createAction("xn--q9jyb4c");
-    action.lockIndex = 2;
-    action.numPublishLocks = 4;
-    action.domains = ImmutableSet.of("example.xn--q9jyb4c");
     LockHandler mockLockHandler = mock(LockHandler.class);
     when(mockLockHandler.executeWithLocks(any(), any(), any(), any())).thenReturn(true);
-    action.lockHandler = mockLockHandler;
+    action =
+        createActionWithCustomLocks(
+            "xn--q9jyb4c",
+            ImmutableSet.of("example.xn--q9jyb4c"),
+            ImmutableSet.of(),
+            2,
+            4,
+            mockLockHandler);
 
     action.run();
 
@@ -175,11 +231,13 @@ public class PublishDnsUpdatesActionTest {
 
   @Test
   void testPublish_commitFails() {
-    action = createAction("xn--q9jyb4c");
+    ImmutableSet<String> hosts =
-    action.domains = ImmutableSet.of("example.xn--q9jyb4c", "example2.xn--q9jyb4c");
-    action.hosts =
         ImmutableSet.of(
             "ns1.example.xn--q9jyb4c", "ns2.example.xn--q9jyb4c", "ns1.example2.xn--q9jyb4c");
+    action =
+        createAction(
+            "xn--q9jyb4c", ImmutableSet.of("example.xn--q9jyb4c", "example2.xn--q9jyb4c"), hosts);
+
     doThrow(new RuntimeException()).when(dnsWriter).commit();
 
     assertThrows(RuntimeException.class, action::run);
@@ -202,13 +260,153 @@ public class PublishDnsUpdatesActionTest {
     verifyNoMoreInteractions(dnsQueue);
   }
 
+  @Test
+  void testTaskFails_splitsBatch() {
+    ImmutableSet<String> domains =
+        ImmutableSet.of(
+            "example1.xn--q9jyb4c",
+            "example2.xn--q9jyb4c",
+            "example3.xn--q9jyb4c",
+            "example4.xn--q9jyb4c");
+    action = createAction("xn--q9jyb4c", domains, ImmutableSet.of("ns1.example.xn--q9jyb4c"), 3);
+    doThrow(new RuntimeException()).when(dnsWriter).commit();
+    action.run();
+
+    cloudTasksHelper.assertTasksEnqueued(
+        DNS_PUBLISH_PUSH_QUEUE_NAME,
+        new TaskMatcher()
+            .url(PublishDnsUpdatesAction.PATH)
+            .param(PARAM_TLD, "xn--q9jyb4c")
+            .param(PARAM_DNS_WRITER, "correctWriter")
+            .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_DOMAINS, "example1.xn--q9jyb4c,example2.xn--q9jyb4c")
+            .param(PARAM_HOSTS, "")
+            .header("content-type", "application/x-www-form-urlencoded"),
+        new TaskMatcher()
+            .url(PublishDnsUpdatesAction.PATH)
+            .param(PARAM_TLD, "xn--q9jyb4c")
+            .param(PARAM_DNS_WRITER, "correctWriter")
+            .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_DOMAINS, "example3.xn--q9jyb4c,example4.xn--q9jyb4c")
+            .param(PARAM_HOSTS, "ns1.example.xn--q9jyb4c")
+            .header("content-type", "application/x-www-form-urlencoded"));
+  }
+
+  @Test
+  void testTaskFails_splitsBatch5Names() {
+    ImmutableSet<String> domains =
+        ImmutableSet.of(
+            "example1.xn--q9jyb4c",
+            "example2.xn--q9jyb4c",
+            "example3.xn--q9jyb4c",
+            "example4.xn--q9jyb4c",
+            "example5.xn--q9jyb4c");
+    action = createAction("xn--q9jyb4c", domains, ImmutableSet.of("ns1.example.xn--q9jyb4c"), 3);
+    doThrow(new RuntimeException()).when(dnsWriter).commit();
+    action.run();
+
+    cloudTasksHelper.assertTasksEnqueued(
+        DNS_PUBLISH_PUSH_QUEUE_NAME,
+        new TaskMatcher()
+            .url(PublishDnsUpdatesAction.PATH)
+            .param(PARAM_TLD, "xn--q9jyb4c")
+            .param(PARAM_DNS_WRITER, "correctWriter")
+            .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_DOMAINS, "example1.xn--q9jyb4c,example2.xn--q9jyb4c")
+            .param(PARAM_HOSTS, "")
+            .header("content-type", "application/x-www-form-urlencoded"),
+        new TaskMatcher()
+            .url(PublishDnsUpdatesAction.PATH)
+            .param(PARAM_TLD, "xn--q9jyb4c")
+            .param(PARAM_DNS_WRITER, "correctWriter")
+            .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_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"));
+  }
+
+  @Test
+  void testTaskFails_singleHostSingleDomain() {
+    action =
+        createAction(
+            "xn--q9jyb4c",
+            ImmutableSet.of("example1.xn--q9jyb4c"),
+            ImmutableSet.of("ns1.example.xn--q9jyb4c"),
+            3);
+    doThrow(new RuntimeException()).when(dnsWriter).commit();
+    action.run();
+
+    cloudTasksHelper.assertTasksEnqueued(
+        DNS_PUBLISH_PUSH_QUEUE_NAME,
+        new TaskMatcher()
+            .url(PublishDnsUpdatesAction.PATH)
+            .param(PARAM_TLD, "xn--q9jyb4c")
+            .param(PARAM_DNS_WRITER, "correctWriter")
+            .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_DOMAINS, "example1.xn--q9jyb4c")
+            .param(PARAM_HOSTS, "")
+            .header("content-type", "application/x-www-form-urlencoded"),
+        new TaskMatcher()
+            .url(PublishDnsUpdatesAction.PATH)
+            .param(PARAM_TLD, "xn--q9jyb4c")
+            .param(PARAM_DNS_WRITER, "correctWriter")
+            .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_DOMAINS, "")
+            .param(PARAM_HOSTS, "ns1.example.xn--q9jyb4c")
+            .header("content-type", "application/x-www-form-urlencoded"));
+  }
+
+  @Test
+  void testTaskFailsAfterTenRetries_doesNotRetry() {
+    action =
+        createAction(
+            "xn--q9jyb4c", ImmutableSet.of(), ImmutableSet.of("ns1.example.xn--q9jyb4c"), 10);
+    doThrow(new RuntimeException()).when(dnsWriter).commit();
+    action.run();
+    cloudTasksHelper.assertNoTasksEnqueued(DNS_PUBLISH_PUSH_QUEUE_NAME);
+    assertThat(response.getStatus()).isEqualTo(SC_ACCEPTED);
+  }
+
+  @Test
+  void testTaskMissingRetryHeaders_throwsException() {
+    IllegalStateException thrown =
+        assertThrows(
+            IllegalStateException.class,
+            () ->
+                createAction(
+                    "xn--q9jyb4c",
+                    ImmutableSet.of(),
+                    ImmutableSet.of("ns1.example.xn--q9jyb4c"),
+                    null));
+    assertThat(thrown).hasMessageThat().contains("Missing a valid retry count header");
+  }
+
   @Test
   void testHostAndDomain_published() {
-    action = createAction("xn--q9jyb4c");
+    ImmutableSet<String> hosts =
-    action.domains = ImmutableSet.of("example.xn--q9jyb4c", "example2.xn--q9jyb4c");
-    action.hosts =
         ImmutableSet.of(
             "ns1.example.xn--q9jyb4c", "ns2.example.xn--q9jyb4c", "ns1.example2.xn--q9jyb4c");
+    action =
+        createAction(
+            "xn--q9jyb4c", ImmutableSet.of("example.xn--q9jyb4c", "example2.xn--q9jyb4c"), hosts);
 
     action.run();
 
@@ -239,9 +437,11 @@ public class PublishDnsUpdatesActionTest {
 
   @Test
   void testWrongTld_notPublished() {
-    action = createAction("xn--q9jyb4c");
+    action =
-    action.domains = ImmutableSet.of("example.com", "example2.com");
+        createAction(
-    action.hosts = ImmutableSet.of("ns1.example.com", "ns2.example.com", "ns1.example2.com");
+            "xn--q9jyb4c",
+            ImmutableSet.of("example.com", "example2.com"),
+            ImmutableSet.of("ns1.example.com", "ns2.example.com", "ns1.example2.com"));
 
     action.run();
 
@@ -267,10 +467,14 @@ public class PublishDnsUpdatesActionTest {
 
   @Test
   void testLockIsntAvailable() {
-    action = createAction("xn--q9jyb4c");
+    action =
-    action.domains = ImmutableSet.of("example.com", "example2.com");
+        createActionWithCustomLocks(
-    action.hosts = ImmutableSet.of("ns1.example.com", "ns2.example.com", "ns1.example2.com");
+            "xn--q9jyb4c",
-    action.lockHandler = new FakeLockHandler(false);
+            ImmutableSet.of("example.com", "example2.com"),
+            ImmutableSet.of("ns1.example.com", "ns2.example.com", "ns1.example2.com"),
+            1,
+            1,
+            new FakeLockHandler(false));
 
     ServiceUnavailableException thrown =
         assertThrows(ServiceUnavailableException.class, action::run);
@@ -292,12 +496,14 @@ public class PublishDnsUpdatesActionTest {
   @Test
   void testParam_invalidLockIndex() {
     persistResource(Registry.get("xn--q9jyb4c").asBuilder().setNumDnsPublishLocks(4).build());
-    action = createAction("xn--q9jyb4c");
+    action =
-    action.domains = ImmutableSet.of("example.com");
+        createActionWithCustomLocks(
-    action.hosts = ImmutableSet.of("ns1.example.com");
+            "xn--q9jyb4c",
-    action.lockIndex = 5;
+            ImmutableSet.of("example.com"),
-    action.numPublishLocks = 4;
+            ImmutableSet.of("ns1.example.com"),
-
+            5,
+            4,
+            lockHandler);
     action.run();
 
     verifyNoMoreInteractions(dnsWriter);
@@ -318,12 +524,14 @@ public class PublishDnsUpdatesActionTest {
   @Test
   void testRegistryParam_mismatchedMaxLocks() {
     persistResource(Registry.get("xn--q9jyb4c").asBuilder().setNumDnsPublishLocks(4).build());
-    action = createAction("xn--q9jyb4c");
+    action =
-    action.domains = ImmutableSet.of("example.com");
+        createActionWithCustomLocks(
-    action.hosts = ImmutableSet.of("ns1.example.com");
+            "xn--q9jyb4c",
-    action.lockIndex = 3;
+            ImmutableSet.of("example.com"),
-    action.numPublishLocks = 5;
+            ImmutableSet.of("ns1.example.com"),
-
+            3,
+            5,
+            lockHandler);
     action.run();
 
     verifyNoMoreInteractions(dnsWriter);
@@ -343,11 +551,11 @@ public class PublishDnsUpdatesActionTest {
 
   @Test
   void testWrongDnsWriter() {
-    action = createAction("xn--q9jyb4c");
+    action =
-    action.domains = ImmutableSet.of("example.com", "example2.com");
+        createActionBadDnsWriter(
-    action.hosts = ImmutableSet.of("ns1.example.com", "ns2.example.com", "ns1.example2.com");
+            "xn--q9jyb4c",
-    action.dnsWriter = "wrongWriter";
+            ImmutableSet.of("example.com", "example2.com"),
-
+            ImmutableSet.of("ns1.example.com", "ns2.example.com", "ns1.example2.com"));
     action.run();
 
     verifyNoMoreInteractions(dnsWriter);