Make read-only transactions more performant (#2233)

Since the replica SQL instance is read-only, any transaction performed
on it should be explicitly read-only, which would allow PostgreSQL to
optimize away (some) use of predicate locks.

Also changed the EPP cache to read from the replica. The foreign key
cache already behaves this way.

See: https://www.postgresql.org/docs/current/transaction-iso.html

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

@@ -20,6 +20,7 @@ import static com.google.common.collect.Sets.difference;
 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.replicaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.util.CollectionUtils.nullToEmpty;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
@@ -357,13 +358,13 @@ public abstract class EppResource extends UpdateAutoTimestampEntity implements B
 
         @Override
         public EppResource load(VKey<? extends EppResource> key) {
-          return tm().reTransact(() -> tm().loadByKey(key));
+          return replicaTm().reTransact(() -> replicaTm().loadByKey(key));
         }
 
         @Override
         public Map<VKey<? extends EppResource>, EppResource> loadAll(
             Iterable<? extends VKey<? extends EppResource>> keys) {
-          return tm().reTransact(() -> tm().loadByKeys(keys));
+          return replicaTm().reTransact(() -> replicaTm().loadByKeys(keys));
         }
       };
 

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

@@ -165,10 +165,11 @@ public final class ForeignKeyUtils {
    * foreign keys should not use this cache.
    *
    * <p>Note that here the key of the {@link LoadingCache} is of type {@code VKey<? extends
-   * EppResource>}, but they are not legal {VKey}s to {@link EppResource}s, whose keys are the SQL
-   * primary keys, i.e. the {@code repoId}s. Instead, their keys are the foreign keys used to query
-   * the database. We use {@link VKey} here because it is a convenient composite class that contains
-   * both the resource type and the foreign key, which are needed to for the query and caching.
+   * EppResource>}, but they are not legal {@link VKey}s to {@link EppResource}s, whose keys are the
+   * SQL primary keys, i.e. the {@code repoId}s. Instead, their keys are the foreign keys used to
+   * query the database. We use {@link VKey} here because it is a convenient composite class that
+   * contains both the resource type and the foreign key, which are needed to for the query and
+   * caching.
    *
    * <p>Also note that the value type of this cache is {@link Optional} because the foreign keys in
    * question are coming from external commands, and thus don't necessarily represent entities in

core/src/main/java/google/registry/persistence/PersistenceModule.java

@@ -267,7 +267,7 @@ public abstract class PersistenceModule {
         name -> overrides.put(HIKARI_DS_CLOUD_SQL_INSTANCE, name));
     overrides.put(
         Environment.ISOLATION, TransactionIsolationLevel.TRANSACTION_REPEATABLE_READ.name());
-    return new JpaTransactionManagerImpl(create(overrides), clock);
+    return new JpaTransactionManagerImpl(create(overrides), clock, true);
   }
 
   @Provides
@@ -283,7 +283,7 @@ public abstract class PersistenceModule {
         name -> overrides.put(HIKARI_DS_CLOUD_SQL_INSTANCE, name));
     overrides.put(
         Environment.ISOLATION, TransactionIsolationLevel.TRANSACTION_REPEATABLE_READ.name());
-    return new JpaTransactionManagerImpl(create(overrides), clock);
+    return new JpaTransactionManagerImpl(create(overrides), clock, true);
   }
 
   /** Constructs the {@link EntityManagerFactory} instance. */

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

@@ -85,13 +85,19 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   // EntityManagerFactory is thread safe.
   private final EntityManagerFactory emf;
   private final Clock clock;
+  private final boolean readOnly;
 
   private static final ThreadLocal<TransactionInfo> transactionInfo =
       ThreadLocal.withInitial(TransactionInfo::new);
 
