Allow multiple DNS writers on TLDs

This completes the data/functionality migration for multiple DNS writers.

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

java/google/registry/dns/DnsModule.java

@@ -22,13 +22,11 @@ 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;
@@ -62,9 +60,8 @@ public abstract class DnsModule {
 
   @Provides
   @Parameter(PARAM_DNS_WRITER)
-  static Optional<String> provideDnsWriter(HttpServletRequest req) {
+  static String provideDnsWriter(HttpServletRequest req) {
-    // TODO(b/63385597): Make this required after DNS writers migration is complete.
+    return extractRequiredParameter(req, PARAM_DNS_WRITER);
-    return extractOptionalParameter(req, PARAM_DNS_WRITER);
   }
 
   @Provides

java/google/registry/dns/DnsWriterProxy.java

@@ -15,7 +15,6 @@
 package google.registry.dns;
 
 import static com.google.common.base.Preconditions.checkState;
-import static com.google.common.collect.Iterables.getOnlyElement;
 import static google.registry.util.FormattingLogger.getLoggerForCallerClass;
 
 import com.google.common.collect.ImmutableMap;
@@ -37,13 +36,6 @@ public final class DnsWriterProxy {
     this.dnsWriters = ImmutableMap.copyOf(dnsWriters);
   }
 
-  /** 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) {
-    return getByClassNameForTld(getOnlyElement(Registry.get(tld).getDnsWriters()), tld);
-  }
-
   /** Returns the instance of {@link DnsWriter} by class name. */
   public DnsWriter getByClassNameForTld(String className, String tld) {
     if (!Registry.get(tld).getDnsWriters().contains(className)) {

java/google/registry/dns/PublishDnsUpdatesAction.java

@@ -19,7 +19,6 @@ 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;
@@ -65,8 +64,7 @@ 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).
    */
-  // TODO(b/63385597): Make this non-optional DNS once writers migration is complete.
+  @Inject @Parameter(PARAM_DNS_WRITER) String dnsWriter;
-  @Inject @Parameter(PARAM_DNS_WRITER) Optional<String> dnsWriter;
 
   @Inject @Parameter(PARAM_DOMAINS) Set<String> domains;
   @Inject @Parameter(PARAM_HOSTS) Set<String> hosts;
@@ -95,12 +93,7 @@ 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() {
-    // TODO(b/63385597): After all old DNS task queue items that did not have a DNS writer on them
+    try (DnsWriter writer = dnsWriterProxy.getByClassNameForTld(dnsWriter, tld)) {
-    // 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/model/registry/Registry.java

@@ -676,8 +676,6 @@ public class Registry extends ImmutableObject implements Buildable {
     }
 
     public Builder setDnsWriters(ImmutableSet<String> dnsWriters) {
-      // TODO(b/63385597): Remove this restriction once DNS task queue migration is complete.
-      checkArgument(dnsWriters.size() == 1, "Multiple DNS writers are not yet supported");
       getInstance().dnsWriters = dnsWriters;
       return this;
     }

javatests/google/registry/dns/PublishDnsUpdatesActionTest.java

@@ -23,7 +23,6 @@ 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;
@@ -80,7 +79,7 @@ public class PublishDnsUpdatesActionTest {
     action.tld = tld;
     action.hosts = ImmutableSet.<String>of();
     action.domains = ImmutableSet.<String>of();
-    action.dnsWriter = Optional.<String>absent();
+    action.dnsWriter = "mock";
     action.dnsWriterProxy = new DnsWriterProxy(ImmutableMap.of("mock", dnsWriter));
     action.dnsMetrics = dnsMetrics;
     return action;
@@ -104,7 +103,6 @@ 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

@@ -33,7 +33,7 @@ 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.ImmutableMultimap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.net.InternetDomainName;
 import google.registry.dns.DnsConstants.TargetType;
@@ -123,12 +123,12 @@ public class ReadDnsQueueActionTest {
     return options.param("tld", tld);
   }
 
-  private void assertTldsEnqueuedInPushQueue(ImmutableMap<String, String> tldsToDnsWriters)
+  private void assertTldsEnqueuedInPushQueue(ImmutableMultimap<String, String> tldsToDnsWriters)
       throws Exception {
     assertTasksEnqueued(
         DNS_PUBLISH_PUSH_QUEUE_NAME,
         transform(
-            tldsToDnsWriters.entrySet().asList(),
+            tldsToDnsWriters.entries().asList(),
             new Function<Entry<String, String>, TaskMatcher>() {
               @Override
               public TaskMatcher apply(Entry<String, String> tldToDnsWriter) {
@@ -163,7 +163,20 @@ public class ReadDnsQueueActionTest {
     run(false);
     assertNoTasksEnqueued(DNS_PULL_QUEUE_NAME);
     assertTldsEnqueuedInPushQueue(
-        ImmutableMap.of("com", "comWriter", "net", "netWriter", "example", "exampleWriter"));
+        ImmutableMultimap.of("com", "comWriter", "net", "netWriter", "example", "exampleWriter"));
+  }
+
+  @Test
+  public void testSuccess_twoDnsWriters() throws Exception {
+    persistResource(
+        Registry.get("com")
+            .asBuilder()
+            .setDnsWriters(ImmutableSet.of("comWriter", "otherWriter"))
+            .build());
+    dnsQueue.addDomainRefreshTask("domain.com");
+    run(false);
+    assertNoTasksEnqueued(DNS_PULL_QUEUE_NAME);
+    assertTldsEnqueuedInPushQueue(ImmutableMultimap.of("com", "comWriter", "com", "otherWriter"));
   }
 
   @Test
@@ -175,7 +188,7 @@ public class ReadDnsQueueActionTest {
         TaskQueueHelper.getQueueInfo(DNS_PULL_QUEUE_NAME).getTaskInfo();
     run(true);
     assertTldsEnqueuedInPushQueue(
-        ImmutableMap.of("com", "comWriter", "net", "netWriter", "example", "exampleWriter"));
+        ImmutableMultimap.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);
   }
@@ -189,7 +202,7 @@ public class ReadDnsQueueActionTest {
     run(false);
     assertTasksEnqueued(DNS_PULL_QUEUE_NAME, new TaskMatcher());
     assertTldsEnqueuedInPushQueue(
-        ImmutableMap.of("com", "comWriter", "example", "exampleWriter"));
+        ImmutableMultimap.of("com", "comWriter", "example", "exampleWriter"));
   }
 
   @Test
@@ -250,6 +263,4 @@ 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
 }

javatests/google/registry/model/registry/RegistryTest.java

@@ -32,6 +32,7 @@ import static org.joda.money.CurrencyUnit.USD;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.googlecode.objectify.Key;
+import google.registry.dns.writer.VoidDnsWriter;
 import google.registry.model.EntityTestCase;
 import google.registry.model.registry.Registry.RegistryNotFoundException;
 import google.registry.model.registry.Registry.TldState;
@@ -229,6 +230,14 @@ public class RegistryTest extends EntityTestCase {
     assertThat(registry.getLordnUsername()).isEqualTo("username");
   }
 
+  @Test
+  public void testSettingDnsWriters() {
+    Registry registry = Registry.get("tld");
+    assertThat(registry.getDnsWriters()).containsExactly(VoidDnsWriter.NAME);
+    registry = registry.asBuilder().setDnsWriters(ImmutableSet.of("baz", "bang")).build();
+    assertThat(registry.getDnsWriters()).containsExactly("baz", "bang");
+  }
+
   @Test
   public void testPdtLooksLikeGa() {
     Registry registry = Registry.get("tld").asBuilder()

javatests/google/registry/tools/UpdateTldCommandTest.java

@@ -183,9 +183,12 @@ public class UpdateTldCommandTest extends CommandTestCase<UpdateTldCommand> {
   }
 
   @Test
-  public void testFailure_multipleDnsWritersNotYetSupported() throws Exception {
+  public void testSuccess_multipleDnsWriters() throws Exception {
-    thrown.expect(IllegalArgumentException.class, "Multiple DNS writers are not yet supported");
+    assertThat(Registry.get("xn--q9jyb4c").getDnsWriters()).containsExactly("VoidDnsWriter");
+
     runCommandForced("--dns_writers=FooDnsWriter,VoidDnsWriter", "xn--q9jyb4c");
+    assertThat(Registry.get("xn--q9jyb4c").getDnsWriters())
+        .containsExactly("FooDnsWriter", "VoidDnsWriter");
   }
 
   @Test