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,29 +93,30 @@ 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)) {
-      for (String domain : nullToEmpty(domains)) {
-        if (!DomainNameUtils.isUnder(
-            InternetDomainName.from(domain), InternetDomainName.from(tld))) {
-          dnsMetrics.incrementPublishDomainRequests(tld, Status.REJECTED);
-          logger.severefmt("%s: skipping domain %s not under tld", tld, domain);
-        } else {
-          dnsMetrics.incrementPublishDomainRequests(tld, Status.ACCEPTED);
-          writer.publishDomain(domain);
-          logger.infofmt("%s: published domain %s", tld, domain);
-        }
-      }
-      for (String host : nullToEmpty(hosts)) {
-        if (!DomainNameUtils.isUnder(
-            InternetDomainName.from(host), InternetDomainName.from(tld))) {
-          dnsMetrics.incrementPublishHostRequests(tld, Status.REJECTED);
-          logger.severefmt("%s: skipping host %s not under tld", tld, host);
-        } else {
-          dnsMetrics.incrementPublishHostRequests(tld, Status.ACCEPTED);
-          writer.publishHost(host);
-          logger.infofmt("%s: published host %s", tld, host);
-        }
+    DnsWriter writer = dnsWriterProxy.getByClassNameForTld(dnsWriter, tld);
+    for (String domain : nullToEmpty(domains)) {
+      if (!DomainNameUtils.isUnder(
+          InternetDomainName.from(domain), InternetDomainName.from(tld))) {
+        dnsMetrics.incrementPublishDomainRequests(tld, Status.REJECTED);
+        logger.severefmt("%s: skipping domain %s not under tld", tld, domain);
+      } else {
+        dnsMetrics.incrementPublishDomainRequests(tld, Status.ACCEPTED);
+        writer.publishDomain(domain);
+        logger.infofmt("%s: published domain %s", tld, domain);
       }
     }
+    for (String host : nullToEmpty(hosts)) {
+      if (!DomainNameUtils.isUnder(
+          InternetDomainName.from(host), InternetDomainName.from(tld))) {
+        dnsMetrics.incrementPublishHostRequests(tld, Status.REJECTED);
+        logger.severefmt("%s: skipping host %s not under tld", tld, host);
+      } else {
+        dnsMetrics.incrementPublishHostRequests(tld, Status.ACCEPTED);
+        writer.publishHost(host);
+        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.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
-  void close();
+  /**
+   * Commits the updates to the DNS server atomically.
+   *
+   * <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() {
-    close(desiredRecordsBuilder.build());
+  public void commit() {
+    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 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.
+ * UPDATE messages are sent via a supplied "transport" class.
+ *
+ * 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
- * is thrown, expecting the caller to retry the update later. The SOA record serial number is
- * implicitly incremented by the server on each UPDATE message, as required by RFC 2136. Care must
- * be taken to make sure the SOA serial number does not go backwards if the entire TLD (zone) is
- * "reset" to empty and republished.
+ * <p>Each commit call is treated as an atomic update to the DNS. If a commit fails an exception
+ * is thrown. The SOA record serial number is implicitly incremented by the server on each UPDATE
+ * message, as required by RFC 2136. Care must be taken to make sure the SOA serial number does not
+ * go backwards if the entire TLD (zone) is "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,26 +128,15 @@ 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.
-        deleteSubordinateHostAddressSet(domain, requestingHostName, update);
-        if (domain.shouldPublishToDns()) {
-          addInBailiwickNameServerSet(domain, update);
-          update.add(makeNameServerSet(domain));
-          update.add(makeDelegationSignerSet(domain));
-        }
+    update.delete(toAbsoluteName(domainName), Type.ANY);
+    if (domain != null) {
+      // As long as the domain exists, orphan glues should be cleaned.
+      deleteSubordinateHostAddressSet(domain, requestingHostName, update);
+      if (domain.shouldPublishToDns()) {
+        addInBailiwickNameServerSet(domain, update);
+        update.add(makeNameServerSet(domain));
+        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);
     }
   }
 
@@ -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)
-      throws TextParseException {
+  private void addInBailiwickNameServerSet(DomainResource domain, Update update) {
     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) {
-    return Registries.findTldForNameOrThrow(InternetDomainName.from(name)).toString();
-  }
-
-  private Name toAbsoluteName(String name) throws TextParseException {
-    return Name.fromString(name, Name.root);
+  private Name toAbsoluteName(String name) {
+    try {
+      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);
+    }
   }
 }