Update the Registries cache to leverage/populate the Registry cache (#382)

* Update the Registries cache to leverage/populate the Registry cache

This is accomplished by also providing a loadAll() method on the Registry cache
that can be used to load an entire batch of Registry objects at once.

This improves efficiency, because now, any operation on Registries that loads
all the Registry entities (getTlds(), getTldsOfType(), and getTldEntities()), is
plumbed through the Registry cache, therefore loading it from that cache if it
exists and only hitting the DB if not. If not, this populates the Registry cache
upon loading, so that subsequent calls to Registry.get() will now hit the cache
instead of the DB.

To give a concrete example, the following code:

    for (String tld : Registries.getTlds()) {
      // ...
      doSomethingWith(Registry.get(tld));
      // ...
    }

is now much more efficient, because the initial call to Registries.getTlds()
populates all the entities in cache, and the subsequent individual calls to
Registry.get(tld) now retrieve them from the cache. Prior to this change,
Registries.getTlds() did not populate the Registry cache, and each subsequent
Registry.get(tld) had the potential to trigger an individual round-trip to the
DB, which is obviously bad for performance.

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

@@ -19,6 +19,7 @@ import static com.google.common.base.Predicates.equalTo;
 import static com.google.common.base.Predicates.in;
 import static com.google.common.base.Predicates.not;
 import static com.google.common.base.Strings.emptyToNull;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.collect.Maps.filterValues;
 import static google.registry.model.CacheUtils.memoizeWithShortExpiration;
@@ -31,10 +32,12 @@ import com.google.common.base.Joiner;
 import com.google.common.base.Supplier;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Maps;
 import com.google.common.collect.Streams;
 import com.google.common.net.InternetDomainName;
 import com.googlecode.objectify.Key;
 import google.registry.model.registry.Registry.TldType;
+import java.util.Map;
 import java.util.Optional;
 
 /** Utilities for finding and listing {@link Registry} entities. */
@@ -54,16 +57,21 @@ public final class Registries {
   private static Supplier<ImmutableMap<String, TldType>> createFreshCache() {
     return memoizeWithShortExpiration(
         () ->
-            tm()
+            tm().doTransactionless(
-                .doTransactionless(
                     () -> {
-                      ImmutableMap.Builder<String, TldType> builder =
+                      ImmutableSet<String> tlds =
-                          new ImmutableMap.Builder<>();
+                          ofy()
-                      for (Registry registry :
+                              .load()
-                          ofy().load().type(Registry.class).ancestor(getCrossTldKey())) {
+                              .type(Registry.class)
-                        builder.put(registry.getTldStr(), registry.getTldType());
+                              .ancestor(getCrossTldKey())
-                      }
+                              .keys()
-                      return builder.build();
+                              .list()
+                              .stream()
+                              .map(Key::getName)
+                              .collect(toImmutableSet());
+                      return Registry.getAll(tlds).stream()
+                          .map(e -> Maps.immutableEntry(e.getTldStr(), e.getTldType()))
+                          .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
                     }));
   }
 
@@ -83,11 +91,7 @@ public final class Registries {
 
   /** Returns the Registry entities themselves of the given type loaded fresh from Datastore. */
   public static ImmutableSet<Registry> getTldEntitiesOfType(TldType type) {
-    ImmutableSet<Key<Registry>> keys =
+    return Registry.getAll(filterValues(cache.get(), equalTo(type)).keySet());
-        filterValues(cache.get(), equalTo(type)).keySet().stream()
-            .map(tld -> Key.create(getCrossTldKey(), Registry.class, tld))
-            .collect(toImmutableSet());
-    return ImmutableSet.copyOf(tm().doTransactionless(() -> ofy().load().keys(keys).values()));
   }
 
   /** Pass-through check that the specified TLD exists, otherwise throw an IAE. */

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

@@ -18,6 +18,9 @@ import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Predicates.equalTo;
 import static com.google.common.base.Predicates.not;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static com.google.common.collect.Maps.toMap;
 import static google.registry.config.RegistryConfig.getSingletonCacheRefreshDuration;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.ObjectifyService.ofy;
@@ -30,9 +33,11 @@ import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static org.joda.money.CurrencyUnit.USD;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Joiner;
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.Iterables;
@@ -58,8 +63,10 @@ import google.registry.model.domain.fee.Fee;
 import google.registry.model.registry.label.PremiumList;
 import google.registry.model.registry.label.ReservedList;
 import google.registry.util.Idn;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.ExecutionException;
 import java.util.function.Predicate;
 import java.util.regex.Pattern;
 import javax.annotation.Nullable;
@@ -201,6 +208,25 @@ public class Registry extends ImmutableObject implements Buildable {
     return registry;
   }
 
+  /** Returns the registry entities for the given TLD strings, throwing if any don't exist. */
+  static ImmutableSet<Registry> getAll(Set<String> tlds) {
+    try {
+      ImmutableMap<String, Optional<Registry>> registries = CACHE.getAll(tlds);
+      ImmutableSet<String> missingRegistries =
+          registries.entrySet().stream()
+              .filter(e -> !e.getValue().isPresent())
+              .map(Map.Entry::getKey)
+              .collect(toImmutableSet());
+      if (missingRegistries.isEmpty()) {
+        return registries.values().stream().map(Optional::get).collect(toImmutableSet());
+      } else {
+        throw new RegistryNotFoundException(missingRegistries);
+      }
+    } catch (ExecutionException e) {
+      throw new RuntimeException("Unexpected error retrieving TLDs " + tlds, e);
+    }
+  }
+
   /**
    * Invalidates the cache entry.
    *
@@ -220,16 +246,31 @@ public class Registry extends ImmutableObject implements Buildable {
               new CacheLoader<String, Optional<Registry>>() {
                 @Override
                 public Optional<Registry> load(final String tld) {
-                  // Enter a transactionless context briefly; we don't want to enroll every TLD in a
+                  // Enter a transaction-less context briefly; we don't want to enroll every TLD in
-                  // transaction that might be wrapping this call.
+                  // a transaction that might be wrapping this call.
                   return Optional.ofNullable(
-                      tm()
+                      tm().doTransactionless(
-                          .doTransactionless(
+                              () ->
-                              () -> ofy()
+                                  ofy()
                                       .load()
                                       .key(Key.create(getCrossTldKey(), Registry.class, tld))
                                       .now()));
                 }
+
+                @Override
+                public Map<String, Optional<Registry>> loadAll(Iterable<? extends String> tlds) {
+                  ImmutableMap<String, Key<Registry>> keysMap =
+                      toMap(
+                          ImmutableSet.copyOf(tlds),
+                          tld -> Key.create(getCrossTldKey(), Registry.class, tld));
+                  Map<Key<Registry>, Registry> entities =
+                      tm().doTransactionless(() -> ofy().load().keys(keysMap.values()));
+                  return keysMap.entrySet().stream()
+                      .collect(
+                          toImmutableMap(
+                              Map.Entry::getKey,
+                              e -> Optional.ofNullable(entities.getOrDefault(e.getValue(), null))));
+                }
               });
 
   /**
@@ -883,10 +924,14 @@ public class Registry extends ImmutableObject implements Buildable {
     }
   }
 
-  /** Exception to throw when no Registry is found for a given tld. */
+  /** Exception to throw when no Registry entity is found for given TLD string(s). */
   public static class RegistryNotFoundException extends RuntimeException {
+    RegistryNotFoundException(ImmutableSet<String> tlds) {
+      super("No registry object(s) found for " + Joiner.on(", ").join(tlds));
+    }
+
     RegistryNotFoundException(String tld) {
-      super("No registry object found for " + tld);
+      this(ImmutableSet.of(tld));
     }
   }
 }

core/src/test/java/google/registry/export/ExportReservedTermsActionTest.java

@@ -48,10 +48,7 @@ import org.junit.runners.JUnit4;
 @RunWith(JUnit4.class)
 public class ExportReservedTermsActionTest {
 
-  @Rule
+  @Rule public final AppEngineRule appEngine = AppEngineRule.builder().withDatastore().build();
-  public final AppEngineRule appEngine = AppEngineRule.builder()
-      .withDatastore()
-      .build();
 
   private final DriveConnection driveConnection = mock(DriveConnection.class);
   private final Response response = mock(Response.class);
@@ -133,6 +130,6 @@ public class ExportReservedTermsActionTest {
     assertThat(thrown)
         .hasCauseThat()
         .hasMessageThat()
-        .isEqualTo("No registry object found for fakeTld");
+        .isEqualTo("No registry object(s) found for fakeTld");
   }
 }

core/src/test/java/google/registry/model/registry/RegistryTest.java

@@ -51,6 +51,7 @@ import org.junit.Test;
 
 /** Unit tests for {@link Registry}. */
 public class RegistryTest extends EntityTestCase {
+
   Registry registry;
 
   @Before
@@ -146,6 +147,19 @@ public class RegistryTest extends EntityTestCase {
     assertThat(registry.getReservedLists()).isEmpty();
   }
 
+  @Test
+  public void testGetAll() {
+    createTld("foo");
+    assertThat(Registry.getAll(ImmutableSet.of("foo", "tld")))
+        .containsExactlyElementsIn(
+            ofy()
+                .load()
+                .keys(
+                    Key.create(getCrossTldKey(), Registry.class, "foo"),
+                    Key.create(getCrossTldKey(), Registry.class, "tld"))
+                .values());
+  }
+
   @Test
   public void testSetReservedLists() {
     ReservedList rl5 = persistReservedList(