Add capability to sync DNS using multiple writers if configured

This is written in such a way that it can safely handle task items in the
old format so long as the DNS writer to use for the given TLD is unambiguous
(which it is for now, until we allow multiple DNS writers to be configured).

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=162293412

java/google/registry/dns/DnsModule.java

@@ -16,16 +16,19 @@ package google.registry.dns;
 
 import static google.registry.dns.DnsConstants.DNS_PUBLISH_PUSH_QUEUE_NAME;
 import static google.registry.dns.DnsConstants.DNS_PULL_QUEUE_NAME;
-import static google.registry.dns.PublishDnsUpdatesAction.DOMAINS_PARAM;
-import static google.registry.dns.PublishDnsUpdatesAction.HOSTS_PARAM;
-import static google.registry.dns.ReadDnsQueueAction.KEEP_TASKS_PARAM;
+import static google.registry.dns.PublishDnsUpdatesAction.PARAM_DNS_WRITER;
+import static google.registry.dns.PublishDnsUpdatesAction.PARAM_DOMAINS;
+import static google.registry.dns.PublishDnsUpdatesAction.PARAM_HOSTS;
+import static google.registry.dns.ReadDnsQueueAction.PARAM_KEEP_TASKS;
 import static google.registry.request.RequestParameters.extractBooleanParameter;
 import static google.registry.request.RequestParameters.extractEnumParameter;
+import static google.registry.request.RequestParameters.extractOptionalParameter;
 import static google.registry.request.RequestParameters.extractRequiredParameter;
 import static google.registry.request.RequestParameters.extractSetOfParameters;
 
 import com.google.appengine.api.taskqueue.Queue;
 import com.google.appengine.api.taskqueue.QueueFactory;
+import com.google.common.base.Optional;
 import dagger.Binds;
 import dagger.Module;
 import dagger.Provides;
@@ -58,21 +61,28 @@ public abstract class DnsModule {
   }
 
   @Provides
-  @Parameter(DOMAINS_PARAM)
+  @Parameter(PARAM_DNS_WRITER)
+  static Optional<String> provideDnsWriter(HttpServletRequest req) {
+    // TODO(b/63385597): Make this required after DNS writers migration is complete.
+    return extractOptionalParameter(req, PARAM_DNS_WRITER);
+  }
+
+  @Provides
+  @Parameter(PARAM_DOMAINS)
   static Set<String> provideDomains(HttpServletRequest req) {
-    return extractSetOfParameters(req, DOMAINS_PARAM);
+    return extractSetOfParameters(req, PARAM_DOMAINS);
   }
 
   @Provides
-  @Parameter(HOSTS_PARAM)
+  @Parameter(PARAM_HOSTS)
   static Set<String> provideHosts(HttpServletRequest req) {
-    return extractSetOfParameters(req, HOSTS_PARAM);
+    return extractSetOfParameters(req, PARAM_HOSTS);
   }
 
   @Provides
-  @Parameter(KEEP_TASKS_PARAM)
+  @Parameter(PARAM_KEEP_TASKS)
   static boolean provideKeepTasks(HttpServletRequest req) {
-    return extractBooleanParameter(req, KEEP_TASKS_PARAM);
+    return extractBooleanParameter(req, PARAM_KEEP_TASKS);
   }
 
   @Provides

java/google/registry/dns/DnsWriterProxy.java

@@ -15,16 +15,20 @@
 package google.registry.dns;
 
 import static com.google.common.base.Preconditions.checkState;
+import static google.registry.util.FormattingLogger.getLoggerForCallerClass;
 
 import com.google.common.collect.ImmutableMap;
 import google.registry.dns.writer.DnsWriter;
 import google.registry.model.registry.Registry;
