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));
   }

javatests/google/registry/model/EppResourceTest.java

@@ -15,7 +15,6 @@
 package google.registry.model;
 
 import static com.google.common.truth.Truth.assertThat;
-import static google.registry.model.EppResource.CACHE_LOADER;
 import static google.registry.testing.DatastoreHelper.persistActiveContact;
 import static google.registry.testing.DatastoreHelper.persistActiveHost;
 import static google.registry.testing.DatastoreHelper.persistResource;
@@ -64,7 +63,6 @@ public class EppResourceTest extends EntityTestCase {
   }
 
   private static void setNonZeroCachingInterval() {
-    EppResource.cacheEppResources =
-        CacheBuilder.newBuilder().expireAfterWrite(1L, DAYS).build(CACHE_LOADER);
+    EppResource.setCacheForTest(CacheBuilder.newBuilder().expireAfterWrite(1L, DAYS));
   }
 }

javatests/google/registry/model/index/ForeignKeyIndexTest.java

@@ -15,7 +15,6 @@
 package google.registry.model.index;
 
 import static com.google.common.truth.Truth.assertThat;
-import static google.registry.model.index.ForeignKeyIndex.CACHE_LOADER;
 import static google.registry.model.ofy.ObjectifyService.ofy;
 import static google.registry.testing.DatastoreHelper.createTld;
 import static google.registry.testing.DatastoreHelper.deleteResource;
@@ -232,7 +231,6 @@ public class ForeignKeyIndexTest extends EntityTestCase {
   }
 
   private static void setNonZeroCachingInterval() {
-    ForeignKeyIndex.cacheForeignKeyIndexes =
-        CacheBuilder.newBuilder().expireAfterWrite(1L, DAYS).build(CACHE_LOADER);
+    ForeignKeyIndex.setCacheForTest(CacheBuilder.newBuilder().expireAfterWrite(1L, DAYS));
   }
 }

javatests/google/registry/whois/WhoisActionTest.java

@@ -15,10 +15,12 @@
 package google.registry.whois;
 
 import static com.google.common.truth.Truth.assertThat;
+import static google.registry.model.EppResourceUtils.loadByForeignKeyCached;
 import static google.registry.model.registrar.Registrar.State.ACTIVE;
 import static google.registry.model.registrar.Registrar.Type.PDT;
 import static google.registry.model.registry.Registries.getTlds;
 import static google.registry.testing.DatastoreHelper.createTlds;
+import static google.registry.testing.DatastoreHelper.persistActiveDomain;
 import static google.registry.testing.DatastoreHelper.persistResource;
 import static google.registry.testing.DatastoreHelper.persistSimpleResources;
 import static google.registry.testing.FullFieldsTestEntityHelper.makeContactResource;
@@ -27,6 +29,7 @@ import static google.registry.testing.FullFieldsTestEntityHelper.makeHostResourc
 import static google.registry.testing.FullFieldsTestEntityHelper.makeRegistrar;
 import static google.registry.testing.FullFieldsTestEntityHelper.makeRegistrarContacts;
 import static google.registry.whois.WhoisTestData.loadFile;
+import static java.util.concurrent.TimeUnit.DAYS;
 import static javax.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
 import static javax.servlet.http.HttpServletResponse.SC_NOT_FOUND;
 import static javax.servlet.http.HttpServletResponse.SC_OK;
@@ -38,7 +41,14 @@ import static org.mockito.Mockito.when;
 
 import com.google.appengine.api.datastore.DatastoreFailureException;
 import com.google.appengine.api.datastore.DatastoreTimeoutException;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.net.InetAddresses;
+import google.registry.model.EppResource;
+import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainResource;
+import google.registry.model.host.HostResource;
+import google.registry.model.index.ForeignKeyIndex;
 import google.registry.model.ofy.Ofy;
 import google.registry.model.registrar.Registrar;
 import google.registry.model.registry.Registry;