-  public JpaTransactionManagerImpl(EntityManagerFactory emf, Clock clock) {
+  public JpaTransactionManagerImpl(EntityManagerFactory emf, Clock clock, boolean readOnly) {
     this.emf = emf;
     this.clock = clock;
+    this.readOnly = readOnly;
+  }
+
+  public JpaTransactionManagerImpl(EntityManagerFactory emf, Clock clock) {
+    this(emf, clock, false);
   }
 
   @Override
@@ -200,6 +206,10 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     try {
       txn.begin();
       txnInfo.start(clock);
+      if (readOnly) {
+        getEntityManager().createNativeQuery("SET TRANSACTION READ ONLY").executeUpdate();
+        logger.atInfo().log("Using read-only SQL replica");
+      }
       if (isolationLevel != null && isolationLevel != getDefaultTransactionIsolationLevel()) {
         getEntityManager()
             .createNativeQuery(

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

@@ -219,10 +219,10 @@ public abstract class JpaTransactionManagerExtension
       recreateSchema();
     }
     JpaTransactionManagerImpl txnManager = new JpaTransactionManagerImpl(emf, clock);
+    JpaTransactionManagerImpl readOnlyTxnManager = new JpaTransactionManagerImpl(emf, clock, true);
     cachedTm = TransactionManagerFactory.tm();
     TransactionManagerFactory.setJpaTm(Suppliers.ofInstance(txnManager));
-    TransactionManagerFactory.setReplicaJpaTm(
-        Suppliers.ofInstance(new ReplicaSimulatingJpaTransactionManager(txnManager)));
+    TransactionManagerFactory.setReplicaJpaTm(Suppliers.ofInstance(readOnlyTxnManager));
     // Reset SQL Sequence based id allocation so that ids are deterministic in tests.
     TransactionManagerFactory.tm()
         .transact(

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

@@ -20,6 +20,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.persistence.PersistenceModule.TransactionIsolationLevel.TRANSACTION_READ_COMMITTED;
 import static google.registry.persistence.PersistenceModule.TransactionIsolationLevel.TRANSACTION_READ_UNCOMMITTED;
 import static google.registry.persistence.PersistenceModule.TransactionIsolationLevel.TRANSACTION_REPEATABLE_READ;
+import static google.registry.persistence.transaction.TransactionManagerFactory.replicaTm;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.assertDetachedFromEntityManager;
 import static google.registry.testing.DatabaseHelper.existsInDb;
@@ -107,6 +108,44 @@ class JpaTransactionManagerImplTest {
     assertCompanyExist("Bar");
   }
 
+  @Test
+  void transact_replica_failureOnWrite() {
+    assertPersonEmpty();
+    assertCompanyEmpty();
+    DatabaseException thrown =
+        assertThrows(
+            DatabaseException.class,
+            () ->
+                replicaTm()
+                    .transact(
+                        () -> {
+                          insertPerson(10);
+                        }));
+    assertThat(thrown)
+        .hasMessageThat()
+        .contains("cannot execute INSERT in a read-only transaction");
+  }
+
+  @Test
+  void transact_replica_successOnRead() {
+    assertPersonEmpty();
+    assertCompanyEmpty();
+    tm().transact(
+            () -> {
+              insertPerson(10);
+            });
+    replicaTm()
+        .transact(
+            () -> {
+              EntityManager em = replicaTm().getEntityManager();
+              Integer maybeAge =
+                  (Integer)
+                      em.createNativeQuery("SELECT age FROM Person WHERE age = 10")
+                          .getSingleResult();
+              assertThat(maybeAge).isEqualTo(10);
+            });
+  }
+
   @Test
   void transact_setIsolationLevel() {
     // If not specified, run at the default isolation level.

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

@@ -1,289 +0,0 @@
-// Copyright 2022 The Nomulus Authors. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package google.registry.persistence.transaction;
-
-import static com.google.common.base.Throwables.throwIfUnchecked;
-import static google.registry.persistence.transaction.DatabaseException.throwIfSqlException;
-
-import com.google.common.collect.ImmutableCollection;
-import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
-import google.registry.model.ImmutableObject;
-import google.registry.persistence.PersistenceModule.TransactionIsolationLevel;
-import google.registry.persistence.VKey;
-import java.util.Optional;
-import java.util.concurrent.Callable;
-import java.util.stream.Stream;
-import javax.persistence.EntityManager;
-import javax.persistence.Query;
-import javax.persistence.TypedQuery;
-import javax.persistence.criteria.CriteriaQuery;
-import org.joda.time.DateTime;
-
-/**
- * A {@link JpaTransactionManager} that simulates a read-only replica SQL instance.
- *
- * <p>We accomplish this by delegating all calls to the standard transaction manager except for
- * calls that start transactions. For these, we create a transaction like normal but set it to READ
- * ONLY mode before doing any work. This is similar to how the read-only Postgres replica works; it
- * treats all transactions as read-only transactions.
- */
-public class ReplicaSimulatingJpaTransactionManager implements JpaTransactionManager {
-
-  private final JpaTransactionManager delegate;
-
-  public ReplicaSimulatingJpaTransactionManager(JpaTransactionManager delegate) {
-    this.delegate = delegate;
-  }
-
-  @Override
-  public void teardown() {
-    delegate.teardown();
-  }
-
-  @Override
-  public TransactionIsolationLevel getDefaultTransactionIsolationLevel() {
-    return delegate.getDefaultTransactionIsolationLevel();
-  }
-
-  @Override
-  public TransactionIsolationLevel getCurrentTransactionIsolationLevel() {
-    return delegate.getCurrentTransactionIsolationLevel();
-  }
-
-  @Override
-  public EntityManager getStandaloneEntityManager() {
-    return delegate.getStandaloneEntityManager();
-  }
-
-  @Override
-  public EntityManager getEntityManager() {
-    return delegate.getEntityManager();
-  }
-
-  @Override
-  public <T> TypedQuery<T> query(String sqlString, Class<T> resultClass) {
-    return delegate.query(sqlString, resultClass);
-  }
-
-  @Override
-  public <T> TypedQuery<T> criteriaQuery(CriteriaQuery<T> criteriaQuery) {
-    return delegate.criteriaQuery(criteriaQuery);
-  }
-
-  @Override
-  public Query query(String sqlString) {
-    return delegate.query(sqlString);
-  }
-
-  @Override
-  public boolean inTransaction() {
-    return delegate.inTransaction();
-  }
-
-  @Override
-  public void assertInTransaction() {
-    delegate.assertInTransaction();
-  }
-
-  @Override
-  public <T> T transact(Callable<T> work, TransactionIsolationLevel isolationLevel) {
-    if (inTransaction()) {
-      try {
-        return work.call();
-      } catch (Exception e) {
-        throwIfSqlException(e);
-        throwIfUnchecked(e);
-        throw new RuntimeException(e);
-      }
-    }
-    return delegate.transact(
-        () -> {
-          delegate
-              .getEntityManager()
-              .createNativeQuery("SET TRANSACTION READ ONLY")
-              .executeUpdate();
-          return work.call();
-        },
-        isolationLevel);
-  }
-
-  @Override
-  public <T> T reTransact(Callable<T> work) {
-    return transact(work);
-  }
-
-  @Override
-  public <T> T transact(Callable<T> work) {
-    return transact(work, null);
-  }
-
-  @Override
-  public void transact(ThrowingRunnable work, TransactionIsolationLevel isolationLevel) {
-    transact(
-        () -> {
-          work.run();
-          return null;
-        },
-        isolationLevel);
-  }
-
-  @Override
-  public void reTransact(ThrowingRunnable work) {
-    transact(work);
-  }
-
-  @Override
-  public void transact(ThrowingRunnable work) {
-    transact(work, null);
-  }
-
-  @Override
-  public DateTime getTransactionTime() {
-    return delegate.getTransactionTime();
-  }
-
-  @Override
-  public void insert(Object entity) {
-    delegate.insert(entity);
-  }
-
-  @Override
-  public void insertAll(ImmutableCollection<?> entities) {
-    delegate.insertAll(entities);
-  }
-
-  @Override
-  public void insertAll(ImmutableObject... entities) {
-    delegate.insertAll(entities);
-  }
-
-  @Override
-  public void put(Object entity) {
-    delegate.put(entity);
-  }
-
-  @Override
-  public void putAll(ImmutableObject... entities) {
-    delegate.putAll(entities);
-  }
-
-  @Override
-  public void putAll(ImmutableCollection<?> entities) {
-    delegate.putAll(entities);
-  }
-
-  @Override
-  public void update(Object entity) {
-    delegate.update(entity);
-  }
-
-  @Override
-  public void updateAll(ImmutableCollection<?> entities) {
-    delegate.updateAll(entities);
-  }
-
-  @Override
-  public void updateAll(ImmutableObject... entities) {
-    delegate.updateAll(entities);
-  }
-
-  @Override
-  public <T> boolean exists(VKey<T> key) {
-    return delegate.exists(key);
-  }
-
-  @Override
-  public boolean exists(Object entity) {
-    return delegate.exists(entity);
-  }
-
-  @Override
-  public <T> Optional<T> loadByKeyIfPresent(VKey<T> key) {
-    return delegate.loadByKeyIfPresent(key);
-  }
-
-  @Override
-  public <T> ImmutableMap<VKey<? extends T>, T> loadByKeysIfPresent(
-      Iterable<? extends VKey<? extends T>> vKeys) {
-    return delegate.loadByKeysIfPresent(vKeys);
-  }
-
-  @Override
-  public <T> ImmutableList<T> loadByEntitiesIfPresent(Iterable<T> entities) {
-    return delegate.loadByEntitiesIfPresent(entities);
-  }
-
-  @Override
-  public <T> T loadByKey(VKey<T> key) {
-    return delegate.loadByKey(key);
-  }
-
-  @Override
-  public <T> ImmutableMap<VKey<? extends T>, T> loadByKeys(
-      Iterable<? extends VKey<? extends T>> vKeys) {
-    return delegate.loadByKeys(vKeys);
-  }
-
-  @Override
-  public <T> T loadByEntity(T entity) {
-    return delegate.loadByEntity(entity);
-  }
-
-  @Override
-  public <T> ImmutableList<T> loadByEntities(Iterable<T> entities) {
-    return delegate.loadByEntities(entities);
-  }
-
-  @Override
-  public <T> ImmutableList<T> loadAllOf(Class<T> clazz) {
-    return delegate.loadAllOf(clazz);
-  }
-
-  @Override
-  public <T> Stream<T> loadAllOfStream(Class<T> clazz) {
-    return delegate.loadAllOfStream(clazz);
-  }
-
-  @Override
-  public <T> Optional<T> loadSingleton(Class<T> clazz) {
-    return delegate.loadSingleton(clazz);
-  }
-
-  @Override
-  public void delete(VKey<?> key) {
-    delegate.delete(key);
-  }
-
-  @Override
-  public void delete(Iterable<? extends VKey<?>> vKeys) {
-    delegate.delete(vKeys);
-  }
-
-  @Override
-  public <T> T delete(T entity) {
-    return delegate.delete(entity);
-  }
-
-  @Override
-  public <T> QueryComposer<T> createQueryComposer(Class<T> entity) {
-    return delegate.createQueryComposer(entity);
-  }
-
-  @Override
-  public <T> void assertDelete(VKey<T> key) {
-    delegate.assertDelete(key);
-  }
-}