diff --git a/core/src/main/java/google/registry/model/EppResource.java b/core/src/main/java/google/registry/model/EppResource.java index 208e3852b..77f7e682d 100644 --- a/core/src/main/java/google/registry/model/EppResource.java +++ b/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 key) { - return tm().doTransactionless(() -> tm().loadByKey(key)); + return replicaTm().doTransactionless(() -> replicaTm().loadByKey(key)); } @Override public Map, EppResource> loadAll( Iterable> keys) { - return tm().doTransactionless(() -> tm().loadByKeys(keys)); + return replicaTm().doTransactionless(() -> replicaTm().loadByKeys(keys)); } }; diff --git a/core/src/main/java/google/registry/model/index/ForeignKeyIndex.java b/core/src/main/java/google/registry/model/index/ForeignKeyIndex.java index 0437931e2..015f3b60b 100644 --- a/core/src/main/java/google/registry/model/index/ForeignKeyIndex.java +++ b/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 extends BackupGroup */ public static ImmutableMap> load( Class clazz, Collection 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 extends BackupGroup */ private static ImmutableMap> loadIndexesFromStore( - Class clazz, Collection foreignKeys, boolean inTransaction) { + Class clazz, + Collection foreignKeys, + boolean inTransaction, + boolean useReplicaJpaTm) { if (tm().isOfy()) { Class> fkiClass = mapToFkiClass(clazz); return ImmutableMap.copyOf( @@ -226,17 +231,18 @@ public abstract class ForeignKeyIndex 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> indexes = - tm().transact( - () -> - jpaTm() - .criteriaQuery( - CriteriaQueryBuilder.create(clazz) - .whereFieldIsIn(property, foreignKeys) - .build()) - .getResultStream() - .map(e -> ForeignKeyIndex.create(e, e.getDeletionTime())) - .collect(toImmutableList())); + jpaTmToUse.transact( + () -> + jpaTmToUse + .criteriaQuery( + CriteriaQueryBuilder.create(clazz) + .whereFieldIsIn(property, foreignKeys) + .build()) + .getResultStream() + .map(e -> ForeignKeyIndex.create(e, e.getDeletionTime())) + .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 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 extends BackupGroup Streams.stream(keys).map(v -> v.getSqlKey().toString()).collect(toImmutableSet()); ImmutableSet>> typedKeys = ImmutableSet.copyOf(keys); ImmutableMap> 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 extends BackupGroup // Safe to cast VKey> to VKey> @SuppressWarnings("unchecked") ImmutableList>> fkiVKeys = - Streams.stream(foreignKeys) + foreignKeys.stream() .map(fk -> (VKey>) VKey.create(fkiClass, fk)) .collect(toImmutableList()); try { diff --git a/core/src/main/java/google/registry/persistence/transaction/TransactionManagerFactory.java b/core/src/main/java/google/registry/persistence/transaction/TransactionManagerFactory.java index a2c4da2d5..891748d96 100644 --- a/core/src/main/java/google/registry/persistence/transaction/TransactionManagerFactory.java +++ b/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 jpaTm = Suppliers.memoize(TransactionManagerFactory::createJpaTransactionManager); + @NonFinalForTesting + private static Supplier 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. + * + *

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 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 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 jpaTmSupplier) { - checkNotNull(jpaTmSupplier, "jpaTmSupplier"); + checkArgumentNotNull(jpaTmSupplier, "jpaTmSupplier"); jpaTm = Suppliers.memoize(jpaTmSupplier::get); onBeam = true; } diff --git a/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerExtension.java b/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerExtension.java index 77bd456e1..4cdbb96c4 100644 --- a/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerExtension.java +++ b/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; diff --git a/core/src/test/java/google/registry/persistence/transaction/ReplicaSimulatingJpaTransactionManager.java b/core/src/test/java/google/registry/persistence/transaction/ReplicaSimulatingJpaTransactionManager.java index fe9f2a88b..cfbea25ae 100644 --- a/core/src/test/java/google/registry/persistence/transaction/ReplicaSimulatingJpaTransactionManager.java +++ b/core/src/test/java/google/registry/persistence/transaction/ReplicaSimulatingJpaTransactionManager.java @@ -91,9 +91,15 @@ public class ReplicaSimulatingJpaTransactionManager implements JpaTransactionMan @Override public T transact(Supplier 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(); }); }