Add sharded DNS publishing capability

This enables sharded DNS publishing on a per-TLD basis. Instead of a TLD-wide lock, the sharded scheme locks each update on the shard number, allowing parallel writes to DNS.

We allow N (the number of shards) to be 0 or 1 for no sharding, and N > 1 for an N-way sharding scheme. Unless explicitly set, all TLDs default to a numShards of 0, so we don't have to reload all registry objects explicitly.

WARNING: This will change the lock name upon deployment for the PublishDnsAction from "<TLD> Dns Updates" to "<TLD> Dns Updates shard 0". This may cause concurrency issues if the underlying DNSWriter is not parallel-write tolerant (currently all production usages are ZonemanWriter, which is parallel-tolerant, so no issues are expected).

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

java/google/registry/dns/DnsMetrics.java

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

java/google/registry/dns/DnsModule.java

@@ -19,15 +19,20 @@ import static google.registry.dns.DnsConstants.DNS_PULL_QUEUE_NAME;
 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.PublishDnsUpdatesAction.PARAM_LOCK_INDEX;
+import static google.registry.dns.PublishDnsUpdatesAction.PARAM_NUM_PUBLISH_LOCKS;
 import static google.registry.dns.PublishDnsUpdatesAction.PARAM_PUBLISH_TASK_ENQUEUED;
 import static google.registry.dns.PublishDnsUpdatesAction.PARAM_REFRESH_REQUEST_CREATED;
 import static google.registry.request.RequestParameters.extractEnumParameter;
+import static google.registry.request.RequestParameters.extractOptionalIntParameter;
 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.hash.HashFunction;
+import com.google.common.hash.Hashing;
 import dagger.Binds;
 import dagger.Module;
 import dagger.Provides;
