From 6dd6ebce75e337ce3b7d17a1a24f9dff5ee1c0be Mon Sep 17 00:00:00 2001
From: Weimin Yu <weiminyu@google.com>
Date: Tue, 15 Mar 2022 22:02:05 -0400
Subject: [PATCH] Improve cache loading in Registries.java (#1558)

* Improve cache loading in Registries.java

The loader for the TLD cache in Registries.java unnecessarily reads from
another cache when running with SQL, potentially triggering additional
database access. This code runs in the whois query path, and contributes
to the high latency in sandbox.
---
 .../google/registry/model/tld/Registries.java | 45 ++++++++++++-------
 1 file changed, 28 insertions(+), 17 deletions(-)

diff --git a/core/src/main/java/google/registry/model/tld/Registries.java b/core/src/main/java/google/registry/model/tld/Registries.java
index 98fbe6f9d..40cadae87 100644
--- a/core/src/main/java/google/registry/model/tld/Registries.java
+++ b/core/src/main/java/google/registry/model/tld/Registries.java
@@ -24,6 +24,7 @@ import static com.google.common.collect.Maps.filterValues;
 import static google.registry.model.CacheUtils.memoizeWithShortExpiration;
 import static google.registry.model.common.EntityGroupRoot.getCrossTldKey;
 import static google.registry.model.ofy.ObjectifyService.auditedOfy;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.CollectionUtils.entriesToImmutableMap;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
@@ -38,6 +39,8 @@ import com.google.common.net.InternetDomainName;
 import com.googlecode.objectify.Key;
 import google.registry.model.tld.Registry.TldType;
 import java.util.Optional;
+import java.util.stream.Stream;
+import javax.persistence.EntityManager;
 
 /** Utilities for finding and listing {@link Registry} entities. */
 public final class Registries {
@@ -58,23 +61,31 @@ public final class Registries {
         () ->
             tm().doTransactionless(
                     () -> {
-                      ImmutableSet<String> tlds =
-                          tm().isOfy()
-                              ? auditedOfy()
-                                  .load()
-                                  .type(Registry.class)
-                                  .ancestor(getCrossTldKey())
-                                  .keys()
-                                  .list()
-                                  .stream()
-                                  .map(Key::getName)
-                                  .collect(toImmutableSet())
-                              : tm().loadAllOf(Registry.class).stream()
-                                  .map(Registry::getTldStr)
-                                  .collect(toImmutableSet());
-                      return Registry.getAll(tlds).stream()
-                          .map(e -> Maps.immutableEntry(e.getTldStr(), e.getTldType()))
-                          .collect(entriesToImmutableMap());
+                      if (tm().isOfy()) {
+                        ImmutableSet<String> tlds =
+                            auditedOfy()
+                                .load()
+                                .type(Registry.class)
+                                .ancestor(getCrossTldKey())
+                                .keys()
+                                .list()
+                                .stream()
+                                .map(Key::getName)
+                                .collect(toImmutableSet());
+                        return Registry.getAll(tlds).stream()
+                            .map(e -> Maps.immutableEntry(e.getTldStr(), e.getTldType()))
+                            .collect(entriesToImmutableMap());
+                      } else {
+                        EntityManager entityManager = jpaTm().getEntityManager();
+                        Stream<?> resultStream =
+                            entityManager
+                                .createQuery("SELECT tldStrId, tldType FROM Tld")
+                                .getResultStream();
+                        return resultStream
+                            .map(e -> ((Object[]) e))
+                            .map(e -> Maps.immutableEntry((String) e[0], ((TldType) e[1])))
+                            .collect(entriesToImmutableMap());
+                      }
                     }));
   }