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.
This commit is contained in:
gbrodman 2022-01-28 18:05:18 -05:00 committed by GitHub
parent 1253fa479a
commit b24670f33a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 75 additions and 21 deletions

View file

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

View file

@ -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(
() ->
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<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 {

View file

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

View file

@ -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;

View file

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