Cache domains, contacts, and hosts in WHOIS queries

This should prevent having issues with hot key paths on entities that
experience a heavy WHOIS volume (e.g. contacts that registrars reuse on
many domains).

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

java/google/registry/model/EppResource.java

@@ -341,12 +341,17 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable {
    * directly should of course never use the cache.
    */
   @NonFinalForTesting
-  static LoadingCache<Key<? extends EppResource>, EppResource> cacheEppResources =
+  private static LoadingCache<Key<? extends EppResource>, EppResource> cacheEppResources =
       CacheBuilder.newBuilder()
           .expireAfterWrite(getEppResourceCachingDuration().getMillis(), MILLISECONDS)
           .maximumSize(getEppResourceMaxCachedEntries())
           .build(CACHE_LOADER);
 
+  @VisibleForTesting
+  public static void setCacheForTest(CacheBuilder<Object, Object> cacheBuilder) {
+    cacheEppResources = cacheBuilder.build(CACHE_LOADER);
+  }
+
   private static ImmutableMap<Key<? extends EppResource>, EppResource> loadMultiple(
       Iterable<? extends Key<? extends EppResource>> keys) {
     // This cast is safe because, in Objectify, Key<? extends EppResource> can also be
@@ -368,7 +373,7 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable {
   }
 
   /**
-   * Loads a given EppResource by its key using the cache (if enabled).
+   * Loads the given EppResources by their keys using the cache (if enabled).
    *
    * <p>Don't use this unless you really need it for performance reasons, and be sure that you are
    * OK with the trade-offs in loss of transactional consistency.
@@ -384,4 +389,24 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable {
       throw new RuntimeException("Error loading cached EppResources", e.getCause());
     }
   }
+
+  /**
+   * Loads a given EppResource by its key using the cache (if enabled).
+   *
+   * <p>Don't use this unless you really need it for performance reasons, and be sure that you are
+   * OK with the trade-offs in loss of transactional consistency.
+   */
+  public static <T extends EppResource> T loadCached(Key<T> key) {
+    if (!RegistryConfig.isEppResourceCachingEnabled()) {
+      return ofy().load().key(key).now();
+    }
+    try {
+      // Safe to cast because loading a Key<T> returns an entity of type T.
+      @SuppressWarnings("unchecked")
+      T resource = (T) cacheEppResources.get(key);
+      return resource;
+    } catch (ExecutionException e) {
+      throw new RuntimeException("Error loading cached EppResources", e.getCause());
+    }
+  }
 }

java/google/registry/model/EppResourceUtils.java

@@ -21,10 +21,12 @@ import static google.registry.util.DateTimeUtils.isAtOrAfter;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 import static google.registry.util.DateTimeUtils.latestOf;
 
+import com.google.common.collect.ImmutableList;
 import com.googlecode.objectify.Key;
 import com.googlecode.objectify.Result;
 import com.googlecode.objectify.cmd.Query;
 import com.googlecode.objectify.util.ResultNow;
+import google.registry.config.RegistryConfig;
 import google.registry.model.EppResource.Builder;
 import google.registry.model.EppResource.BuilderWithTransferData;
 import google.registry.model.EppResource.ForeignKeyedEppResource;
@@ -80,10 +82,10 @@ public final class EppResourceUtils {
    * it may have various expirable conditions and status values that might implicitly change its
    * state as time progresses even if it has not been updated in Datastore. Rather, the resource
    * must be combined with a timestamp to view its current state. We use a global last updated
-   * timestamp on the entire entity group (which is essentially free since all writes to the entity
-   * group must be serialized anyways) to guarantee monotonically increasing write times, so
-   * forwarding our projected time to the greater of "now", and this update timestamp guarantees
-   * that we're not projecting into the past.
+   * timestamp on the resource's entity group (which is essentially free since all writes to the
+   * entity group must be serialized anyways) to guarantee monotonically increasing write times, and
+   * forward our projected time to the greater of this timestamp or "now". This guarantees that
+   * we're not projecting into the past.
    *
    * @param clazz the resource type to load
    * @param foreignKey id to match
@@ -92,16 +94,58 @@ public final class EppResourceUtils {
   @Nullable
   public static <T extends EppResource> T loadByForeignKey(
       Class<T> clazz, String foreignKey, DateTime now) {
+    return loadByForeignKeyHelper(clazz, foreignKey, now, false);
+  }
+
+  /**
+   * Loads the last created version of an {@link EppResource} from Datastore by foreign key, using a
+   * cache.
+   *
+   * <p>Returns null if no resource with this foreign key was ever created, or if the most recently
+   * created resource was deleted before time "now".
+   *
+   * <p>Loading an {@link EppResource} by itself is not sufficient to know its current state since
+   * it may have various expirable conditions and status values that might implicitly change its
+   * state as time progresses even if it has not been updated in Datastore. Rather, the resource
+   * must be combined with a timestamp to view its current state. We use a global last updated
+   * timestamp on the resource's entity group (which is essentially free since all writes to the
+   * entity group must be serialized anyways) to guarantee monotonically increasing write times, and
+   * forward our projected time to the greater of this timestamp or "now". This guarantees that
+   * we're not projecting into the past.
+   *
+   * <p>Do not call this cached version for anything that needs transactional consistency. It should
+   * only be used when it's OK if the data is potentially being out of date, e.g. WHOIS.
+   *
+   * @param clazz the resource type to load
+   * @param foreignKey id to match
+   * @param now the current logical time to project resources at
+   */
+  @Nullable
+  public static <T extends EppResource> T loadByForeignKeyCached(
+      Class<T> clazz, String foreignKey, DateTime now) {
+    return loadByForeignKeyHelper(
+        clazz, foreignKey, now, RegistryConfig.isEppResourceCachingEnabled());
+  }
+
+  @Nullable
+  private static <T extends EppResource> T loadByForeignKeyHelper(
+      Class<T> clazz, String foreignKey, DateTime now, boolean useCache) {
     checkArgument(
         ForeignKeyedEppResource.class.isAssignableFrom(clazz),
         "loadByForeignKey may only be called for foreign keyed EPP resources");
     ForeignKeyIndex<T> fki =
-        ofy().load().type(ForeignKeyIndex.mapToFkiClass(clazz)).id(foreignKey).now();
+        useCache
+            ? ForeignKeyIndex.loadCached(clazz, ImmutableList.of(foreignKey), now)
+                .getOrDefault(foreignKey, null)
+            : ofy().load().type(ForeignKeyIndex.mapToFkiClass(clazz)).id(foreignKey).now();
     // The value of fki.getResourceKey() might be null for hard-deleted prober data.
     if (fki == null || isAtOrAfter(now, fki.getDeletionTime()) || fki.getResourceKey() == null) {
       return null;
     }
-    T resource = ofy().load().key(fki.getResourceKey()).now();
+    T resource =
+        useCache
+            ? EppResource.loadCached(fki.getResourceKey())
+            : ofy().load().key(fki.getResourceKey()).now();
     if (resource == null || isAtOrAfter(now, resource.getDeletionTime())) {
       return null;
     }

java/google/registry/model/index/ForeignKeyIndex.java

@@ -22,6 +22,7 @@ import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.util.TypeUtils.instantiate;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
@@ -218,13 +219,18 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
    * given IDs (blah) don't exist."
    */
   @NonFinalForTesting
-  static LoadingCache<Key<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>>
+  private static LoadingCache<Key<ForeignKeyIndex<?>>, Optional<ForeignKeyIndex<?>>>
       cacheForeignKeyIndexes =
           CacheBuilder.newBuilder()
               .expireAfterWrite(getEppResourceCachingDuration().getMillis(), MILLISECONDS)
               .maximumSize(getEppResourceMaxCachedEntries())
               .build(CACHE_LOADER);
 
+  @VisibleForTesting
+  public static void setCacheForTest(CacheBuilder<Object, Object> cacheBuilder) {
+    cacheForeignKeyIndexes = cacheBuilder.build(CACHE_LOADER);
+  }
+
   /**
    * Load a list of {@link ForeignKeyIndex} instances by class and id strings that are active at or
    * after the specified moment in time, using the cache if enabled.

java/google/registry/whois/DomainLookupCommand.java

@@ -14,7 +14,7 @@
 
 package google.registry.whois;
 
-import static google.registry.model.EppResourceUtils.loadByForeignKey;
+import static google.registry.model.EppResourceUtils.loadByForeignKeyCached;
 
 import com.google.common.net.InternetDomainName;
 import google.registry.model.domain.DomainResource;
@@ -31,7 +31,7 @@ public class DomainLookupCommand extends DomainOrHostLookupCommand {
   @Override
   protected Optional<WhoisResponse> getResponse(InternetDomainName domainName, DateTime now) {
     final DomainResource domainResource =
-        loadByForeignKey(DomainResource.class, domainName.toString(), now);
+        loadByForeignKeyCached(DomainResource.class, domainName.toString(), now);
     return Optional.ofNullable(
         domainResource == null ? null : new DomainWhoisResponse(domainResource, now));
   }

java/google/registry/whois/DomainWhoisResponse.java

@@ -16,13 +16,13 @@ package google.registry.whois;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
-import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.util.CollectionUtils.isNullOrEmpty;
 import static google.registry.xml.UtcDateTimeAdapter.getFormattedString;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.googlecode.objectify.Key;
+import google.registry.model.EppResource;
 import google.registry.model.contact.ContactPhoneNumber;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.contact.PostalInfo;
@@ -134,7 +134,7 @@ final class DomainWhoisResponse extends WhoisResponseImpl {
       // If we refer to a contact that doesn't exist, that's a bug. It means referential integrity
       // has somehow been broken. We skip the rest of this contact, but log it to hopefully bring it
       // someone's attention.
-      ContactResource contactResource = ofy().load().key(contact).now();
+      ContactResource contactResource = EppResource.loadCached(contact);
       if (contactResource == null) {
         logger.severefmt(
             "(BUG) Broken reference found from domain %s to contact %s",

java/google/registry/whois/NameserverLookupByHostCommand.java

@@ -14,7 +14,7 @@
 
 package google.registry.whois;
 
-import static google.registry.model.EppResourceUtils.loadByForeignKey;
+import static google.registry.model.EppResourceUtils.loadByForeignKeyCached;
 
 import com.google.common.net.InternetDomainName;
 import google.registry.model.host.HostResource;
@@ -31,7 +31,7 @@ public class NameserverLookupByHostCommand extends DomainOrHostLookupCommand {
   @Override
   protected Optional<WhoisResponse> getResponse(InternetDomainName hostName, DateTime now) {
     final HostResource hostResource =
-        loadByForeignKey(HostResource.class, hostName.toString(), now);
+        loadByForeignKeyCached(HostResource.class, hostName.toString(), now);
     return Optional.ofNullable(
         hostResource == null ? null : new NameserverWhoisResponse(hostResource, now));
   }