Make EppResource.loadCached() use batched fetch (#652)

* Make EppResource.loadCached() use batched fetch

Use a batched fetch (ofy().load().keys(...)) from datastore in
EppResource.loadCached().

To support this, convert TransactionManager.load(Iterable<VKey>) to accept the
more flexible generic parameters and return a map.

* Simplify datastore key streaming

* Changes requested in review.

core/src/main/java/google/registry/model/EppResource.java

@@ -16,7 +16,6 @@ package google.registry.model;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
-import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.Sets.difference;
 import static com.google.common.collect.Sets.union;
 import static google.registry.config.RegistryConfig.getEppResourceCachingDuration;
@@ -47,7 +46,6 @@ import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ExecutionException;
-import java.util.stream.StreamSupport;
 import javax.persistence.Access;
 import javax.persistence.AccessType;
 import javax.persistence.Column;
@@ -357,7 +355,7 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable {
         @Override
         public Map<VKey<? extends EppResource>, EppResource> loadAll(
             Iterable<? extends VKey<? extends EppResource>> keys) {
-          return tm().doTransactionless(() -> loadAsMap(keys));
+          return tm().doTransactionless(() -> tm().load(keys));
         }
       };
 
@@ -397,7 +395,7 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable {
   public static ImmutableMap<VKey<? extends EppResource>, EppResource> loadCached(
       Iterable<VKey<? extends EppResource>> keys) {
     if (!RegistryConfig.isEppResourceCachingEnabled()) {
-      return loadAsMap(keys);
+      return tm().load(keys);
     }
     try {
       return cacheEppResources.getAll(keys);
@@ -425,17 +423,4 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable {
       throw new RuntimeException("Error loading cached EppResources", e.getCause());
     }
   }
-
-  private static ImmutableMap<VKey<? extends EppResource>, EppResource> loadAsMap(
-      Iterable<? extends VKey<? extends EppResource>> keys) {
-    return StreamSupport.stream(keys.spliterator(), false)
-        // It's possible for us to receive the same key more than once which causes
-        // the immutable map build to break with a duplicate key, so we have to ensure key
-        // uniqueness.
-        .distinct()
-        // We have to use "key -> key" here instead of the identity() function, because
-        // the latter breaks the fairly complicated generic type checking required by the
-        // caching interface.
-        .collect(toImmutableMap(key -> key, key -> tm().load(key)));
-  }
 }

core/src/main/java/google/registry/model/ofy/DatastoreTransactionManager.java

@@ -15,15 +15,18 @@
 package google.registry.model.ofy;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
+import com.google.common.base.Functions;
 import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.googlecode.objectify.Key;
 import google.registry.persistence.VKey;
 import google.registry.persistence.transaction.TransactionManager;
-import java.util.Iterator;
+import java.util.Map.Entry;
 import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.function.Supplier;
@@ -156,12 +159,15 @@ public class DatastoreTransactionManager implements TransactionManager {
   }
 
   @Override
-  public <T> ImmutableList<T> load(Iterable<VKey<T>> keys) {
+  public <T> ImmutableMap<VKey<? extends T>, T> load(Iterable<? extends VKey<? extends T>> keys) {
-    Iterator<Key<T>> iter =
+    // Keep track of the Key -> VKey mapping so we can translate them back.
-        StreamSupport.stream(keys.spliterator(), false).map(VKey::getOfyKey).iterator();
+    ImmutableMap<Key<T>, VKey<? extends T>> keyMap =
+        StreamSupport.stream(keys.spliterator(), false)
+            .distinct()
+            .collect(toImmutableMap(key -> (Key<T>) key.getOfyKey(), Functions.identity()));
 
-    // The lambda argument to keys() effectively converts Iterator -> Iterable.
+    return getOfy().load().keys(keyMap.keySet()).entrySet().stream()
-    return ImmutableList.copyOf(getOfy().load().keys(() -> iter).values());
+        .collect(ImmutableMap.toImmutableMap(entry -> keyMap.get(entry.getKey()), Entry::getValue));
   }
 
   @Override

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

@@ -15,18 +15,21 @@
 package google.registry.persistence.transaction;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.collect.ImmutableList.toImmutableList;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
+import static java.util.AbstractMap.SimpleEntry;
 import static java.util.stream.Collectors.joining;
 
 import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.FluentLogger;
 import google.registry.persistence.VKey;
 import google.registry.util.Clock;
 import java.lang.reflect.Field;
+import java.util.Map.Entry;
 import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.function.Supplier;
@@ -253,20 +256,18 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   }
 
   @Override
-  public <T> ImmutableList<T> load(Iterable<VKey<T>> keys) {
+  public <T> ImmutableMap<VKey<? extends T>, T> load(Iterable<? extends VKey<? extends T>> keys) {
     checkArgumentNotNull(keys, "keys must be specified");
     assertInTransaction();
     return StreamSupport.stream(keys.spliterator(), false)
+        // Accept duplicate keys.
+        .distinct()
         .map(
-            key -> {
+            key ->
-              T entity = getEntityManager().find(key.getKind(), key.getSqlKey());
+                new SimpleEntry<VKey<? extends T>, T>(
-              if (entity == null) {
+                    key, getEntityManager().find(key.getKind(), key.getSqlKey())))
-                throw new NoSuchElementException(
+        .filter(entry -> entry.getValue() != null)
-                    key.getKind().getName() + " with key " + key.getSqlKey() + " not found.");
+        .collect(toImmutableMap(Entry::getKey, Entry::getValue));
-              }
-              return entity;
-            })
-        .collect(toImmutableList());
   }
 
   @Override

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

@@ -16,6 +16,7 @@ package google.registry.persistence.transaction;
 
 import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import google.registry.persistence.VKey;
 import java.util.NoSuchElementException;
 import java.util.Optional;
@@ -119,7 +120,7 @@ public interface TransactionManager {
    *
    * @throws NoSuchElementException if any of the keys are not found.
    */
-  <T> ImmutableList<T> load(Iterable<VKey<T>> keys);
+  <T> ImmutableMap<VKey<? extends T>, T> load(Iterable<? extends VKey<? extends T>> keys);
 
   /** Loads all entities of the given type, returns empty if there is no such entity. */
   <T> ImmutableList<T> loadAll(Class<T> clazz);

core/src/main/java/google/registry/rdap/RdapJsonFormatter.java

@@ -342,7 +342,7 @@ public class RdapJsonFormatter {
     // Kick off the database loads of the nameservers that we will need, so it can load
     // asynchronously while we load and process the contacts.
     ImmutableSet<HostResource> loadedHosts =
-        ImmutableSet.copyOf(tm().load(domainBase.getNameservers()));
+        ImmutableSet.copyOf(tm().load(domainBase.getNameservers()).values());
     // Load the registrant and other contacts and add them to the data.
     Map<Key<ContactResource>, ContactResource> loadedContacts =
         ofy()

core/src/main/java/google/registry/tools/UniformRapidSuspensionCommand.java

@@ -149,7 +149,7 @@ final class UniformRapidSuspensionCommand extends MutatingEppToolCommand {
 
   private ImmutableSortedSet<String> getExistingNameservers(DomainBase domain) {
     ImmutableSortedSet.Builder<String> nameservers = ImmutableSortedSet.naturalOrder();
-    for (HostResource host : tm().load(domain.getNameservers())) {
+    for (HostResource host : tm().load(domain.getNameservers()).values()) {
       nameservers.add(host.getForeignKey());
     }
     return nameservers.build();

core/src/main/java/google/registry/tools/server/GenerateZoneFilesAction.java

@@ -214,7 +214,7 @@ public class GenerateZoneFilesAction implements Runnable, JsonActionRunner.JsonA
     private void emitForSubordinateHosts(DomainBase domain) {
       ImmutableSet<String> subordinateHosts = domain.getSubordinateHosts();
       if (!subordinateHosts.isEmpty()) {
-        for (HostResource unprojectedHost : tm().load(domain.getNameservers())) {
+        for (HostResource unprojectedHost : tm().load(domain.getNameservers()).values()) {
           HostResource host = loadAtPointInTime(unprojectedHost, exportTime).now();
           // A null means the host was deleted (or not created) at this time.
           if ((host != null) && subordinateHosts.contains(host.getHostName())) {
@@ -283,7 +283,7 @@ public class GenerateZoneFilesAction implements Runnable, JsonActionRunner.JsonA
       Duration dnsDefaultDsTtl) {
     StringBuilder result = new StringBuilder();
     String domainLabel = stripTld(domain.getDomainName(), domain.getTld());
-    for (HostResource nameserver : tm().load(domain.getNameservers())) {
+    for (HostResource nameserver : tm().load(domain.getNameservers()).values()) {
       result.append(
           String.format(
               NS_FORMAT,

core/src/test/java/google/registry/persistence/transaction/TransactionManagerTest.java

@@ -21,6 +21,7 @@ import static google.registry.persistence.transaction.TransactionManagerFactory.
 import static org.junit.Assert.assertThrows;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Maps;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.annotation.Entity;
 import com.googlecode.objectify.annotation.Id;
@@ -35,6 +36,7 @@ import google.registry.testing.InjectRule;
 import java.util.List;
 import java.util.NoSuchElementException;
 import java.util.Set;
+import java.util.stream.Stream;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.TestTemplate;
 import org.junit.jupiter.api.extension.RegisterExtension;
@@ -272,6 +274,40 @@ public class TransactionManagerTest {
     assertAllEntitiesNotExist(moreEntities);
   }
 
+  @TestTemplate
+  void load_multi() {
+    assertAllEntitiesNotExist(moreEntities);
+    tm().transact(() -> tm().saveAllNew(moreEntities));
+    List<VKey<TestEntity>> keys =
+        moreEntities.stream().map(TestEntity::key).collect(toImmutableList());
+    assertThat(tm().transact(() -> tm().load(keys)))
+        .isEqualTo(Maps.uniqueIndex(moreEntities, TestEntity::key));
+  }
+
+  @TestTemplate
+  void load_multiWithDuplicateKeys() {
+    assertAllEntitiesNotExist(moreEntities);
+    tm().transact(() -> tm().saveAllNew(moreEntities));
+    ImmutableList<VKey<TestEntity>> keys =
+        moreEntities.stream().map(TestEntity::key).collect(toImmutableList());
+    ImmutableList<VKey<TestEntity>> doubleKeys =
+        Stream.concat(keys.stream(), keys.stream()).collect(toImmutableList());
+    assertThat(tm().transact(() -> tm().load(doubleKeys)))
+        .isEqualTo(Maps.uniqueIndex(moreEntities, TestEntity::key));
+  }
+
+  @TestTemplate
+  void load_multiMissingKeys() {
+    assertAllEntitiesNotExist(moreEntities);
+    tm().transact(() -> tm().saveAllNew(moreEntities));
+    List<VKey<TestEntity>> keys =
+        Stream.concat(moreEntities.stream(), Stream.of(new TestEntity("dark", "matter")))
+            .map(TestEntity::key)
+            .collect(toImmutableList());
+    assertThat(tm().transact(() -> tm().load(keys)))
+        .isEqualTo(Maps.uniqueIndex(moreEntities, TestEntity::key));
+  }
+
   private static void assertEntityExists(TestEntity entity) {
     assertThat(tm().transact(() -> tm().checkExists(entity))).isTrue();
   }