+import google.registry.util.FormattingLogger;
 import java.util.Map;
 import javax.inject.Inject;
 
 /** Proxy for retrieving {@link DnsWriter} implementations. */
 public final class DnsWriterProxy {
 
+  private static final FormattingLogger logger = getLoggerForCallerClass();
+
   private final ImmutableMap<String, DnsWriter> dnsWriters;
 
   @Inject
@@ -32,11 +36,22 @@ public final class DnsWriterProxy {
     this.dnsWriters = ImmutableMap.copyOf(dnsWriters);
   }
 
-  /** Return the {@link DnsWriter} for the given tld. */
+  /** Returns the {@link DnsWriter} for the given tld. */
+  // TODO(b/63385597): Delete this method after DNS writers migration is complete.
+  @Deprecated
   public DnsWriter getForTld(String tld) {
-    String clazz = Registry.get(tld).getDnsWriter();
-    DnsWriter dnsWriter = dnsWriters.get(clazz);
-    checkState(dnsWriter != null, "Could not load DnsWriter %s for TLD %s", clazz, tld);
+    return getByClassNameForTld(Registry.get(tld).getDnsWriter(), tld);
+  }
+
+  /** Returns the instance of {@link DnsWriter} by class name. */
+  public DnsWriter getByClassNameForTld(String className, String tld) {
+    if (!Registry.get(tld).getDnsWriters().contains(className)) {
+      logger.warningfmt(
+          "Loaded potentially stale DNS writer %s which is no longer active on TLD %s.",
+          className, tld);
+    }
+    DnsWriter dnsWriter = dnsWriters.get(className);
+    checkState(dnsWriter != null, "Could not load DnsWriter %s for TLD %s", className, tld);
     return dnsWriter;
   }
 }

java/google/registry/dns/PublishDnsUpdatesAction.java

@@ -16,16 +16,18 @@ package google.registry.dns;
 
 import static google.registry.model.server.Lock.executeWithLocks;
 import static google.registry.request.Action.Method.POST;
+import static google.registry.request.RequestParameters.PARAM_TLD;
 import static google.registry.util.CollectionUtils.nullToEmpty;
 
+import com.google.common.base.Optional;
 import com.google.common.net.InternetDomainName;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.dns.DnsMetrics.Status;
 import google.registry.dns.writer.DnsWriter;
+import google.registry.model.registry.Registry;
 import google.registry.request.Action;
 import google.registry.request.HttpException.ServiceUnavailableException;
 import google.registry.request.Parameter;
-import google.registry.request.RequestParameters;
 import google.registry.request.auth.Auth;
 import google.registry.util.DomainNameUtils;
 import google.registry.util.FormattingLogger;
