Batch DNS refresh requests on host renames

core/src/main/java/google/registry/dns/RefreshDnsOnHostRenameAction.java

@@ -14,23 +14,30 @@
 
 package google.registry.dns;
 
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.dns.DnsUtils.requestDomainDnsRefresh;
 import static google.registry.dns.RefreshDnsOnHostRenameAction.PATH;
 import static google.registry.model.EppResourceUtils.getLinkedDomainKeys;
+import static google.registry.model.EppResourceUtils.isDeleted;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static jakarta.servlet.http.HttpServletResponse.SC_NO_CONTENT;
+import static jakarta.servlet.http.HttpServletResponse.SC_OK;
 
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
 import com.google.common.net.MediaType;
-import google.registry.model.EppResourceUtils;
 import google.registry.model.domain.Domain;
 import google.registry.model.host.Host;
 import google.registry.persistence.VKey;
 import google.registry.request.Action;
 import google.registry.request.Parameter;
 import google.registry.request.Response;
 import google.registry.request.auth.Auth;
+import google.registry.util.Clock;
 import jakarta.inject.Inject;
 import java.time.Instant;
+import java.util.List;
+import java.util.Optional;
 
 @Action(
     service = Action.Service.BACKEND,
@@ -43,45 +50,52 @@ public class RefreshDnsOnHostRenameAction implements Runnable {
   public static final String PARAM_HOST_KEY = "hostKey";
   public static final String PATH = "/_dr/task/refreshDnsOnHostRename";
 
+  private static final int DNS_REFRESH_BATCH_SIZE = 1000;
+
   private final VKey<Host> hostKey;
   private final Response response;
+  private final Clock clock;
 
   @Inject
-  RefreshDnsOnHostRenameAction(@Parameter(PARAM_HOST_KEY) String hostKey, Response response) {
+  RefreshDnsOnHostRenameAction(
+      @Parameter(PARAM_HOST_KEY) String hostKey, Response response, Clock clock) {
     this.hostKey = VKey.createEppVKeyFromString(hostKey);
     this.response = response;
+    this.clock = clock;
   }
 
   @Override
   public void run() {
-    tm().transact(
-            () -> {
-              Instant now = tm().getTxTime();
-              Host host = tm().loadByKeyIfPresent(hostKey).orElse(null);
-              boolean hostValid = true;
-              String failureMessage = null;
-              if (host == null) {
-                hostValid = false;
-                failureMessage = String.format("Host to refresh does not exist: %s", hostKey);
-              } else if (EppResourceUtils.isDeleted(host, now)) {
-                hostValid = false;
-                failureMessage =
-                    String.format("Host to refresh is already deleted: %s", host.getHostName());
-              } else {
-                getLinkedDomainKeys(
-                        host.createVKey(), host.getUpdateTimestamp().getTimestamp(), null)
-                    .stream()
-                    .map(domainKey -> tm().loadByKey(domainKey))
-                    .filter(Domain::shouldPublishToDns)
-                    .forEach(domain -> requestDomainDnsRefresh(domain.getDomainName()));
-              }
+    Optional<Host> optionalHost = tm().transact(() -> tm().loadByKeyIfPresent(hostKey));
+    if (optionalHost.isEmpty()) {
+      setFailedStatus(String.format("Host to refresh does not exist: %s", hostKey));
+      return;
+    }
+    Instant now = clock.now();
+    Host host = optionalHost.get();
+    if (isDeleted(host, now)) {
+      setFailedStatus(String.format("Host to refresh is already deleted: %s", host.getHostName()));
+      return;
+    }
+    ImmutableSet<VKey<Domain>> linkedDomainKeys =
+        getLinkedDomainKeys(hostKey, host.getUpdateTimestamp().getTimestamp(), null);
+    for (List<VKey<Domain>> batch : Iterables.partition(linkedDomainKeys, DNS_REFRESH_BATCH_SIZE)) {
+      tm().transact(
+              () -> {
+                ImmutableSet<String> domainNames =
+                    tm().loadByKeysIfPresent(batch).values().stream()
+                        .filter(Domain::shouldPublishToDns)
+                        .map(Domain::getDomainName)
+                        .collect(toImmutableSet());
+                requestDomainDnsRefresh(domainNames);
+              });
+    }
+    response.setStatus(SC_OK);
+  }
 
-              if (!hostValid) {
-                // Set the response status code to be 204 so to not retry.
-                response.setContentType(MediaType.PLAIN_TEXT_UTF_8);
-                response.setStatus(SC_NO_CONTENT);
-                response.setPayload(failureMessage);
-              }
-            });
+  private void setFailedStatus(String message) {
+    response.setContentType(MediaType.PLAIN_TEXT_UTF_8);
+    response.setStatus(SC_NO_CONTENT);
+    response.setPayload(message);
   }
 }

core/src/main/java/google/registry/persistence/transaction/TransactionManager.java

@@ -271,7 +271,7 @@ <T> ImmutableMap<VKey<? extends T>, T> loadByKeysIfPresent(
    * A runnable that allows for checked exceptions to be thrown.
    *
    * <p>This makes it easier to write lambdas without having to worry about wrapping and re-throwing
-   * checked excpetions as unchecked ones.
+   * checked exceptions as unchecked ones.
    */
   @FunctionalInterface
   interface ThrowingRunnable {

core/src/test/java/google/registry/dns/RefreshDnsOnHostRenameActionTest.java

@@ -28,6 +28,7 @@
 import static jakarta.servlet.http.HttpServletResponse.SC_OK;
 
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
 import google.registry.model.eppcommon.StatusValue;
 import google.registry.model.host.Host;
 import google.registry.persistence.transaction.JpaTestExtensions;
@@ -52,7 +53,7 @@ public class RefreshDnsOnHostRenameActionTest {
   private RefreshDnsOnHostRenameAction action;
 
   private void createAction(String hostKey) {
-    action = new RefreshDnsOnHostRenameAction(hostKey, response);
+    action = new RefreshDnsOnHostRenameAction(hostKey, response, clock);
   }
 
   @BeforeEach
@@ -99,4 +100,28 @@ void testFailure_deletedHost() {
     assertThat(response.getPayload())
         .isEqualTo("Host to refresh is already deleted: ns1.example.tld");
   }
+
+  @Test
+  void testSuccess_multipleBatches() {
+    Host host = persistActiveHost("ns1.example.tld");
+    ImmutableSet.Builder<String> domainNamesBuilder = new ImmutableSet.Builder<>();
+    for (int i = 1; i <= 1001; i++) {
+      String domainName = "example" + i + ".tld";
+      domainNamesBuilder.add(domainName);
+      persistResource(newDomain(domainName, host));
+    }
+    createAction(host.createVKey().stringify());
+    action.run();
+    assertDomainDnsRequests(Iterables.toArray(domainNamesBuilder.build(), String.class));
+    assertThat(response.getStatus()).isEqualTo(SC_OK);
+  }
+
+  @Test
+  void testSuccess_noLinkedDomains() {
+    Host host = persistActiveHost("ns1.example.tld");
+    createAction(host.createVKey().stringify());
+    action.run();
+    assertNoDnsRequests();
+    assertThat(response.getStatus()).isEqualTo(SC_OK);
+  }
 }