Make DnsWriter truly atomic

Right now - if there's an error during DnsWriter.publish*, all the publish from
before that error will be committed, while all the publish after that error
will not.

More than that - in some writers partial publishes can be committed, depending
on implementation.

This defines a new contract that publish* are only committed when .commit is
called. That way any error will simply mean no publish is committed.

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

java/google/registry/dns/PublishDnsUpdatesAction.java

@@ -93,7 +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() {
-    try (DnsWriter writer = dnsWriterProxy.getByClassNameForTld(dnsWriter, tld)) {
+    DnsWriter writer = dnsWriterProxy.getByClassNameForTld(dnsWriter, tld);
     for (String domain : nullToEmpty(domains)) {
       if (!DomainNameUtils.isUnder(
           InternetDomainName.from(domain), InternetDomainName.from(tld))) {
@@ -116,6 +116,7 @@ public final class PublishDnsUpdatesAction implements Runnable, Callable<Void> {
         logger.infofmt("%s: published host %s", tld, host);
       }
     }
-    }
+    // If we got here it means we managed to stage the entire batch without any errors.
+    writer.commit();
   }
 }

java/google/registry/dns/writer/DnsWriter.java

@@ -17,16 +17,18 @@ package google.registry.dns.writer;
 /**
  * Transaction object for sending an atomic batch of updates for a single zone to the DNS server.
  *
+ * <p>All updates are tentative until commit is called. If commit isn't called, no change will
+ * happen.
+ *
  * <p>Here's an example of how you would publish updates for a domain and host:
  * <pre>
  * &#064;Inject Provider&lt;DnsWriter&gt; dnsWriter;
- * try (DnsWriter writer = dnsWriter.get()) {
  * writer.publishDomain(domainName);
  * writer.publishHost(hostName);
- * }
+ * writer.commit();
  * </pre>
  */