@@ -44,8 +46,9 @@ import org.joda.time.Duration;
 public final class PublishDnsUpdatesAction implements Runnable, Callable<Void> {
 
   public static final String PATH = "/_dr/task/publishDnsUpdates";
-  public static final String DOMAINS_PARAM = "domains";
-  public static final String HOSTS_PARAM = "hosts";
+  public static final String PARAM_DNS_WRITER = "dnsWriter";
+  public static final String PARAM_DOMAINS = "domains";
+  public static final String PARAM_HOSTS = "hosts";
 
   private static final FormattingLogger logger = FormattingLogger.getLoggerForCallerClass();
 
@@ -53,9 +56,21 @@ public final class PublishDnsUpdatesAction implements Runnable, Callable<Void> {
   @Inject DnsWriterProxy dnsWriterProxy;
   @Inject DnsMetrics dnsMetrics;
   @Inject @Config("dnsWriteLockTimeout") Duration timeout;
-  @Inject @Parameter(RequestParameters.PARAM_TLD) String tld;
-  @Inject @Parameter(DOMAINS_PARAM) Set<String> domains;
-  @Inject @Parameter(HOSTS_PARAM) Set<String> hosts;
+
+  /**
+   * The DNS writer to use for this batch.
+   *
+   * <p>This comes from the fanout in {@link ReadDnsQueueAction} which dispatches each batch to be
+   * published by each DNS writer on the TLD. So this field contains the value of one of the DNS
+   * writers configured in {@link Registry#getDnsWriters()}, as of the time the batch was written
+   * out (and not necessarily currently).
+   */
+  // TODO(b/63385597): Make this non-optional DNS once writers migration is complete.
+  @Inject @Parameter(PARAM_DNS_WRITER) Optional<String> dnsWriter;
+
+  @Inject @Parameter(PARAM_DOMAINS) Set<String> domains;
+  @Inject @Parameter(PARAM_HOSTS) Set<String> hosts;
+  @Inject @Parameter(PARAM_TLD) String tld;
   @Inject PublishDnsUpdatesAction() {}
 
   /** Runs the task. */
@@ -80,7 +95,12 @@ public final class PublishDnsUpdatesAction implements Runnable, Callable<Void> {
 
   /** Steps through the domain and host refreshes contained in the parameters and processes them. */
   private void processBatch() {
-    try (DnsWriter writer = dnsWriterProxy.getForTld(tld)) {
+    // TODO(b/63385597): After all old DNS task queue items that did not have a DNS writer on them
+    // are finished being processed, then remove handling for when dnsWriter is absent.
+    try (DnsWriter writer =
+        (dnsWriter.isPresent())
+            ? dnsWriterProxy.getByClassNameForTld(dnsWriter.get(), tld)
+            : dnsWriterProxy.getForTld(tld)) {
       for (String domain : nullToEmpty(domains)) {
         if (!DomainNameUtils.isUnder(
             InternetDomainName.from(domain), InternetDomainName.from(tld))) {

java/google/registry/dns/ReadDnsQueueAction.java

@@ -72,17 +72,17 @@ import org.joda.time.Duration;
 )
 public final class ReadDnsQueueAction implements Runnable {
 
-  public static final String KEEP_TASKS_PARAM = "keepTasks";
+  public static final String PARAM_KEEP_TASKS = "keepTasks";
 
-  private static final String JITTER_SECONDS_PARAM = "jitterSeconds";
+  private static final String PARAM_JITTER_SECONDS = "jitterSeconds";
   private static final Random random = new Random();
   private static final FormattingLogger logger = FormattingLogger.getLoggerForCallerClass();
 
   @Inject @Config("dnsTldUpdateBatchSize") int tldUpdateBatchSize;
   @Inject @Config("dnsWriteLockTimeout") Duration writeLockTimeout;
   @Inject @Named(DNS_PUBLISH_PUSH_QUEUE_NAME) Queue dnsPublishPushQueue;
-  @Inject @Parameter(JITTER_SECONDS_PARAM) Optional<Integer> jitterSeconds;
-  @Inject @Parameter(KEEP_TASKS_PARAM) boolean keepTasks;
+  @Inject @Parameter(PARAM_JITTER_SECONDS) Optional<Integer> jitterSeconds;
+  @Inject @Parameter(PARAM_KEEP_TASKS) boolean keepTasks;
   @Inject DnsQueue dnsQueue;
   @Inject TaskEnqueuer taskEnqueuer;
   @Inject ReadDnsQueueAction() {}
@@ -161,26 +161,31 @@ public final class ReadDnsQueueAction implements Runnable {
     if (!pausedTlds.isEmpty()) {
       logger.infofmt("the dns-pull queue is paused for tlds: %s", pausedTlds);
     }
-    // Loop through the multimap by TLD and generate refresh tasks for the hosts and domains.
+    // Loop through the multimap by TLD and generate refresh tasks for the hosts and domains for
+    // each configured DNS writer.
     for (Map.Entry<String, Collection<RefreshItem>> tldRefreshItemsEntry
         : refreshItemMultimap.asMap().entrySet()) {
-      for (List<RefreshItem> chunk : Iterables.partition(
-          tldRefreshItemsEntry.getValue(), tldUpdateBatchSize)) {
+      String tld = tldRefreshItemsEntry.getKey();
+      for (List<RefreshItem> chunk :
+          Iterables.partition(tldRefreshItemsEntry.getValue(), tldUpdateBatchSize)) {
+        for (String dnsWriter : Registry.get(tld).getDnsWriters()) {
           TaskOptions options = withUrl(PublishDnsUpdatesAction.PATH)
               .countdownMillis(jitterSeconds.isPresent()
                   ? random.nextInt((int) SECONDS.toMillis(jitterSeconds.get()))
                   : 0)
-            .param(RequestParameters.PARAM_TLD, tldRefreshItemsEntry.getKey());
+              .param(RequestParameters.PARAM_TLD, tld)
+              .param(PublishDnsUpdatesAction.PARAM_DNS_WRITER, dnsWriter);
           for (RefreshItem refreshItem : chunk) {
             options.param(
                 (refreshItem.type() == TargetType.HOST)
-                  ? PublishDnsUpdatesAction.HOSTS_PARAM
-                  : PublishDnsUpdatesAction.DOMAINS_PARAM,
+                    ? PublishDnsUpdatesAction.PARAM_HOSTS
+                    : PublishDnsUpdatesAction.PARAM_DOMAINS,
                 refreshItem.name());
           }
           taskEnqueuer.enqueue(dnsPublishPushQueue, options);
         }
       }
+    }
     Set<TaskHandle> tasksToDelete = difference(ImmutableSet.copyOf(tasks), tasksToKeep);
     // In keepTasks mode, never delete any tasks.
     if (keepTasks) {

javatests/google/registry/dns/PublishDnsUpdatesActionTest.java

@@ -23,6 +23,7 @@ import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 
+import com.google.common.base.Optional;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import google.registry.dns.DnsMetrics.Status;
@@ -79,6 +80,7 @@ public class PublishDnsUpdatesActionTest {
     action.tld = tld;
     action.hosts = ImmutableSet.<String>of();
     action.domains = ImmutableSet.<String>of();
+    action.dnsWriter = Optional.<String>absent();
     action.dnsWriterProxy = new DnsWriterProxy(ImmutableMap.of("mock", dnsWriter));
     action.dnsMetrics = dnsMetrics;
     return action;
@@ -102,6 +104,7 @@ public class PublishDnsUpdatesActionTest {
   public void testDomain_published() throws Exception {
     action = createAction("xn--q9jyb4c");
     action.domains = ImmutableSet.of("example.xn--q9jyb4c");
+    action.dnsWriter = Optional.of("mock");
     action.run();
 
     verify(dnsWriter).publishDomain("example.xn--q9jyb4c");

javatests/google/registry/dns/ReadDnsQueueActionTest.java

@@ -24,7 +24,6 @@ import static google.registry.testing.DatastoreHelper.createTlds;
 import static google.registry.testing.DatastoreHelper.persistResource;
 import static google.registry.testing.TaskQueueHelper.assertNoTasksEnqueued;
 import static google.registry.testing.TaskQueueHelper.assertTasksEnqueued;
-import static java.util.Arrays.asList;
 
 import com.google.appengine.api.taskqueue.QueueFactory;
 import com.google.appengine.api.taskqueue.TaskOptions;
@@ -34,6 +33,8 @@ import com.google.common.base.Function;
 import com.google.common.base.Joiner;
 import com.google.common.base.Optional;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.net.InternetDomainName;
 import google.registry.dns.DnsConstants.TargetType;
 import google.registry.model.registry.Registry;
@@ -46,6 +47,7 @@ import google.registry.util.Retrier;
 import google.registry.util.TaskEnqueuer;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map.Entry;
 import org.joda.time.DateTime;
 import org.joda.time.DateTimeZone;
 import org.joda.time.Duration;
@@ -85,7 +87,16 @@ public class ReadDnsQueueActionTest {
   public void before() throws Exception {
     clock.setTo(DateTime.now(DateTimeZone.UTC));
     createTlds("com", "net", "example");
-    persistResource(Registry.get("example").asBuilder().setTldType(TldType.TEST).build());
+    persistResource(
+        Registry.get("com").asBuilder().setDnsWriters(ImmutableSet.of("comWriter")).build());
+    persistResource(
+        Registry.get("net").asBuilder().setDnsWriters(ImmutableSet.of("netWriter")).build());
+    persistResource(
+        Registry.get("example")
+            .asBuilder()
+            .setTldType(TldType.TEST)
+            .setDnsWriters(ImmutableSet.of("exampleWriter"))
+            .build());
     dnsQueue = DnsQueue.create();
   }
 
@@ -112,17 +123,22 @@ public class ReadDnsQueueActionTest {
     return options.param("tld", tld);
   }
 
-  private void assertTldsEnqueuedInPushQueue(String... tlds) throws Exception {
+  private void assertTldsEnqueuedInPushQueue(ImmutableMap<String, String> tldsToDnsWriters)
+      throws Exception {
     assertTasksEnqueued(
         DNS_PUBLISH_PUSH_QUEUE_NAME,
-        transform(asList(tlds), new Function<String, TaskMatcher>() {
+        transform(
+            tldsToDnsWriters.entrySet().asList(),
+            new Function<Entry<String, String>, TaskMatcher>() {
               @Override
-          public TaskMatcher apply(String tld) {
+              public TaskMatcher apply(Entry<String, String> tldToDnsWriter) {
                 return new TaskMatcher()
                     .url(PublishDnsUpdatesAction.PATH)
-                .param("tld", tld)
+                    .param("tld", tldToDnsWriter.getKey())
+                    .param("dnsWriter", tldToDnsWriter.getValue())
                     .header("content-type", "application/x-www-form-urlencoded");
-          }}));
+              }
+            }));
   }
 
   @Test
@@ -146,7 +162,8 @@ public class ReadDnsQueueActionTest {
     dnsQueue.addDomainRefreshTask("domain.example");
     run(false);
     assertNoTasksEnqueued(DNS_PULL_QUEUE_NAME);
-    assertTldsEnqueuedInPushQueue("com", "net", "example");
+    assertTldsEnqueuedInPushQueue(
+        ImmutableMap.of("com", "comWriter", "net", "netWriter", "example", "exampleWriter"));
   }
 
   @Test
@@ -157,7 +174,8 @@ public class ReadDnsQueueActionTest {
     List<TaskStateInfo> preexistingTasks =
         TaskQueueHelper.getQueueInfo(DNS_PULL_QUEUE_NAME).getTaskInfo();
     run(true);
-    assertTldsEnqueuedInPushQueue("com", "net", "example");
+    assertTldsEnqueuedInPushQueue(
+        ImmutableMap.of("com", "comWriter", "net", "netWriter", "example", "exampleWriter"));
     // Check that keepTasks was honored and the pull queue tasks are still present in the queue.
     assertTasksEnqueued(DNS_PULL_QUEUE_NAME, preexistingTasks);
   }
@@ -170,7 +188,8 @@ public class ReadDnsQueueActionTest {
     dnsQueue.addDomainRefreshTask("domain.example");
     run(false);
     assertTasksEnqueued(DNS_PULL_QUEUE_NAME, new TaskMatcher());
-    assertTldsEnqueuedInPushQueue("com", "example");
+    assertTldsEnqueuedInPushQueue(
+        ImmutableMap.of("com", "comWriter", "example", "exampleWriter"));
   }
 
   @Test
@@ -180,13 +199,10 @@ public class ReadDnsQueueActionTest {
     dnsQueue.addZoneRefreshTask("example");
     run(false);
     assertNoTasksEnqueued(DNS_PULL_QUEUE_NAME);
-    assertTasksEnqueued(DNS_PUBLISH_PUSH_QUEUE_NAME,
-        new TaskMatcher()
-            .url(PublishDnsUpdatesAction.PATH)
-            .param("domains", "domain.net"),
-        new TaskMatcher()
-            .url(PublishDnsUpdatesAction.PATH)
-            .param("hosts", "ns1.domain.com"));
+    assertTasksEnqueued(
+        DNS_PUBLISH_PUSH_QUEUE_NAME,
+        new TaskMatcher().url(PublishDnsUpdatesAction.PATH).param("domains", "domain.net"),
+        new TaskMatcher().url(PublishDnsUpdatesAction.PATH).param("hosts", "ns1.domain.com"));
   }
 
   @Test
@@ -234,4 +250,6 @@ public class ReadDnsQueueActionTest {
     assertNoTasksEnqueued(DNS_PULL_QUEUE_NAME);
     assertTasksEnqueued(DNS_PUBLISH_PUSH_QUEUE_NAME, expectedTasks);
   }
+
+  // TODO(b/63385597): Add a test that DNS tasks are processed for multiple DNS writers once allowed
 }