@@ -49,6 +54,17 @@ public abstract class DnsModule {
   @DnsWriterZone
   abstract String provideZoneName(@Parameter(RequestParameters.PARAM_TLD) String tld);
 
+  /**
+   * Provides a HashFunction used for generating sharded DNS publish queue tasks.
+   *
+   * <p>We use murmur3_32 because it isn't subject to change, and is fast (non-cryptographic, which
+   * would be overkill in this situation.)
+   */
+  @Provides
+  static HashFunction provideHashFunction() {
+    return Hashing.murmur3_32();
+  }
+
   @Provides
   @Named(DNS_PULL_QUEUE_NAME)
   static Queue provideDnsPullQueue() {
@@ -81,6 +97,20 @@ public abstract class DnsModule {
     return extractRequiredParameter(req, PARAM_DNS_WRITER);
   }
 
+  @Provides
+  @Parameter(PARAM_LOCK_INDEX)
+  static int provideLockIndex(HttpServletRequest req) {
+    // TODO(b/72150053): Make non-optional once this cl has reached production for an hour
+    return extractOptionalIntParameter(req, PARAM_LOCK_INDEX).orElse(1);
+  }
+
+  @Provides
+  @Parameter(PARAM_NUM_PUBLISH_LOCKS)
+  static int provideMaxNumLocks(HttpServletRequest req) {
+    // TODO(b/72150053): Make non-optional once this cl has reached production for an hour
+    return extractOptionalIntParameter(req, PARAM_NUM_PUBLISH_LOCKS).orElse(1);
+  }
+
   @Provides
   @Parameter(PARAM_DOMAINS)
   static Set<String> provideDomains(HttpServletRequest req) {

java/google/registry/dns/PublishDnsUpdatesAction.java

@@ -51,6 +51,8 @@ public final class PublishDnsUpdatesAction implements Runnable, Callable<Void> {
 
   public static final String PATH = "/_dr/task/publishDnsUpdates";
   public static final String PARAM_DNS_WRITER = "dnsWriter";
+  public static final String PARAM_LOCK_INDEX = "lockIndex";
+  public static final String PARAM_NUM_PUBLISH_LOCKS = "numPublishLocks";
   public static final String PARAM_DOMAINS = "domains";
   public static final String PARAM_HOSTS = "hosts";
   public static final String PARAM_PUBLISH_TASK_ENQUEUED = "enqueued";
@@ -78,6 +80,8 @@ public final class PublishDnsUpdatesAction implements Runnable, Callable<Void> {
   // TODO(b/73343464): make not-optional once transition has ended
   @Inject @Parameter(PARAM_REFRESH_REQUEST_CREATED) Optional<DateTime> itemsCreateTime;
 
+  @Inject @Parameter(PARAM_LOCK_INDEX) int lockIndex;
+  @Inject @Parameter(PARAM_NUM_PUBLISH_LOCKS) int numPublishLocks;
   @Inject @Parameter(PARAM_DOMAINS) Set<String> domains;
   @Inject @Parameter(PARAM_HOSTS) Set<String> hosts;
   @Inject @Parameter(PARAM_TLD) String tld;
@@ -108,11 +112,20 @@ public final class PublishDnsUpdatesAction implements Runnable, Callable<Void> {
   /** Runs the task. */
   @Override
   public void run() {
+    if (!validLockParams()) {
+      recordActionResult(ActionStatus.BAD_LOCK_INDEX);
+      requeueBatch();
+      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
     // in the update task being dequeued and dropped. A message will already have been logged
     // to indicate the problem.
-    if (!lockHandler.executeWithLocks(this, tld, timeout, LOCK_NAME)) {
+    if (!lockHandler.executeWithLocks(
+        this,
+        tld,
+        timeout,
+        String.format("%s-lock %d of %d", LOCK_NAME, lockIndex, numPublishLocks))) {
       recordActionResult(ActionStatus.LOCK_FAILURE);
       throw new ServiceUnavailableException("Lock failure");
     }
@@ -127,6 +140,7 @@ public final class PublishDnsUpdatesAction implements Runnable, Callable<Void> {
 
   /** Adds all the domains and hosts in the batch back to the queue to be processed later. */
   private void requeueBatch() {
+    logger.infofmt("Requeueing batch for retry");
     for (String domain : nullToEmpty(domains)) {
       dnsQueue.addDomainRefreshTask(domain);
     }
@@ -135,6 +149,25 @@ public final class PublishDnsUpdatesAction implements Runnable, Callable<Void> {
     }
   }
 
+  /** Returns if the lock parameters are valid for this action. */
+  private boolean validLockParams() {
+    // LockIndex should always be within [1, numPublishLocks]
+    if (lockIndex > numPublishLocks || lockIndex <= 0) {
+      logger.severefmt(
+          "Lock index should be within [1,%d], got %d instead", numPublishLocks, lockIndex);
+      return false;
+    }
+    // Check if the Registry object's num locks has changed since this task was batched
+    int registryNumPublishLocks = Registry.get(tld).getNumDnsPublishLocks();
+    if (registryNumPublishLocks != numPublishLocks) {
+      logger.warningfmt(
+          "Registry numDnsPublishLocks %d out of sync with parameter %d",
+          registryNumPublishLocks, numPublishLocks);
+      return false;
+    }
+    return true;
+  }
+
   /** Steps through the domain and host refreshes contained in the parameters and processes them. */
   private void processBatch() {
     DateTime timeAtStart = clock.nowUtc();
@@ -142,9 +175,7 @@ public final class PublishDnsUpdatesAction implements Runnable, Callable<Void> {
     DnsWriter writer = dnsWriterProxy.getByClassNameForTld(dnsWriter, tld);
 
     if (writer == null) {
-      logger.warningfmt(
+      logger.warningfmt("Couldn't get writer %s for TLD %s", dnsWriter, tld);
-          "Couldn't get writer %s for TLD %s, pushing domains back to the queue for retry",
-          dnsWriter, tld);
       recordActionResult(ActionStatus.BAD_WRITER);
       requeueBatch();
       return;

java/google/registry/dns/ReadDnsQueueAction.java

@@ -15,11 +15,14 @@
 package google.registry.dns;
 
 import static com.google.appengine.api.taskqueue.TaskOptions.Builder.withUrl;
+import static com.google.common.collect.ImmutableSetMultimap.toImmutableSetMultimap;
 import static com.google.common.collect.Sets.difference;
 import static google.registry.dns.DnsConstants.DNS_PUBLISH_PUSH_QUEUE_NAME;
 import static google.registry.dns.DnsConstants.DNS_TARGET_CREATE_TIME_PARAM;
 import static google.registry.dns.DnsConstants.DNS_TARGET_NAME_PARAM;
 import static google.registry.dns.DnsConstants.DNS_TARGET_TYPE_PARAM;
+import static google.registry.util.DomainNameUtils.getSecondLevelDomain;
+import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.concurrent.TimeUnit.SECONDS;
 
 import com.google.appengine.api.taskqueue.Queue;
@@ -32,6 +35,8 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSetMultimap;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Ordering;
+import com.google.common.hash.HashFunction;
+import com.google.common.hash.Hashing;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.dns.DnsConstants.TargetType;
 import google.registry.model.registry.Registries;
@@ -95,6 +100,7 @@ public final class ReadDnsQueueAction implements Runnable {
   @Inject @Parameter(PARAM_JITTER_SECONDS) Optional<Integer> jitterSeconds;
   @Inject Clock clock;
   @Inject DnsQueue dnsQueue;
+  @Inject HashFunction hashFunction;
   @Inject TaskEnqueuer taskEnqueuer;
   @Inject ReadDnsQueueAction() {}
 
@@ -220,7 +226,7 @@ public final class ReadDnsQueueAction implements Runnable {
       logger.warningfmt(
           "The dns-pull queue has unknown TLDs: %s.", classifiedTasks.unknownTlds());
     }
-    enqueueUpdates(classifiedTasks.refreshItemsByTld());
+    bucketRefreshItems(classifiedTasks.refreshItemsByTld());
     if (!classifiedTasks.tasksToKeep().isEmpty()) {
       logger.warningfmt(
           "Keeping %d DNS update tasks in the queue.", classifiedTasks.tasksToKeep().size());
@@ -297,28 +303,74 @@ public final class ReadDnsQueueAction implements Runnable {
     return classifiedTasksBuilder.build();
   }
 
-  private void enqueueUpdates(ImmutableSetMultimap<String, RefreshItem> refreshItemsByTld) {
+  /**
+   * Subdivides the tld to {@link RefreshItem} multimap into buckets by lock index, if applicable.
+   *
+   * <p>If the tld has numDnsPublishLocks <= 1, we enqueue all updates on the default lock 1 of 1.
+   */
+  private void bucketRefreshItems(ImmutableSetMultimap<String, RefreshItem> refreshItemsByTld) {
     // 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
         : refreshItemsByTld.asMap().entrySet()) {
       String tld = tldRefreshItemsEntry.getKey();
-      for (List<RefreshItem> chunk :
+      int numPublishLocks = Registry.get(tld).getNumDnsPublishLocks();
-          Iterables.partition(tldRefreshItemsEntry.getValue(), tldUpdateBatchSize)) {
+      // 1 lock or less implies no TLD-wide locks, simply enqueue everything under lock 1 of 1
+      if (numPublishLocks <= 1) {
+        enqueueUpdates(tld, 1, 1, tldRefreshItemsEntry.getValue());
+      } else {
+        tldRefreshItemsEntry
+            .getValue()
+            .stream()
+            .collect(
+                toImmutableSetMultimap(
+                    refreshItem -> getLockIndex(tld, numPublishLocks, refreshItem),
+                    refreshItem -> refreshItem))
+            .asMap()
+            .entrySet()
+            .forEach(
+                entry -> enqueueUpdates(tld, entry.getKey(), numPublishLocks, entry.getValue()));
+      }
+    }
+  }
+
+  /**
+   * Returns the lock index for a given refreshItem.
+   *
+   * <p>We hash the second level domain domain for all records, to group in-balliwick hosts (the
+   * only ones we refresh DNS for) with their superordinate domains. We use consistent hashing to
+   * determine the lock index because it gives us [0,N) bucketing properties out of the box, then
+   * add 1 to make indexes within [1,N].
+   */
+  private int getLockIndex(String tld, int numPublishLocks, RefreshItem refreshItem) {
+    String domain = getSecondLevelDomain(refreshItem.name(), tld);
+    return Hashing.consistentHash(hashFunction.hashString(domain, UTF_8), numPublishLocks) + 1;
+  }
+
+  /**
+   * Creates DNS refresh tasks for all writers for the tld within a lock index and batches large
+   * updates into smaller chunks.
+   */
+  private void enqueueUpdates(
+      String tld, int lockIndex, int numPublishLocks, Collection<RefreshItem> items) {
+    for (List<RefreshItem> chunk : Iterables.partition(items, tldUpdateBatchSize)) {
       DateTime earliestCreateTime =
           chunk.stream().map(RefreshItem::creationTime).min(Comparator.naturalOrder()).get();
       for (String dnsWriter : Registry.get(tld).getDnsWriters()) {
         TaskOptions options =
             withUrl(PublishDnsUpdatesAction.PATH)
                 .countdownMillis(
-                      jitterSeconds.isPresent()
+                    jitterSeconds
-                          ? random.nextInt((int) SECONDS.toMillis(jitterSeconds.get()))
+                        .map(seconds -> random.nextInt((int) SECONDS.toMillis(seconds)))
-                          : 0)
+                        .orElse(0))
                 .param(RequestParameters.PARAM_TLD, tld)
                 .param(PublishDnsUpdatesAction.PARAM_DNS_WRITER, dnsWriter)
+                .param(PublishDnsUpdatesAction.PARAM_LOCK_INDEX, Integer.toString(lockIndex))
                 .param(
-                      PublishDnsUpdatesAction.PARAM_PUBLISH_TASK_ENQUEUED,
+                    PublishDnsUpdatesAction.PARAM_NUM_PUBLISH_LOCKS,
-                      clock.nowUtc().toString())
+                    Integer.toString(numPublishLocks))
+                .param(
+                    PublishDnsUpdatesAction.PARAM_PUBLISH_TASK_ENQUEUED, clock.nowUtc().toString())
                 .param(
                     PublishDnsUpdatesAction.PARAM_REFRESH_REQUEST_CREATED,
                     earliestCreateTime.toString());
@@ -333,5 +385,4 @@ public final class ReadDnsQueueAction implements Runnable {
       }
     }
   }
-  }
 }

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

@@ -17,6 +17,7 @@ package google.registry.dns.writer.clouddns;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.model.EppResourceUtils.loadByForeignKey;
+import static google.registry.util.DomainNameUtils.getSecondLevelDomain;
 
 import com.google.api.client.googleapis.json.GoogleJsonError.ErrorInfo;
 import com.google.api.client.googleapis.json.GoogleJsonResponseException;
@@ -24,7 +25,6 @@ import com.google.api.services.dns.Dns;
 import com.google.api.services.dns.model.Change;
 import com.google.api.services.dns.model.ResourceRecordSet;
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -255,17 +255,9 @@ public class CloudDnsWriter extends BaseDnsWriter {
       return;
     }
 
-    // Extract the superordinate domain name. The TLD and host may have several dots so this
-    // must calculate a sublist.
-    ImmutableList<String> hostParts = host.parts();
-    ImmutableList<String> tldParts = tld.get().parts();
-    ImmutableList<String> domainParts =
-        hostParts.subList(hostParts.size() - tldParts.size() - 1, hostParts.size());
-    String domain = Joiner.on(".").join(domainParts);
-
     // Refresh the superordinate domain, since we shouldn't be publishing glue records if we are not
     // authoritative for the superordinate domain.
-    publishDomain(domain);
+    publishDomain(getSecondLevelDomain(hostName, tld.get().toString()));
   }
 
   /**

java/google/registry/model/registry/Registry.java

@@ -46,6 +46,7 @@ import com.googlecode.objectify.annotation.Embed;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
 import com.googlecode.objectify.annotation.Mapify;
+import com.googlecode.objectify.annotation.OnLoad;
 import com.googlecode.objectify.annotation.OnSave;
 import com.googlecode.objectify.annotation.Parent;
 import google.registry.model.Buildable;
@@ -278,6 +279,41 @@ public class Registry extends ImmutableObject implements Buildable {
    */
   Set<String> dnsWriters;
 
+  /**
+   * The number of locks we allow at once for {@link google.registry.dns.PublishDnsUpdatesAction}.
+   *
+   * <p>This should always be a positive integer- use 1 for TLD-wide locks. All {@link Registry}
+   * objects have this value default to 1.
+   *
+   * <p>WARNING: changing this parameter changes the lock name for subsequent DNS updates, and thus
+   * invalidates the locking scheme for enqueued DNS publish updates. If the {@link
+   * google.registry.dns.writer.DnsWriter} you use is not parallel-write tolerant, you must follow
+   * this procedure to change this value:
+   *
+   * <ol>
+   *   <li>Pause the DNS queue via {@link google.registry.tools.UpdateTldCommand}
+   *   <li>Change this number
+   *   <li>Let the Registry caches expire (currently 5 minutes) and drain the DNS publish queue
+   *   <li>Unpause the DNS queue
+   * </ol>
+   *
+   * <p>Failure to do so can result in parallel writes to the {@link
+   * google.registry.dns.writer.DnsWriter}, which may be dangerous depending on your implementation.
+   */
+  int numDnsPublishLocks;
+
+  /**
+   * Updates an unset numDnsPublishLocks (0) to the standard default of 1.
+   *
+   * <p>TODO(b/74010245): Remove OnLoad once all Registry objects have this value set to 1.
+   */
+  @OnLoad
+  void setDefaultNumDnsPublishLocks() {
+    if (numDnsPublishLocks == 0) {
+      numDnsPublishLocks = 1;
+    }
+  }
+
   /**
    * The unicode-aware representation of the TLD associated with this {@link Registry}.
    *
@@ -598,6 +634,11 @@ public class Registry extends ImmutableObject implements Buildable {
     return ImmutableSet.copyOf(dnsWriters);
   }
 
+  /** Returns the number of simultaneous DNS publish operations we allow at once. */
+  public int getNumDnsPublishLocks() {
+    return numDnsPublishLocks;
+  }
+
   public ImmutableSet<String> getAllowedRegistrantContactIds() {
     return nullToEmptyImmutableCopy(allowedRegistrantContactIds);
   }
@@ -686,6 +727,14 @@ public class Registry extends ImmutableObject implements Buildable {
       return this;
     }
 
+    public Builder setNumDnsPublishLocks(int numDnsPublishLocks) {
+      checkArgument(
+          numDnsPublishLocks > 0,
+          "numDnsPublishLocks must be positive when set explicitly (use 1 for TLD-wide locks)");
+      getInstance().numDnsPublishLocks = numDnsPublishLocks;
+      return this;
+    }
+
     public Builder setAddGracePeriodLength(Duration addGracePeriodLength) {
       checkArgument(
           addGracePeriodLength.isLongerThan(Duration.ZERO),
@@ -947,6 +996,11 @@ public class Registry extends ImmutableObject implements Buildable {
           instance.dnsWriters != null && !instance.dnsWriters.isEmpty(),
           "At least one DNS writer must be specified."
               + " VoidDnsWriter can be used if DNS writing isn't desired");
+      // If not set explicitly, numDnsPublishLocks defaults to 1.
+      instance.setDefaultNumDnsPublishLocks();
+      checkArgument(
+          instance.numDnsPublishLocks > 0,
+          "Number of DNS publish locks must be positive. Use 1 for TLD-wide locks.");
       instance.tldStrId = tldName;
       instance.tldUnicode = Idn.toUnicode(tldName);
       return super.build();

java/google/registry/tools/CreateOrUpdateTldCommand.java

@@ -232,6 +232,16 @@ abstract class CreateOrUpdateTldCommand extends MutatingCommand {
     description = "A comma-separated list of DnsWriter implementations to use")
   List<String> dnsWriters;
 
+  @Nullable
+  @Parameter(
+    names = {"--num_dns_publish_locks"},
+    description =
+        "The number of publish locks we allow in parallel for DNS updates under this tld "
+            + "(1 for TLD-wide locks)",
+    arity = 1
+  )
+  Integer numDnsPublishShards;
+
   @Nullable
   @Parameter(
       names = "--lrp_period",
@@ -347,6 +357,7 @@ abstract class CreateOrUpdateTldCommand extends MutatingCommand {
       Optional.ofNullable(lordnUsername).ifPresent(u -> builder.setLordnUsername(u.orElse(null)));
       Optional.ofNullable(claimsPeriodEnd).ifPresent(builder::setClaimsPeriodEnd);
       Optional.ofNullable(domainCreateRestricted).ifPresent(builder::setDomainCreateRestricted);
+      Optional.ofNullable(numDnsPublishShards).ifPresent(builder::setNumDnsPublishLocks);
       Optional.ofNullable(lrpPeriod).ifPresent(p -> builder.setLrpPeriod(p.orElse(null)));
 
       if (premiumListName != null) {
@@ -410,7 +421,7 @@ abstract class CreateOrUpdateTldCommand extends MutatingCommand {
           Joiner.on(", ").join(invalidNames),
           tld);
       if (overrideReservedListRules) {
-        System.err.println("Error overriden: " + errMsg);
+        System.err.println("Error overridden: " + errMsg);
       } else {
         throw new IllegalArgumentException(errMsg);
       }

java/google/registry/util/DomainNameUtils.java

@@ -18,7 +18,9 @@ import static com.google.common.base.Preconditions.checkArgument;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
 import com.google.common.base.Ascii;
+import com.google.common.base.Joiner;
 import com.google.common.base.Strings;
+import com.google.common.collect.ImmutableList;
 import com.google.common.net.InternetDomainName;
 
 /** Utility methods related to domain names. */
@@ -78,5 +80,34 @@ public final class DomainNameUtils {
     return domainName.parent().toString();
   }
 
+  /**
+   * Returns the second level domain name for a fully qualified host name under a given tld.
+   *
+   * <p>This function is merely a string parsing utility, and does not verify if the tld is operated
+   * by the registry.
+   *
+   * @throws IllegalArgumentException if either argument is null or empty, or the domain name is not
+   *     under the tld
+   */
+  public static String getSecondLevelDomain(String hostName, String tld) {
+    checkArgument(
+        !Strings.isNullOrEmpty(hostName),
+        "hostName cannot be null or empty");
+    checkArgument(!Strings.isNullOrEmpty(tld), "tld cannot be null or empty");
+    ImmutableList<String> domainParts = InternetDomainName.from(hostName).parts();
+    ImmutableList<String> tldParts = InternetDomainName.from(tld).parts();
+    checkArgument(
+        domainParts.size() > tldParts.size(),
+        "hostName must be at least one level below the tld");
+    checkArgument(
+        domainParts
+            .subList(domainParts.size() - tldParts.size(), domainParts.size())
+            .equals(tldParts),
+        "hostName must be under the tld");
+    ImmutableList<String> sldParts =
+        domainParts.subList(domainParts.size() - tldParts.size() - 1, domainParts.size());
+    return Joiner.on(".").join(sldParts);
+  }
+
   private DomainNameUtils() {}
 }

javatests/google/registry/dns/BUILD

@@ -20,6 +20,7 @@ java_library(
         "//java/google/registry/model",
         "//java/google/registry/module/backend",
         "//java/google/registry/request",
+        "//java/google/registry/request/lock",
         "//java/google/registry/util",
         "//javatests/google/registry/testing",
         "//third_party/objectify:objectify-v4_1",

javatests/google/registry/dns/PublishDnsUpdatesActionTest.java

@@ -20,10 +20,12 @@ import static google.registry.testing.DatastoreHelper.persistActiveDomain;
 import static google.registry.testing.DatastoreHelper.persistActiveSubordinateHost;
 import static google.registry.testing.DatastoreHelper.persistResource;
 import static google.registry.testing.JUnitBackports.assertThrows;
+import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -35,6 +37,7 @@ import google.registry.model.domain.DomainResource;
 import google.registry.model.ofy.Ofy;
 import google.registry.model.registry.Registry;
 import google.registry.request.HttpException.ServiceUnavailableException;
+import google.registry.request.lock.LockHandler;
 import google.registry.testing.AppEngineRule;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeLockHandler;
@@ -96,6 +99,8 @@ public class PublishDnsUpdatesActionTest {
     action.dnsWriterProxy = new DnsWriterProxy(ImmutableMap.of("correctWriter", dnsWriter));
     action.dnsMetrics = dnsMetrics;
     action.dnsQueue = dnsQueue;
+    action.lockIndex = 1;
+    action.numPublishLocks = 1;
     action.lockHandler = lockHandler;
     action.clock = clock;
     return action;
@@ -105,12 +110,12 @@ public class PublishDnsUpdatesActionTest {
   public void testHost_published() throws Exception {
     action = createAction("xn--q9jyb4c");
     action.hosts = ImmutableSet.of("ns1.example.xn--q9jyb4c");
+
     action.run();
 
     verify(dnsWriter).publishHost("ns1.example.xn--q9jyb4c");
     verify(dnsWriter).commit();
     verifyNoMoreInteractions(dnsWriter);
-
     verify(dnsMetrics).incrementPublishDomainRequests(0, PublishStatus.ACCEPTED);
     verify(dnsMetrics).incrementPublishDomainRequests(0, PublishStatus.REJECTED);
     verify(dnsMetrics).incrementPublishHostRequests(1, PublishStatus.ACCEPTED);
@@ -124,7 +129,6 @@ public class PublishDnsUpdatesActionTest {
             Duration.standardHours(2),
             Duration.standardHours(1));
     verifyNoMoreInteractions(dnsMetrics);
-
     verifyNoMoreInteractions(dnsQueue);
   }
 
@@ -132,12 +136,12 @@ public class PublishDnsUpdatesActionTest {
   public void testDomain_published() throws Exception {
     action = createAction("xn--q9jyb4c");
     action.domains = ImmutableSet.of("example.xn--q9jyb4c");
+
     action.run();
 
     verify(dnsWriter).publishDomain("example.xn--q9jyb4c");
     verify(dnsWriter).commit();
     verifyNoMoreInteractions(dnsWriter);
-
     verify(dnsMetrics).incrementPublishDomainRequests(1, PublishStatus.ACCEPTED);
     verify(dnsMetrics).incrementPublishDomainRequests(0, PublishStatus.REJECTED);
     verify(dnsMetrics).incrementPublishHostRequests(0, PublishStatus.ACCEPTED);
@@ -151,10 +155,27 @@ public class PublishDnsUpdatesActionTest {
             Duration.standardHours(2),
             Duration.standardHours(1));
     verifyNoMoreInteractions(dnsMetrics);
-
     verifyNoMoreInteractions(dnsQueue);
   }
 
+  @Test
+  public void testAction_acquiresCorrectLock() throws Exception {
+    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.run();
+
+    verify(mockLockHandler)
+        .executeWithLocks(
+            action, "xn--q9jyb4c", Duration.standardSeconds(10), "DNS updates-lock 2 of 4");
+  }
+
   @Test
   public void testPublish_commitFails() throws Exception {
     action = createAction("xn--q9jyb4c");
@@ -163,6 +184,7 @@ public class PublishDnsUpdatesActionTest {
         ImmutableSet.of(
             "ns1.example.xn--q9jyb4c", "ns2.example.xn--q9jyb4c", "ns1.example2.xn--q9jyb4c");
     doThrow(new RuntimeException()).when(dnsWriter).commit();
+
     assertThrows(RuntimeException.class, action::run);
 
     verify(dnsMetrics).incrementPublishDomainRequests(2, PublishStatus.ACCEPTED);
@@ -178,7 +200,6 @@ public class PublishDnsUpdatesActionTest {
             Duration.standardHours(2),
             Duration.standardHours(1));
     verifyNoMoreInteractions(dnsMetrics);
-
     verifyNoMoreInteractions(dnsQueue);
   }
 
@@ -188,6 +209,7 @@ public class PublishDnsUpdatesActionTest {
     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.run();
 
     verify(dnsWriter).publishDomain("example.xn--q9jyb4c");
@@ -197,7 +219,6 @@ public class PublishDnsUpdatesActionTest {
     verify(dnsWriter).publishHost("ns1.example2.xn--q9jyb4c");
     verify(dnsWriter).commit();
     verifyNoMoreInteractions(dnsWriter);
-
     verify(dnsMetrics).incrementPublishDomainRequests(2, PublishStatus.ACCEPTED);
     verify(dnsMetrics).incrementPublishDomainRequests(0, PublishStatus.REJECTED);
     verify(dnsMetrics).incrementPublishHostRequests(3, PublishStatus.ACCEPTED);
@@ -211,7 +232,6 @@ public class PublishDnsUpdatesActionTest {
             Duration.standardHours(2),
             Duration.standardHours(1));
     verifyNoMoreInteractions(dnsMetrics);
-
     verifyNoMoreInteractions(dnsQueue);
   }
 
@@ -220,11 +240,11 @@ public class PublishDnsUpdatesActionTest {
     action = createAction("xn--q9jyb4c");
     action.domains = ImmutableSet.of("example.com", "example2.com");
     action.hosts = ImmutableSet.of("ns1.example.com", "ns2.example.com", "ns1.example2.com");
+
     action.run();
 
     verify(dnsWriter).commit();
     verifyNoMoreInteractions(dnsWriter);
-
     verify(dnsMetrics).incrementPublishDomainRequests(0, PublishStatus.ACCEPTED);
     verify(dnsMetrics).incrementPublishDomainRequests(2, PublishStatus.REJECTED);
     verify(dnsMetrics).incrementPublishHostRequests(0, PublishStatus.ACCEPTED);
@@ -238,7 +258,6 @@ public class PublishDnsUpdatesActionTest {
             Duration.standardHours(2),
             Duration.standardHours(1));
     verifyNoMoreInteractions(dnsMetrics);
-
     verifyNoMoreInteractions(dnsQueue);
   }
 
@@ -248,12 +267,12 @@ public class PublishDnsUpdatesActionTest {
     action.domains = ImmutableSet.of("example.com", "example2.com");
     action.hosts = ImmutableSet.of("ns1.example.com", "ns2.example.com", "ns1.example2.com");
     action.lockHandler = new FakeLockHandler(false);
+
     ServiceUnavailableException thrown =
         assertThrows(ServiceUnavailableException.class, action::run);
+
     assertThat(thrown).hasMessageThat().contains("Lock failure");
-
     verifyNoMoreInteractions(dnsWriter);
-
     verify(dnsMetrics)
         .recordActionResult(
             "correctWriter",
@@ -262,7 +281,56 @@ public class PublishDnsUpdatesActionTest {
             Duration.standardHours(2),
             Duration.standardHours(1));
     verifyNoMoreInteractions(dnsMetrics);
+    verifyNoMoreInteractions(dnsQueue);
+  }
 
+  @Test
+  public void testParam_invalidLockIndex() throws Exception {
+    persistResource(Registry.get("xn--q9jyb4c").asBuilder().setNumDnsPublishLocks(4).build());
+    action = createAction("xn--q9jyb4c");
+    action.domains = ImmutableSet.of("example.com");
+    action.hosts = ImmutableSet.of("ns1.example.com");
+    action.lockIndex = 5;
+    action.numPublishLocks = 4;
+
+    action.run();
+
+    verifyNoMoreInteractions(dnsWriter);
+    verify(dnsMetrics)
+        .recordActionResult(
+            "correctWriter",
+            ActionStatus.BAD_LOCK_INDEX,
+            2,
+            Duration.standardHours(2),
+            Duration.standardHours(1));
+    verifyNoMoreInteractions(dnsMetrics);
+    verify(dnsQueue).addDomainRefreshTask("example.com");
+    verify(dnsQueue).addHostRefreshTask("ns1.example.com");
+    verifyNoMoreInteractions(dnsQueue);
+  }
+
+  @Test
+  public void testRegistryParam_mismatchedMaxLocks() throws Exception {
+    persistResource(Registry.get("xn--q9jyb4c").asBuilder().setNumDnsPublishLocks(4).build());
+    action = createAction("xn--q9jyb4c");
+    action.domains = ImmutableSet.of("example.com");
+    action.hosts = ImmutableSet.of("ns1.example.com");
+    action.lockIndex = 3;
+    action.numPublishLocks = 5;
+
+    action.run();
+
+    verifyNoMoreInteractions(dnsWriter);
+    verify(dnsMetrics)
+        .recordActionResult(
+            "correctWriter",
+            ActionStatus.BAD_LOCK_INDEX,
+            2,
+            Duration.standardHours(2),
+            Duration.standardHours(1));
+    verifyNoMoreInteractions(dnsMetrics);
+    verify(dnsQueue).addDomainRefreshTask("example.com");
+    verify(dnsQueue).addHostRefreshTask("ns1.example.com");
     verifyNoMoreInteractions(dnsQueue);
   }
 
@@ -272,10 +340,10 @@ public class PublishDnsUpdatesActionTest {
     action.domains = ImmutableSet.of("example.com", "example2.com");
     action.hosts = ImmutableSet.of("ns1.example.com", "ns2.example.com", "ns1.example2.com");
     action.dnsWriter = "wrongWriter";
+
     action.run();
 
     verifyNoMoreInteractions(dnsWriter);
-
     verify(dnsMetrics)
         .recordActionResult(
             "wrongWriter",
@@ -284,7 +352,6 @@ public class PublishDnsUpdatesActionTest {
             Duration.standardHours(2),
             Duration.standardHours(1));
     verifyNoMoreInteractions(dnsMetrics);
-
     verify(dnsQueue).addDomainRefreshTask("example.com");
     verify(dnsQueue).addDomainRefreshTask("example2.com");
     verify(dnsQueue).addHostRefreshTask("ns1.example.com");

javatests/google/registry/dns/ReadDnsQueueActionTest.java

@@ -38,6 +38,7 @@ import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.hash.Hashing;
 import com.google.common.net.InternetDomainName;
 import google.registry.dns.DnsConstants.TargetType;
 import google.registry.model.registry.Registry;
@@ -91,7 +92,7 @@ public class ReadDnsQueueActionTest {
     // Because of b/73372999 - the FakeClock can't be in the past, or the TaskQueues stop working.
     // To make sure it's never in the past, we set the date far-far into the future
     clock.setTo(DateTime.parse("3000-01-01TZ"));
-    createTlds("com", "net", "example");
+    createTlds("com", "net", "example", "multilock.uk");
     persistResource(
         Registry.get("com").asBuilder().setDnsWriters(ImmutableSet.of("comWriter")).build());
     persistResource(
@@ -102,6 +103,12 @@ public class ReadDnsQueueActionTest {
             .setTldType(TldType.TEST)
             .setDnsWriters(ImmutableSet.of("exampleWriter"))
             .build());
+    persistResource(
+        Registry.get("multilock.uk")
+            .asBuilder()
+            .setNumDnsPublishLocks(1000)
+            .setDnsWriters(ImmutableSet.of("multilockWriter"))
+            .build());
     dnsQueue = DnsQueue.createForTesting(clock);
   }
 
@@ -112,10 +119,12 @@ public class ReadDnsQueueActionTest {
     action.clock = clock;
     action.dnsQueue = dnsQueue;
     action.dnsPublishPushQueue = QueueFactory.getQueue(DNS_PUBLISH_PUSH_QUEUE_NAME);
+    action.hashFunction = Hashing.murmur3_32();
     action.taskEnqueuer = new TaskEnqueuer(new Retrier(null, 1));
     action.jitterSeconds = Optional.empty();
     // Advance the time a little, to ensure that leaseTasks() returns all tasks.
     clock.advanceBy(Duration.standardHours(1));
+
     action.run();
   }
 
@@ -149,6 +158,9 @@ public class ReadDnsQueueActionTest {
                     .param("dnsWriter", tldToDnsWriter.getValue())
                     .param("itemsCreated", "3000-01-01T00:00:00.000Z")
                     .param("enqueued", "3000-01-01T01:00:00.000Z")
+                    // Single-lock TLDs should use lock 1 of 1 by default
+                    .param("lockIndex", "1")
+                    .param("numPublishLocks", "1")
                     .header("content-type", "application/x-www-form-urlencoded")));
   }
 
@@ -157,7 +169,9 @@ public class ReadDnsQueueActionTest {
     dnsQueue.addDomainRefreshTask("domain.com");
     dnsQueue.addDomainRefreshTask("domain.net");
     dnsQueue.addDomainRefreshTask("domain.example");
+
     run();
+
     assertNoTasksEnqueued(DNS_PULL_QUEUE_NAME);
     assertTasksEnqueued(
         DNS_PUBLISH_PUSH_QUEUE_NAME,
@@ -167,11 +181,13 @@ public class ReadDnsQueueActionTest {
   }
 
   @Test
-  public void testSuccess_allTlds() throws Exception {
+  public void testSuccess_allSingleLockTlds() throws Exception {
     dnsQueue.addDomainRefreshTask("domain.com");
     dnsQueue.addDomainRefreshTask("domain.net");
     dnsQueue.addDomainRefreshTask("domain.example");
+
     run();
+
     assertNoTasksEnqueued(DNS_PULL_QUEUE_NAME);
     assertTldsEnqueuedInPushQueue(
         ImmutableMultimap.of("com", "comWriter", "net", "netWriter", "example", "exampleWriter"));
@@ -186,7 +202,9 @@ public class ReadDnsQueueActionTest {
             .mapToObj(i -> String.format("domain_%04d.com", i))
             .collect(toImmutableList());
     domains.forEach(dnsQueue::addDomainRefreshTask);
+
     run();
+
     assertNoTasksEnqueued(DNS_PULL_QUEUE_NAME);
     ImmutableList<ImmutableMultimap<String, String>> queuedParams =
         getQueuedParams(DNS_PUBLISH_PUSH_QUEUE_NAME);
@@ -209,7 +227,9 @@ public class ReadDnsQueueActionTest {
             .setDnsWriters(ImmutableSet.of("comWriter", "otherWriter"))
             .build());
     dnsQueue.addDomainRefreshTask("domain.com");
+
     run();
+
     assertNoTasksEnqueued(DNS_PULL_QUEUE_NAME);
     assertTldsEnqueuedInPushQueue(ImmutableMultimap.of("com", "comWriter", "com", "otherWriter"));
   }
@@ -222,7 +242,9 @@ public class ReadDnsQueueActionTest {
     dnsQueue.addDomainRefreshTask("domain2.com");
     clock.setTo(DateTime.parse("3000-02-05TZ"));
     dnsQueue.addDomainRefreshTask("domain3.com");
+
     run();
+
     assertNoTasksEnqueued(DNS_PULL_QUEUE_NAME);
     assertThat(getQueuedParams(DNS_PUBLISH_PUSH_QUEUE_NAME))
         .containsExactly(
@@ -234,6 +256,8 @@ public class ReadDnsQueueActionTest {
                 .put("domains", "domain1.com")
                 .put("domains", "domain2.com")
                 .put("domains", "domain3.com")
+                .put("lockIndex", "1")
+                .put("numPublishLocks", "1")
                 .build());
   }
 
@@ -243,7 +267,9 @@ public class ReadDnsQueueActionTest {
     dnsQueue.addDomainRefreshTask("domain.com");
     dnsQueue.addDomainRefreshTask("domain.net");
     dnsQueue.addDomainRefreshTask("domain.example");
+
     run();
+
     assertTasksEnqueued(DNS_PULL_QUEUE_NAME, createDomainRefreshTaskMatcher("domain.net"));
     assertTldsEnqueuedInPushQueue(
         ImmutableMultimap.of("com", "comWriter", "example", "exampleWriter"));
@@ -260,7 +286,9 @@ public class ReadDnsQueueActionTest {
                 .param(DNS_TARGET_TYPE_PARAM, TargetType.DOMAIN.toString())
                 .param(DNS_TARGET_NAME_PARAM, "domain.unknown")
                 .param(PARAM_TLD, "unknown"));
+
     run();
+
     assertTasksEnqueued(DNS_PULL_QUEUE_NAME, createDomainRefreshTaskMatcher("domain.unknown"));
     assertTldsEnqueuedInPushQueue(
         ImmutableMultimap.of("com", "comWriter", "example", "exampleWriter"));
@@ -279,7 +307,9 @@ public class ReadDnsQueueActionTest {
                 .param(DNS_TARGET_NAME_PARAM, "domain.wrongtld")
                 .param(DNS_TARGET_CREATE_TIME_PARAM, "3000-01-01TZ")
                 .param(PARAM_TLD, "net"));
+
     run();
+
     assertNoTasksEnqueued(DNS_PULL_QUEUE_NAME);
     assertTldsEnqueuedInPushQueue(
         ImmutableMultimap.of("com", "comWriter", "example", "exampleWriter", "net", "netWriter"));
@@ -295,7 +325,9 @@ public class ReadDnsQueueActionTest {
                 .method(Method.PULL)
                 .param(DNS_TARGET_TYPE_PARAM, TargetType.DOMAIN.toString())
                 .param(DNS_TARGET_NAME_PARAM, "domain.net"));
+
     run();
+
     // The corrupt task isn't in the pull queue, but also isn't in the push queue
     assertNoTasksEnqueued(DNS_PULL_QUEUE_NAME);
     assertTldsEnqueuedInPushQueue(
@@ -312,7 +344,9 @@ public class ReadDnsQueueActionTest {
                 .method(Method.PULL)
                 .param(DNS_TARGET_TYPE_PARAM, TargetType.DOMAIN.toString())
                 .param(PARAM_TLD, "net"));
+
     run();
+
     // The corrupt task isn't in the pull queue, but also isn't in the push queue
     assertNoTasksEnqueued(DNS_PULL_QUEUE_NAME);
     assertTldsEnqueuedInPushQueue(
@@ -329,7 +363,9 @@ public class ReadDnsQueueActionTest {
                 .method(Method.PULL)
                 .param(DNS_TARGET_NAME_PARAM, "domain.net")
                 .param(PARAM_TLD, "net"));
+
     run();
+
     // The corrupt task isn't in the pull queue, but also isn't in the push queue
     assertNoTasksEnqueued(DNS_PULL_QUEUE_NAME);
     assertTldsEnqueuedInPushQueue(
@@ -347,7 +383,9 @@ public class ReadDnsQueueActionTest {
                 .param(DNS_TARGET_TYPE_PARAM, "Wrong type")
                 .param(DNS_TARGET_NAME_PARAM, "domain.net")
                 .param(PARAM_TLD, "net"));
+
     run();
+
     // The corrupt task isn't in the pull queue, but also isn't in the push queue
     assertNoTasksEnqueued(DNS_PULL_QUEUE_NAME);
     assertTldsEnqueuedInPushQueue(
@@ -359,7 +397,9 @@ public class ReadDnsQueueActionTest {
     dnsQueue.addHostRefreshTask("ns1.domain.com");
     dnsQueue.addDomainRefreshTask("domain.net");
     dnsQueue.addZoneRefreshTask("example");
+
     run();
+
     assertNoTasksEnqueued(DNS_PULL_QUEUE_NAME);
     assertTasksEnqueued(
         DNS_PUBLISH_PUSH_QUEUE_NAME,
@@ -408,8 +448,47 @@ public class ReadDnsQueueActionTest {
         }
       }
     }
+
     run();
+
     assertNoTasksEnqueued(DNS_PULL_QUEUE_NAME);
     assertTasksEnqueued(DNS_PUBLISH_PUSH_QUEUE_NAME, expectedTasks);
   }
+
+  @Test
+  public void testSuccess_lockGroupsHostBySuperordinateDomain() throws Exception {
+    dnsQueue.addDomainRefreshTask("hello.multilock.uk");
+    dnsQueue.addHostRefreshTask("ns1.abc.hello.multilock.uk");
+    dnsQueue.addHostRefreshTask("ns2.hello.multilock.uk");
+    dnsQueue.addDomainRefreshTask("another.multilock.uk");
+    dnsQueue.addHostRefreshTask("ns3.def.another.multilock.uk");
+    dnsQueue.addHostRefreshTask("ns4.another.multilock.uk");
+
+    run();
+
+    assertNoTasksEnqueued(DNS_PULL_QUEUE_NAME);
+    // Expect two different groups; in-balliwick hosts are locked with their superordinate domains.
+    assertTasksEnqueued(
+        DNS_PUBLISH_PUSH_QUEUE_NAME,
+        new TaskMatcher()
+            .url(PublishDnsUpdatesAction.PATH)
+            .param("tld", "multilock.uk")
+            .param("dnsWriter", "multilockWriter")
+            .param("itemsCreated", "3000-01-01T00:00:00.000Z")
+            .param("enqueued", "3000-01-01T01:00:00.000Z")
+            .param("domains", "hello.multilock.uk")
+            .param("hosts", "ns1.abc.hello.multilock.uk")
+            .param("hosts", "ns2.hello.multilock.uk")
+            .header("content-type", "application/x-www-form-urlencoded"),
+        new TaskMatcher()
+            .url(PublishDnsUpdatesAction.PATH)
+            .param("tld", "multilock.uk")
+            .param("dnsWriter", "multilockWriter")
+            .param("itemsCreated", "3000-01-01T00:00:00.000Z")
+            .param("enqueued", "3000-01-01T01:00:00.000Z")
+            .param("domains", "another.multilock.uk")
+            .param("hosts", "ns3.def.another.multilock.uk")
+            .param("hosts", "ns4.another.multilock.uk")
+            .header("content-type", "application/x-www-form-urlencoded"));
+  }
 }

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

@@ -98,6 +98,19 @@ public class RegistryTest extends EntityTestCase {
     assertThat(registry.getStandardRestoreCost()).isEqualTo(Money.of(USD, 42));
   }
 
+  @Test
+  public void testDefaultNumDnsPublishShards_equalToOne() {
+    Registry registry = Registry.get("tld").asBuilder().build();
+    assertThat(registry.getNumDnsPublishLocks()).isEqualTo(1);
+  }
+
+  @Test
+  public void testSettingNumDnsPublishShards() {
+    Registry registry =
+        Registry.get("tld").asBuilder().setNumDnsPublishLocks(2).build();
+    assertThat(registry.getNumDnsPublishLocks()).isEqualTo(2);
+  }
+
   @Test
   public void testSetReservedList_doesntMutateExistingRegistry() {
     ReservedList rl15 = persistReservedList(
@@ -458,6 +471,26 @@ public class RegistryTest extends EntityTestCase {
     assertThat(thrown).hasMessageThat().contains("restoreBillingCost cannot be negative");
   }
 
+  @Test
+  public void testFailure_nonPositiveNumDnsPublishLocks() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> Registry.get("tld").asBuilder().setNumDnsPublishLocks(-1));
+    assertThat(thrown)
+        .hasMessageThat()
+        .contains(
+            "numDnsPublishLocks must be positive when set explicitly (use 1 for TLD-wide locks)");
+    thrown =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> Registry.get("tld").asBuilder().setNumDnsPublishLocks(0));
+    assertThat(thrown)
+        .hasMessageThat()
+        .contains(
+            "numDnsPublishLocks must be positive when set explicitly (use 1 for TLD-wide locks)");
+  }
+
   @Test
   public void testFailure_negativeServerStatusChangeBillingCost() {
     IllegalArgumentException thrown =

javatests/google/registry/model/schema.txt

@@ -681,6 +681,7 @@ class google.registry.model.registry.Registry {
   google.registry.model.common.TimedTransitionProperty<org.joda.money.Money, google.registry.model.registry.Registry$BillingCostTransition> eapFeeSchedule;
   google.registry.model.common.TimedTransitionProperty<org.joda.money.Money, google.registry.model.registry.Registry$BillingCostTransition> renewBillingCostTransitions;
   google.registry.model.registry.Registry$TldType tldType;
+  int numDnsPublishLocks;
   java.lang.String driveFolderId;
   java.lang.String lordnUsername;
   java.lang.String pricingEngineClassName;

javatests/google/registry/tools/CreateTldCommandTest.java

@@ -540,7 +540,8 @@ public class CreateTldCommandTest extends CommandTestCase<CreateTldCommand> {
     System.setErr(new PrintStream(errContent));
     runReservedListsTestOverride("common_abuse,tld_banned");
     String errMsg =
-        "Error overriden: The reserved list(s) tld_banned cannot be applied to the tld xn--q9jyb4c";
+        "Error overridden: The reserved list(s) tld_banned "
+            + "cannot be applied to the tld xn--q9jyb4c";
     assertThat(errContent.toString()).contains(errMsg);
     System.setOut(null);
   }

javatests/google/registry/tools/UpdateTldCommandTest.java

@@ -964,7 +964,8 @@ public class UpdateTldCommandTest extends CommandTestCase<UpdateTldCommand> {
     System.setErr(new PrintStream(errContent));
     runReservedListsTestOverride("common_abuse,tld_banned");
     String errMsg =
-        "Error overriden: The reserved list(s) tld_banned cannot be applied to the tld xn--q9jyb4c";
+        "Error overridden: The reserved list(s) tld_banned "
+            + "cannot be applied to the tld xn--q9jyb4c";
     assertThat(errContent.toString()).contains(errMsg);
     System.setOut(null);
   }

javatests/google/registry/util/DomainNameUtilsTest.java

@@ -17,6 +17,7 @@ package google.registry.util;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.testing.JUnitBackports.assertThrows;
 import static google.registry.util.DomainNameUtils.canonicalizeDomainName;
+import static google.registry.util.DomainNameUtils.getSecondLevelDomain;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -26,7 +27,7 @@ import org.junit.runners.JUnit4;
 @RunWith(JUnit4.class)
 public class DomainNameUtilsTest {
   @Test
-  public void testCanonicalizeDomainName() throws Exception {
+  public void testCanonicalizeDomainName() {
     assertThat(canonicalizeDomainName("foo")).isEqualTo("foo");
     assertThat(canonicalizeDomainName("FOO")).isEqualTo("foo");
     assertThat(canonicalizeDomainName("foo.tld")).isEqualTo("foo.tld");
@@ -40,7 +41,33 @@ public class DomainNameUtilsTest {
   }
 
   @Test
-  public void testCanonicalizeDomainName_acePrefixUnicodeChars() throws Exception {
+  public void testCanonicalizeDomainName_acePrefixUnicodeChars() {
     assertThrows(IllegalArgumentException.class, () -> canonicalizeDomainName("xn--みんな"));
   }
+
+  @Test
+  public void testGetSecondLevelDomain_returnsProperDomain() {
+    assertThat(getSecondLevelDomain("foo.bar", "bar")).isEqualTo("foo.bar");
+    assertThat(getSecondLevelDomain("ns1.foo.bar", "bar")).isEqualTo("foo.bar");
+    assertThat(getSecondLevelDomain("ns1.abc.foo.bar", "bar")).isEqualTo("foo.bar");
+    assertThat(getSecondLevelDomain("ns1.abc.foo.bar", "foo.bar")).isEqualTo("abc.foo.bar");
+  }
+
+  @Test
+  public void testGetSecondLevelDomain_insufficientDomainNameDepth() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class, () -> getSecondLevelDomain("bar", "bar"));
+    assertThat(thrown)
+        .hasMessageThat()
+        .isEqualTo("hostName must be at least one level below the tld");
+  }
+
+  @Test
+  public void testGetSecondLevelDomain_domainNotUnderTld() {
+    IllegalArgumentException thrown =
+        assertThrows(
+            IllegalArgumentException.class, () -> getSecondLevelDomain("foo.bar", "abc"));
+    assertThat(thrown).hasMessageThat().isEqualTo("hostName must be under the tld");
+  }
 }