-public interface DnsWriter extends AutoCloseable {
+public interface DnsWriter {
 
   /**
    * Loads {@code domainName} from Datastore and publishes its NS/DS records to the DNS server.
@@ -34,6 +36,9 @@ public interface DnsWriter extends AutoCloseable {
    * and a DS record for each delegation signer stored in the registry for the supplied domain name.
    * If the domain is deleted or is in a "non-publish" state then any existing records are deleted.
    *
+   * This must NOT actually perform any action, instead it should stage the action so that it's
+   * performed when {@link #commit()} is called.
+   *
    * @param domainName the fully qualified domain name, with no trailing dot
    */
   void publishDomain(String domainName);
@@ -46,11 +51,28 @@ public interface DnsWriter extends AutoCloseable {
    * the existing records are deleted. Assumes that this method will only be called for in-bailiwick
    * hosts. The registry does not have addresses for other hosts.
    *
+   * This must NOT actually perform any action, instead it should stage the action so that it's
+   * performed when {@link #commit()} is called.
+   *
    * @param hostName the fully qualified host name, with no trailing dot
    */
   void publishHost(String hostName);
 
-  /** Commits the updates to the DNS server atomically. */
+  /**
-  @Override
+   * Commits the updates to the DNS server atomically.
-  void close();
+   *
+   * <p>The user is responsible for making sure commit() isn't called twice. Implementations are
+   * encouraged to throw an error if commit() is called twice.
+   *
+   * <p>Here's an example of how you would do that
+   * <pre>
+   * private boolean committed = false;
+   * void commit() {
+   *   checkState(!committed, "commit() has already been called");
+   *   committed = true;
+   *   // ... actual commit implementation
+   * }
+   * </pre>
+   */
+  void commit();
 }

java/google/registry/dns/writer/VoidDnsWriter.java

@@ -14,6 +14,8 @@
 
 package google.registry.dns.writer;
 
+import static com.google.common.base.Preconditions.checkState;
+
 import com.google.common.base.Joiner;
 import java.util.HashSet;
 import java.util.Set;
@@ -35,6 +37,8 @@ public final class VoidDnsWriter implements DnsWriter {
 
   private static final Logger logger = Logger.getLogger(VoidDnsWriter.class.getName());
 
+  private boolean committed = false;
+
   private final Set<String> names = new HashSet<>();
 
   @Inject
@@ -51,7 +55,10 @@ public final class VoidDnsWriter implements DnsWriter {
   }
 
   @Override
-  public void close() {
+  public void commit() {
+    checkState(!committed, "commit() has already been called");
+    committed = true;
+
     logger.warning("Ignoring DNS zone updates! No DnsWriterFactory implementation specified!\n"
         + Joiner.on('\n').join(names));
   }

java/google/registry/dns/writer/clouddns/CloudDnsWriter.java

@@ -15,6 +15,7 @@
 package google.registry.dns.writer.clouddns;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkState;
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
 
 import com.google.api.client.googleapis.json.GoogleJsonError.ErrorInfo;
@@ -84,6 +85,8 @@ public class CloudDnsWriter implements DnsWriter {
   private final ImmutableMap.Builder<String, ImmutableSet<ResourceRecordSet>>
       desiredRecordsBuilder = new ImmutableMap.Builder<>();
 
+  private boolean committed = false;
+
   @Inject
   CloudDnsWriter(
       Dns dnsConnection,
@@ -270,12 +273,14 @@ public class CloudDnsWriter implements DnsWriter {
    * representation built via this writer.
    */
   @Override
-  public void close() {
+  public void commit() {
-    close(desiredRecordsBuilder.build());
+    checkState(!committed, "commit() has already been called");
+    committed = true;
+    commit(desiredRecordsBuilder.build());
   }
 
   @VisibleForTesting
-  void close(ImmutableMap<String, ImmutableSet<ResourceRecordSet>> desiredRecords) {
+  void commit(ImmutableMap<String, ImmutableSet<ResourceRecordSet>> desiredRecords) {
     retrier.callWithRetry(getMutateZoneCallback(desiredRecords), ZoneStateException.class);
     logger.info("Wrote to Cloud DNS");
   }

java/google/registry/dns/writer/dnsupdate/DnsUpdateWriter.java

@@ -14,6 +14,7 @@
 
 package google.registry.dns.writer.dnsupdate;
 
+import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Verify.verify;
 import static com.google.common.collect.Sets.intersection;
 import static com.google.common.collect.Sets.union;
@@ -26,6 +27,7 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.net.InternetDomainName;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.dns.writer.DnsWriter;
+import google.registry.dns.writer.DnsWriterZone;
 import google.registry.model.domain.DomainResource;
 import google.registry.model.domain.secdns.DelegationSignerData;
 import google.registry.model.host.HostResource;
@@ -54,9 +56,11 @@ import org.xbill.DNS.Update;
  * A DnsWriter that implements the DNS UPDATE protocol as specified in
  * <a href="https://tools.ietf.org/html/rfc2136">RFC 2136</a>. Publishes changes in the
  * domain-registry to a (capable) external DNS server, sometimes called a "hidden master". DNS
- * UPDATE messages are sent via a supplied "transport" class. For each publish call, a single
+ * UPDATE messages are sent via a supplied "transport" class.
- * UPDATE message is created containing the records required to "synchronize" the DNS with the
+ *
- * current (at the time of processing) state of the registry, for the supplied domain/host.
+ * On call to {@link #commit()}, a single UPDATE message is created containing the records required
+ * to "synchronize" the DNS with the current (at the time of processing) state of the registry, for
+ * the supplied domain/host.
  *
  * <p>The general strategy of the publish methods is to delete <em>all</em> resource records of any
  * <em>type</em> that match the exact domain/host name supplied. And then for create/update cases,
@@ -67,11 +71,10 @@ import org.xbill.DNS.Update;
  * <p>Only NS, DS, A, and AAAA records are published, and in particular no DNSSEC signing is done
  * assuming that this will be done by a third party DNS provider.
  *
- * <p>Each publish call is treated as an atomic update to the DNS. If an update fails an exception
+ * <p>Each commit call is treated as an atomic update to the DNS. If a commit fails an exception
- * is thrown, expecting the caller to retry the update later. The SOA record serial number is
+ * is thrown. The SOA record serial number is implicitly incremented by the server on each UPDATE
- * implicitly incremented by the server on each UPDATE message, as required by RFC 2136. Care must
+ * message, as required by RFC 2136. Care must be taken to make sure the SOA serial number does not
- * be taken to make sure the SOA serial number does not go backwards if the entire TLD (zone) is
+ * go backwards if the entire TLD (zone) is "reset" to empty and republished.
- * "reset" to empty and republished.
  */
 public class DnsUpdateWriter implements DnsWriter {
 
@@ -86,6 +89,10 @@ public class DnsUpdateWriter implements DnsWriter {
   private final Duration dnsDefaultDsTtl;
   private final DnsMessageTransport transport;
   private final Clock clock;
+  private final Update update;
+  private final String zoneName;
+
+  private boolean committed = false;
 
   /**
    * Class constructor.
@@ -96,11 +103,14 @@ public class DnsUpdateWriter implements DnsWriter {
    */
   @Inject
   public DnsUpdateWriter(
+      @DnsWriterZone String zoneName,
       @Config("dnsDefaultATtl") Duration dnsDefaultATtl,
       @Config("dnsDefaultNsTtl") Duration dnsDefaultNsTtl,
       @Config("dnsDefaultDsTtl") Duration dnsDefaultDsTtl,
       DnsMessageTransport transport,
       Clock clock) {
+    this.zoneName = zoneName;
+    this.update = new Update(toAbsoluteName(zoneName));
     this.dnsDefaultATtl = dnsDefaultATtl;
     this.dnsDefaultNsTtl = dnsDefaultNsTtl;
     this.dnsDefaultDsTtl = dnsDefaultDsTtl;
@@ -118,8 +128,6 @@ public class DnsUpdateWriter implements DnsWriter {
    */
   private void publishDomain(String domainName, String requestingHostName) {
     DomainResource domain = loadByForeignKey(DomainResource.class, domainName, clock.nowUtc());
-    try {
-      Update update = new Update(toAbsoluteName(findTldFromName(domainName)));
     update.delete(toAbsoluteName(domainName), Type.ANY);
     if (domain != null) {
       // As long as the domain exists, orphan glues should be cleaned.
@@ -130,15 +138,6 @@ public class DnsUpdateWriter implements DnsWriter {
         update.add(makeDelegationSignerSet(domain));
       }
     }
-      Message response = transport.send(update);
-      verify(
-          response.getRcode() == Rcode.NOERROR,
-          "DNS server failed domain update for '%s' rcode: %s",
-          domainName,
-          Rcode.string(response.getRcode()));
-    } catch (IOException e) {
-      throw new RuntimeException("publishDomain failed: " + domainName, e);
-    }
   }
 
   @Override
@@ -168,13 +167,24 @@ public class DnsUpdateWriter implements DnsWriter {
     publishDomain(domain, hostName);
   }
 
-  /**
-   * Does nothing. Publish calls are synchronous and atomic.
-   */
   @Override
-  public void close() {}
+  public void commit() {
+    checkState(!committed, "commit() has already been called");
+    committed = true;
 
-  private RRset makeDelegationSignerSet(DomainResource domain) throws TextParseException {
+    try {
+      Message response = transport.send(update);
+      verify(
+          response.getRcode() == Rcode.NOERROR,
+          "DNS server failed domain update for '%s' rcode: %s",
+          zoneName,
+          Rcode.string(response.getRcode()));
+    } catch (IOException e) {
+      throw new RuntimeException("publishDomain failed for zone: " + zoneName, e);
+    }
+  }
+
+  private RRset makeDelegationSignerSet(DomainResource domain) {
     RRset signerSet = new RRset();
     for (DelegationSignerData signerData : domain.getDsData()) {
       DSRecord dsRecord =
@@ -192,7 +202,7 @@ public class DnsUpdateWriter implements DnsWriter {
   }
 
   private void deleteSubordinateHostAddressSet(
-      DomainResource domain, String additionalHost, Update update) throws TextParseException {
+      DomainResource domain, String additionalHost, Update update) {
     for (String hostName :
         union(
             domain.getSubordinateHosts(),
@@ -203,8 +213,7 @@ public class DnsUpdateWriter implements DnsWriter {
     }
   }
 
-  private void addInBailiwickNameServerSet(DomainResource domain, Update update)
+  private void addInBailiwickNameServerSet(DomainResource domain, Update update) {
-      throws TextParseException {
     for (String hostName :
         intersection(
             domain.loadNameserverFullyQualifiedHostNames(), domain.getSubordinateHosts())) {
@@ -214,7 +223,7 @@ public class DnsUpdateWriter implements DnsWriter {
     }
   }
 
-  private RRset makeNameServerSet(DomainResource domain) throws TextParseException {
+  private RRset makeNameServerSet(DomainResource domain) {
     RRset nameServerSet = new RRset();
     for (String hostName : domain.loadNameserverFullyQualifiedHostNames()) {
       NSRecord record =
@@ -228,7 +237,7 @@ public class DnsUpdateWriter implements DnsWriter {
     return nameServerSet;
   }
 
-  private RRset makeAddressSet(HostResource host) throws TextParseException {
+  private RRset makeAddressSet(HostResource host) {
     RRset addressSet = new RRset();
     for (InetAddress address : host.getInetAddresses()) {
       if (address instanceof Inet4Address) {
@@ -244,7 +253,7 @@ public class DnsUpdateWriter implements DnsWriter {
     return addressSet;
   }
 
-  private RRset makeV6AddressSet(HostResource host) throws TextParseException {
+  private RRset makeV6AddressSet(HostResource host) {
     RRset addressSet = new RRset();
     for (InetAddress address : host.getInetAddresses()) {
       if (address instanceof Inet6Address) {
@@ -260,11 +269,12 @@ public class DnsUpdateWriter implements DnsWriter {
     return addressSet;
   }
 
-  private String findTldFromName(String name) {
+  private Name toAbsoluteName(String name) {
-    return Registries.findTldForNameOrThrow(InternetDomainName.from(name)).toString();
+    try {
-  }
-
-  private Name toAbsoluteName(String name) throws TextParseException {
       return Name.fromString(name, Name.root);
+    } catch (TextParseException e) {
+      throw new RuntimeException(
+          String.format("toAbsoluteName failed for name: %s in zone: %s", name, zoneName), e);
+    }
   }
 }

javatests/google/registry/dns/PublishDnsUpdatesActionTest.java

@@ -92,7 +92,7 @@ public class PublishDnsUpdatesActionTest {
     action.run();
 
     verify(dnsWriter).publishHost("ns1.example.xn--q9jyb4c");
-    verify(dnsWriter).close();
+    verify(dnsWriter).commit();
     verifyNoMoreInteractions(dnsWriter);
 
     verify(dnsMetrics).incrementPublishHostRequests("xn--q9jyb4c", Status.ACCEPTED);
@@ -106,7 +106,7 @@ public class PublishDnsUpdatesActionTest {
     action.run();
 
     verify(dnsWriter).publishDomain("example.xn--q9jyb4c");
-    verify(dnsWriter).close();
+    verify(dnsWriter).commit();
     verifyNoMoreInteractions(dnsWriter);
 
     verify(dnsMetrics).incrementPublishDomainRequests("xn--q9jyb4c", Status.ACCEPTED);
@@ -126,7 +126,7 @@ public class PublishDnsUpdatesActionTest {
     verify(dnsWriter).publishHost("ns1.example.xn--q9jyb4c");
     verify(dnsWriter).publishHost("ns2.example.xn--q9jyb4c");
     verify(dnsWriter).publishHost("ns1.example2.xn--q9jyb4c");
-    verify(dnsWriter).close();
+    verify(dnsWriter).commit();
     verifyNoMoreInteractions(dnsWriter);
 
     verify(dnsMetrics, times(2)).incrementPublishDomainRequests("xn--q9jyb4c", Status.ACCEPTED);
@@ -141,7 +141,7 @@ public class PublishDnsUpdatesActionTest {
     action.hosts = ImmutableSet.of("ns1.example.com", "ns2.example.com", "ns1.example2.com");
     action.run();
 
-    verify(dnsWriter).close();
+    verify(dnsWriter).commit();
     verifyNoMoreInteractions(dnsWriter);
 
     verify(dnsMetrics, times(2)).incrementPublishDomainRequests("xn--q9jyb4c", Status.REJECTED);

javatests/google/registry/dns/writer/clouddns/CloudDnsWriterTest.java

@@ -174,7 +174,7 @@ public class CloudDnsWriterTest {
 
   private void verifyZone(ImmutableSet<ResourceRecordSet> expectedRecords) throws Exception {
     // Trigger zone changes
-    writer.close();
+    writer.commit();
 
     assertThat(stubZone).containsExactlyElementsIn(expectedRecords);
   }
@@ -416,12 +416,12 @@ public class CloudDnsWriterTest {
   @Test
   @SuppressWarnings("unchecked")
   public void retryMutateZoneOnError() throws Exception {
-    try (CloudDnsWriter spyWriter = spy(writer)) {
+    CloudDnsWriter spyWriter = spy(writer);
     when(mutateZoneCallable.call()).thenThrow(ZoneStateException.class).thenReturn(null);
     when(spyWriter.getMutateZoneCallback(
         Matchers.<ImmutableMap<String, ImmutableSet<ResourceRecordSet>>>any()))
         .thenReturn(mutateZoneCallable);
-    }
+    spyWriter.commit();
 
     verify(mutateZoneCallable, times(2)).call();
   }

javatests/google/registry/dns/writer/dnsupdate/DnsUpdateWriterTest.java

@@ -28,6 +28,7 @@ import static google.registry.testing.DatastoreHelper.persistDeletedHost;
 import static google.registry.testing.DatastoreHelper.persistResource;
 import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyZeroInteractions;
 import static org.mockito.Mockito.when;
 
 import com.google.common.base.VerifyException;
@@ -97,7 +98,8 @@ public class DnsUpdateWriterTest {
     createTld("tld");
     when(mockResolver.send(any(Update.class))).thenReturn(messageWithResponseCode(Rcode.NOERROR));
 
-    writer = new DnsUpdateWriter(Duration.ZERO, Duration.ZERO, Duration.ZERO, mockResolver, clock);
+    writer = new DnsUpdateWriter(
+        "tld", Duration.ZERO, Duration.ZERO, Duration.ZERO, mockResolver, clock);
   }
 
   @Test
@@ -112,6 +114,7 @@ public class DnsUpdateWriterTest {
     persistResource(domain);
 
     writer.publishDomain("example.tld");
+    writer.commit();
 
     verify(mockResolver).send(updateCaptor.capture());
     Update update = updateCaptor.getValue();
@@ -121,6 +124,62 @@ public class DnsUpdateWriterTest {
     assertThatTotalUpdateSetsIs(update, 2); // The delete and NS sets
   }
 
+  @Test
+  public void testPublishAtomic_noCommit() throws Exception {
+    HostResource host1 = persistActiveHost("ns.example1.tld");
+    DomainResource domain1 =
+        persistActiveDomain("example1.tld")
+            .asBuilder()
+            .setNameservers(ImmutableSet.of(Key.create(host1)))
+            .build();
+    persistResource(domain1);
+
+    HostResource host2 = persistActiveHost("ns.example2.tld");
+    DomainResource domain2 =
+        persistActiveDomain("example2.tld")
+            .asBuilder()
+            .setNameservers(ImmutableSet.of(Key.create(host2)))
+            .build();
+    persistResource(domain2);
+
+    writer.publishDomain("example1.tld");
+    writer.publishDomain("example2.tld");
+
+    verifyZeroInteractions(mockResolver);
+  }
+
+  @Test
+  public void testPublishAtomic_oneUpdate() throws Exception {
+    HostResource host1 = persistActiveHost("ns.example1.tld");
+    DomainResource domain1 =
+        persistActiveDomain("example1.tld")
+            .asBuilder()
+            .setNameservers(ImmutableSet.of(Key.create(host1)))
+            .build();
+    persistResource(domain1);
+
+    HostResource host2 = persistActiveHost("ns.example2.tld");
+    DomainResource domain2 =
+        persistActiveDomain("example2.tld")
+            .asBuilder()
+            .setNameservers(ImmutableSet.of(Key.create(host2)))
+            .build();
+    persistResource(domain2);
+
+    writer.publishDomain("example1.tld");
+    writer.publishDomain("example2.tld");
+    writer.commit();
+
+    verify(mockResolver).send(updateCaptor.capture());
+    Update update = updateCaptor.getValue();
+    assertThatUpdatedZoneIs(update, "tld.");
+    assertThatUpdateDeletes(update, "example1.tld.", Type.ANY);
+    assertThatUpdateDeletes(update, "example2.tld.", Type.ANY);
+    assertThatUpdateAdds(update, "example1.tld.", Type.NS, "ns.example1.tld.");
+    assertThatUpdateAdds(update, "example2.tld.", Type.NS, "ns.example2.tld.");
+    assertThatTotalUpdateSetsIs(update, 4); // The delete and NS sets for each TLD
+  }
+
   @Test
   public void testPublishDomainCreate_publishesDelegationSigner() throws Exception {
     DomainResource domain =
@@ -134,6 +193,7 @@ public class DnsUpdateWriterTest {
     persistResource(domain);
 
     writer.publishDomain("example.tld");
+    writer.commit();
 
     verify(mockResolver).send(updateCaptor.capture());
     Update update = updateCaptor.getValue();
@@ -155,6 +215,7 @@ public class DnsUpdateWriterTest {
     persistResource(domain);
 
     writer.publishDomain("example.tld");
+    writer.commit();
 
     verify(mockResolver).send(updateCaptor.capture());
     Update update = updateCaptor.getValue();
@@ -168,6 +229,7 @@ public class DnsUpdateWriterTest {
     persistDeletedDomain("example.tld", clock.nowUtc().minusDays(1));
 
     writer.publishDomain("example.tld");
+    writer.commit();
 
     verify(mockResolver).send(updateCaptor.capture());
     Update update = updateCaptor.getValue();
@@ -196,6 +258,7 @@ public class DnsUpdateWriterTest {
             .build());
 
     writer.publishHost("ns1.example.tld");
+    writer.commit();
 
     verify(mockResolver).send(updateCaptor.capture());
     Update update = updateCaptor.getValue();
@@ -214,6 +277,7 @@ public class DnsUpdateWriterTest {
     persistActiveDomain("example.tld");
 
     writer.publishHost("ns1.example.tld");
+    writer.commit();
 
     verify(mockResolver).send(updateCaptor.capture());
     Update update = updateCaptor.getValue();
@@ -233,6 +297,7 @@ public class DnsUpdateWriterTest {
             .build());
 
     writer.publishHost("ns1.example.tld");
+    writer.commit();
 
     verify(mockResolver).send(updateCaptor.capture());
     Update update = updateCaptor.getValue();
@@ -266,6 +331,7 @@ public class DnsUpdateWriterTest {
             .build());
 
     writer.publishDomain("example.tld");
+    writer.commit();
 
     verify(mockResolver).send(updateCaptor.capture());
     Update update = updateCaptor.getValue();
@@ -300,6 +366,7 @@ public class DnsUpdateWriterTest {
             .build());
 
     writer.publishDomain("example.tld");
+    writer.commit();
 
     verify(mockResolver).send(updateCaptor.capture());
     Update update = updateCaptor.getValue();
@@ -325,6 +392,7 @@ public class DnsUpdateWriterTest {
     thrown.expect(VerifyException.class, "SERVFAIL");
 
     writer.publishDomain("example.tld");
+    writer.commit();
   }
 
   @Test
@@ -339,6 +407,7 @@ public class DnsUpdateWriterTest {
     thrown.expect(VerifyException.class, "SERVFAIL");
 
     writer.publishHost("ns1.example.tld");
+    writer.commit();
   }
 
   private void assertThatUpdatedZoneIs(Update update, String zoneName) {