Use the replica jpaTm in FKI and EppResource cache methods (#1503)

The cached methods are only used in situations where we don't really
care about being 100% synchronously up to date (e.g. whois), and they're
not used frequently anyway, so it's safe to use the replica in these
locations.

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

@@ -21,6 +21,7 @@ import static com.google.common.collect.Sets.union;
 import static google.registry.config.RegistryConfig.getEppResourceCachingDuration;
 import static google.registry.config.RegistryConfig.getEppResourceMaxCachedEntries;
 import static google.registry.persistence.transaction.TransactionManagerFactory.ofyTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.replicaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.CollectionUtils.nullToEmpty;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
@@ -380,13 +381,13 @@ public abstract class EppResource extends BackupGroupRoot implements Buildable {
 
         @Override
         public EppResource load(VKey<? extends EppResource> key) {
-          return tm().doTransactionless(() -> tm().loadByKey(key));
+          return replicaTm().doTransactionless(() -> replicaTm().loadByKey(key));
         }
 
         @Override
         public Map<VKey<? extends EppResource>, EppResource> loadAll(
             Iterable<? extends VKey<? extends EppResource>> keys) {
-          return tm().doTransactionless(() -> tm().loadByKeys(keys));
+          return replicaTm().doTransactionless(() -> replicaTm().loadByKeys(keys));
         }
       };
 

core/src/main/java/google/registry/model/index/ForeignKeyIndex.java

@@ -21,6 +21,7 @@ import static google.registry.config.RegistryConfig.getEppResourceCachingDuratio
 import static google.registry.config.RegistryConfig.getEppResourceMaxCachedEntries;
 import static google.registry.model.ofy.ObjectifyService.auditedOfy;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.replicaJpaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.CollectionUtils.entriesToImmutableMap;
 import static google.registry.util.TypeUtils.instantiate;
@@ -51,6 +52,7 @@ import google.registry.model.host.HostResource;
 import google.registry.model.replay.DatastoreOnlyEntity;
 import google.registry.persistence.VKey;
 import google.registry.persistence.transaction.CriteriaQueryBuilder;
+import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.util.NonFinalForTesting;
 import java.util.Collection;
 import java.util.Comparator;
@@ -198,7 +200,7 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
    */
   public static <E extends EppResource> ImmutableMap<String, ForeignKeyIndex<E>> load(
       Class<E> clazz, Collection<String> foreignKeys, final DateTime now) {
-    return loadIndexesFromStore(clazz, foreignKeys, true).entrySet().stream()
+    return loadIndexesFromStore(clazz, foreignKeys, true, false).entrySet().stream()
         .filter(e -> now.isBefore(e.getValue().getDeletionTime()))
         .collect(entriesToImmutableMap());
   }
@@ -217,7 +219,10 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
    */
   private static <E extends EppResource>
       ImmutableMap<String, ForeignKeyIndex<E>> loadIndexesFromStore(
-          Class<E> clazz, Collection<String> foreignKeys, boolean inTransaction) {
+          Class<E> clazz,
+          Collection<String> foreignKeys,
+          boolean inTransaction,
+          boolean useReplicaJpaTm) {
     if (tm().isOfy()) {
       Class<ForeignKeyIndex<E>> fkiClass = mapToFkiClass(clazz);
       return ImmutableMap.copyOf(
@@ -226,17 +231,18 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
               : tm().doTransactionless(() -> auditedOfy().load().type(fkiClass).ids(foreignKeys)));
     } else {
       String property = RESOURCE_CLASS_TO_FKI_PROPERTY.get(clazz);
+      JpaTransactionManager jpaTmToUse = useReplicaJpaTm ? replicaJpaTm() : jpaTm();
       ImmutableList<ForeignKeyIndex<E>> indexes =
-          tm().transact(
+          jpaTmToUse.transact(
-                  () ->
+              () ->
-                      jpaTm()
+                  jpaTmToUse
-                          .criteriaQuery(
+                      .criteriaQuery(
-                              CriteriaQueryBuilder.create(clazz)
+                          CriteriaQueryBuilder.create(clazz)
-                                  .whereFieldIsIn(property, foreignKeys)
+                              .whereFieldIsIn(property, foreignKeys)
-                                  .build())
+                              .build())
-                          .getResultStream()
+                      .getResultStream()
-                          .map(e -> ForeignKeyIndex.create(e, e.getDeletionTime()))
+                      .map(e -> ForeignKeyIndex.create(e, e.getDeletionTime()))
-                          .collect(toImmutableList()));
+                      .collect(toImmutableList()));
       // We need to find and return the entities with the maximum deletionTime for each foreign key.
       return Multimaps.index(indexes, ForeignKeyIndex::getForeignKey).asMap().entrySet().stream()
           .map(
@@ -260,7 +266,8 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
               loadIndexesFromStore(
                       RESOURCE_CLASS_TO_FKI_CLASS.inverse().get(key.getKind()),
                       ImmutableSet.of(foreignKey),
-                      false)
+                      false,
+                      true)
                   .get(foreignKey));
         }
 
@@ -276,7 +283,7 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
               Streams.stream(keys).map(v -> v.getSqlKey().toString()).collect(toImmutableSet());
           ImmutableSet<VKey<ForeignKeyIndex<?>>> typedKeys = ImmutableSet.copyOf(keys);
           ImmutableMap<String, ? extends ForeignKeyIndex<? extends EppResource>> existingFkis =
-              loadIndexesFromStore(resourceClass, foreignKeys, false);
+              loadIndexesFromStore(resourceClass, foreignKeys, false, true);
           // ofy omits keys that don't have values in Datastore, so re-add them in
           // here with Optional.empty() values.
           return Maps.asMap(
@@ -336,7 +343,7 @@ public abstract class ForeignKeyIndex<E extends EppResource> extends BackupGroup
     // Safe to cast VKey<FKI<E>> to VKey<FKI<?>>
     @SuppressWarnings("unchecked")
     ImmutableList<VKey<ForeignKeyIndex<?>>> fkiVKeys =
-        Streams.stream(foreignKeys)
+        foreignKeys.stream()
             .map(fk -> (VKey<ForeignKeyIndex<?>>) VKey.create(fkiClass, fk))
             .collect(toImmutableList());
     try {

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

@@ -14,8 +14,8 @@
 
 package google.registry.persistence.transaction;
 
-import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
+import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 import static org.joda.time.DateTimeZone.UTC;
 
 import com.google.appengine.api.utils.SystemProperty;
@@ -47,6 +47,10 @@ public final class TransactionManagerFactory {
   private static Supplier<JpaTransactionManager> jpaTm =
       Suppliers.memoize(TransactionManagerFactory::createJpaTransactionManager);
 
+  @NonFinalForTesting
+  private static Supplier<JpaTransactionManager> replicaJpaTm =
+      Suppliers.memoize(TransactionManagerFactory::createReplicaJpaTransactionManager);
+
   private static boolean onBeam = false;
 
   private TransactionManagerFactory() {}
@@ -61,6 +65,14 @@ public final class TransactionManagerFactory {
     }
   }
 
+  private static JpaTransactionManager createReplicaJpaTransactionManager() {
+    if (isInAppEngine()) {
+      return DaggerPersistenceComponent.create().readOnlyReplicaJpaTransactionManager();
+    } else {
+      return DummyJpaTransactionManager.create();
+    }
+  }
+
   private static DatastoreTransactionManager createTransactionManager() {
     return new DatastoreTransactionManager(null);
   }
@@ -108,6 +120,21 @@ public final class TransactionManagerFactory {
     return jpaTm.get();
   }
 
+  /** Returns a read-only {@link JpaTransactionManager} instance if configured. */
+  public static JpaTransactionManager replicaJpaTm() {
+    return replicaJpaTm.get();
+  }
+
+  /**
+   * Returns a {@link TransactionManager} that uses a replica database if one exists.
+   *
+   * <p>In Datastore mode, this is unchanged from the regular transaction manager. In SQL mode,
+   * however, this will be a reference to the read-only replica database if one is configured.
+   */
+  public static TransactionManager replicaTm() {
+    return tm().isOfy() ? tm() : replicaJpaTm();
+  }
+
   /** Returns {@link DatastoreTransactionManager} instance. */
   @VisibleForTesting
   public static DatastoreTransactionManager ofyTm() {
@@ -116,7 +143,7 @@ public final class TransactionManagerFactory {
 
   /** Sets the return of {@link #jpaTm()} to the given instance of {@link JpaTransactionManager}. */
   public static void setJpaTm(Supplier<JpaTransactionManager> jpaTmSupplier) {
-    checkNotNull(jpaTmSupplier, "jpaTmSupplier");
+    checkArgumentNotNull(jpaTmSupplier, "jpaTmSupplier");
     checkState(
         RegistryEnvironment.get().equals(RegistryEnvironment.UNITTEST)
             || RegistryToolEnvironment.get() != null,
@@ -124,13 +151,23 @@ public final class TransactionManagerFactory {
     jpaTm = Suppliers.memoize(jpaTmSupplier::get);
   }
 
+  /** Sets the value of {@link #replicaJpaTm()} to the given {@link JpaTransactionManager}. */
+  public static void setReplicaJpaTm(Supplier<JpaTransactionManager> replicaJpaTmSupplier) {
+    checkArgumentNotNull(replicaJpaTmSupplier, "replicaJpaTmSupplier");
+    checkState(
+        RegistryEnvironment.get().equals(RegistryEnvironment.UNITTEST)
+            || RegistryToolEnvironment.get() != null,
+        "setReplicaJpaTm() should only be called by tools and tests.");
+    replicaJpaTm = Suppliers.memoize(replicaJpaTmSupplier::get);
+  }
+
   /**
    * Makes {@link #jpaTm()} return the {@link JpaTransactionManager} instance provided by {@code
    * jpaTmSupplier} from now on. This method should only be called by an implementor of {@link
    * org.apache.beam.sdk.harness.JvmInitializer}.
    */
   public static void setJpaTmOnBeamWorker(Supplier<JpaTransactionManager> jpaTmSupplier) {
-    checkNotNull(jpaTmSupplier, "jpaTmSupplier");
+    checkArgumentNotNull(jpaTmSupplier, "jpaTmSupplier");
     jpaTm = Suppliers.memoize(jpaTmSupplier::get);
     onBeam = true;
   }

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

@@ -217,11 +217,14 @@ abstract class JpaTransactionManagerExtension implements BeforeEachCallback, Aft
     JpaTransactionManagerImpl txnManager = new JpaTransactionManagerImpl(emf, clock);
     cachedTm = TransactionManagerFactory.jpaTm();
     TransactionManagerFactory.setJpaTm(Suppliers.ofInstance(txnManager));
+    TransactionManagerFactory.setReplicaJpaTm(
+        Suppliers.ofInstance(new ReplicaSimulatingJpaTransactionManager(txnManager)));
   }
 
   @Override
   public void afterEach(ExtensionContext context) {
     TransactionManagerFactory.setJpaTm(Suppliers.ofInstance(cachedTm));
+    TransactionManagerFactory.setReplicaJpaTm(Suppliers.ofInstance(cachedTm));
     // Even though we didn't set this, reset it to make sure no other tests are affected
     JpaTransactionManagerImpl.removeReplaySqlToDsOverrideForTest();
     cachedTm = null;

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

@@ -91,9 +91,15 @@ public class ReplicaSimulatingJpaTransactionManager implements JpaTransactionMan
 
   @Override
   public <T> T transact(Supplier<T> work) {
+    if (delegate.inTransaction()) {
+      return work.get();
+    }
     return delegate.transact(
         () -> {
-          delegate.getEntityManager().createQuery("SET TRANSACTION READ ONLY").executeUpdate();
+          delegate
+              .getEntityManager()
+              .createNativeQuery("SET TRANSACTION READ ONLY")
+              .executeUpdate();
           return work.get();
         });
   }