@@ -88,6 +98,10 @@ public class WhoisActionTest {
   public void setUp() throws Exception {
     createTlds("lol", "xn--q9jyb4c", "1.test");
     inject.setStaticField(Ofy.class, "clock", clock);
+
+    // Set caches with long intervals, to test caching.
+    EppResource.setCacheForTest(CacheBuilder.newBuilder().expireAfterWrite(1L, DAYS));
+    ForeignKeyIndex.setCacheForTest(CacheBuilder.newBuilder().expireAfterWrite(1L, DAYS));
   }
 
   @Test
@@ -98,23 +112,60 @@ public class WhoisActionTest {
   }
 
   @Test
-  public void testRun_domainQuery_works() throws Exception {
-    Registrar registrar = persistResource(makeRegistrar(
-        "evilregistrar", "Yes Virginia <script>", ACTIVE));
-    persistResource(makeDomainResource(
-        "cat.lol",
-        persistResource(makeContactResource("5372808-ERL", "Goblin Market", "lol@cat.lol")),
-        persistResource(makeContactResource("5372808-IRL", "Santa Claus", "BOFH@cat.lol")),
-        persistResource(makeContactResource("5372808-TRL", "The Raven", "bog@cat.lol")),
-        persistResource(makeHostResource("ns1.cat.lol", "1.2.3.4")),
-        persistResource(makeHostResource("ns2.cat.lol", "bad:f00d:cafe::15:beef")),
-        registrar));
+  public void testRun_domainQuery_works() {
+    Registrar registrar =
+        persistResource(makeRegistrar("evilregistrar", "Yes Virginia <script>", ACTIVE));
+    persistResource(
+        makeDomainResource(
+            "cat.lol",
+            persistResource(makeContactResource("5372808-ERL", "Goblin Market", "lol@cat.lol")),
+            persistResource(makeContactResource("5372808-IRL", "Santa Claus", "BOFH@cat.lol")),
+            persistResource(makeContactResource("5372808-TRL", "The Raven", "bog@cat.lol")),
+            persistResource(makeHostResource("ns1.cat.lol", "1.2.3.4")),
+            persistResource(makeHostResource("ns2.cat.lol", "bad:f00d:cafe::15:beef")),
+            registrar));
     persistSimpleResources(makeRegistrarContacts(registrar));
     newWhoisAction("domain cat.lol\r\n").run();
     assertThat(response.getStatus()).isEqualTo(200);
     assertThat(response.getPayload()).isEqualTo(loadFile("whois_action_domain.txt"));
   }
 
+  @Test
+  public void testRun_domainQuery_usesCache() {
+    Registrar registrar =
+        persistResource(makeRegistrar("evilregistrar", "Yes Virginia <script>", ACTIVE));
+    persistResource(
+        makeDomainResource(
+            "cat.lol",
+            persistResource(makeContactResource("5372808-ERL", "Goblin Market", "lol@cat.lol")),
+            persistResource(makeContactResource("5372808-IRL", "Santa Claus", "BOFH@cat.lol")),
+            persistResource(makeContactResource("5372808-TRL", "The Raven", "bog@cat.lol")),
+            persistResource(makeHostResource("ns1.cat.lol", "1.2.3.4")),
+            persistResource(makeHostResource("ns2.cat.lol", "bad:f00d:cafe::15:beef")),
+            registrar));
+    persistSimpleResources(makeRegistrarContacts(registrar));
+    // Populate the cache for both the domain and contact.
+    DomainResource domain =
+        loadByForeignKeyCached(DomainResource.class, "cat.lol", clock.nowUtc());
+    ContactResource contact =
+        loadByForeignKeyCached(ContactResource.class, "5372808-ERL", clock.nowUtc());
+    // Make a change to the domain and contact that won't be seen because the cache will be hit.
+    persistResource(domain.asBuilder().setDeletionTime(clock.nowUtc().minusDays(1)).build());
+    persistResource(
+        contact
+            .asBuilder()
+            .setInternationalizedPostalInfo(
+                contact
+                    .getInternationalizedPostalInfo()
+                    .asBuilder()
+                    .setOrg("Two by Two, Hands Blue Inc.")
+                    .build())
+            .build());
+    newWhoisAction("domain cat.lol\r\n").run();
+    assertThat(response.getStatus()).isEqualTo(200);
+    assertThat(response.getPayload()).isEqualTo(loadFile("whois_action_domain.txt"));
+  }
+
   @Test
   public void testRun_idnDomain_works() throws Exception {
     Registrar registrar = persistResource(makeRegistrar(
@@ -152,7 +203,18 @@ public class WhoisActionTest {
   }
 
   @Test
-  public void testRun_domainNotFound_returns200OkAndPlainTextResponse() throws Exception {
+  public void testRun_domainNotFound_returns200OkAndPlainTextResponse() {
+    newWhoisAction("domain cat.lol\r\n").run();
+    assertThat(response.getStatus()).isEqualTo(200);
+    assertThat(response.getPayload()).isEqualTo(loadFile("whois_action_domain_not_found.txt"));
+  }
+
+  @Test
+  public void testRun_domainNotFound_usesCache() {
+    // Populate the cache with the nonexistence of this domain.
+    assertThat(loadByForeignKeyCached(DomainResource.class, "cat.lol", clock.nowUtc())).isNull();
+    // Add a new valid cat.lol domain that won't be found because the cache will be hit instead.
+    persistActiveDomain("cat.lol");
     newWhoisAction("domain cat.lol\r\n").run();
     assertThat(response.getStatus()).isEqualTo(200);
     assertThat(response.getPayload()).isEqualTo(loadFile("whois_action_domain_not_found.txt"));
@@ -260,7 +322,7 @@ public class WhoisActionTest {
   }
 
   @Test
-  public void testRun_idnNameserver_works() throws Exception {
+  public void testRun_idnNameserver_works() {
     persistResource(makeHostResource("ns1.cat.みんな", "1.2.3.4"));
     newWhoisAction("nameserver ns1.cat.みんな\r\n").run();
     assertThat(response.getStatus()).isEqualTo(200);
@@ -268,6 +330,23 @@ public class WhoisActionTest {
     assertThat(response.getPayload()).contains("1.2.3.4");
   }
 
+  @Test
+  public void testRun_nameserver_usesCache() {
+    persistResource(makeHostResource("ns1.cat.xn--q9jyb4c", "1.2.3.4"));
+    // Populate the cache.
+    HostResource host =
+        loadByForeignKeyCached(HostResource.class, "ns1.cat.xn--q9jyb4c", clock.nowUtc());
+    // Make a change to the persisted host that won't be seen because the cache will be hit.
+    persistResource(
+        host.asBuilder()
+            .setInetAddresses(ImmutableSet.of(InetAddresses.forString("8.8.8.8")))
+            .build());
+    newWhoisAction("nameserver ns1.cat.みんな\r\n").run();
+    assertThat(response.getStatus()).isEqualTo(200);
+    assertThat(response.getPayload()).contains("ns1.cat.xn--q9jyb4c");
+    assertThat(response.getPayload()).contains("1.2.3.4");
+  }
+
   @Test
   public void testRun_punycodeNameserver_works() throws Exception {
     persistResource(makeHostResource("ns1.cat.みんな", "1.2.3.4"));