Add retrier to JpaTransactionManager (#693)

* Add retrier to retry exceptions we've seen so far

* Check for nested exceptions

* fix formatting

* check for nested transactions

* add JDBCConnectionException to isFailedQueryRetriable

* Add retrier to methods with suppliers

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

@@ -19,6 +19,7 @@ import com.google.common.collect.ImmutableSet;
 import java.sql.SQLException;
 import java.util.function.Predicate;
 import javax.persistence.OptimisticLockException;
+import org.hibernate.exception.JDBCConnectionException;
 
 /** Helpers for identifying retriable database operations. */
 public final class JpaRetries {
@@ -40,6 +41,13 @@ public final class JpaRetries {
               e instanceof SQLException
                   && RETRIABLE_TXN_SQL_STATE.contains(((SQLException) e).getSQLState()));
 
+  private static final Predicate<Throwable> RETRIABLE_QUERY_PREDICATE =
+      Predicates.or(
+          JDBCConnectionException.class::isInstance,
+          e ->
+              e instanceof SQLException
+                  && RETRIABLE_TXN_SQL_STATE.contains(((SQLException) e).getSQLState()));
+
   public static boolean isFailedTxnRetriable(Throwable throwable) {
     Throwable t = throwable;
     while (t != null) {
@@ -53,6 +61,13 @@ public final class JpaRetries {
 
   public static boolean isFailedQueryRetriable(Throwable throwable) {
     // TODO(weiminyu): check for more error codes.
-    return isFailedTxnRetriable(throwable);
+    Throwable t = throwable;
+    while (t != null) {
+      if (RETRIABLE_QUERY_PREDICATE.test(t)) {
+        return true;
+      }
+      t = t.getCause();
+    }
+    return false;
   }
 }

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

@@ -15,6 +15,7 @@
 package google.registry.persistence.transaction;
 
 import google.registry.persistence.VKey;
+import java.util.function.Supplier;
 import javax.persistence.EntityManager;
 
 /** Sub-interface of {@link TransactionManager} which defines JPA related methods. */
@@ -23,6 +24,12 @@ public interface JpaTransactionManager extends TransactionManager {
   /** Returns the {@link EntityManager} for the current request. */
   EntityManager getEntityManager();
 
+  /** Executes the work in a transaction with no retries and returns the result. */
+  <T> T transactNoRetry(Supplier<T> work);
+
+  /** Executes the work in a transaction with no retries. */
+  void transactNoRetry(Runnable work);
+
   /** Deletes the entity by its id, throws exception if the entity is not deleted. */
   public abstract <T> void assertDelete(VKey<T> key);
 

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

@@ -27,8 +27,11 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.FluentLogger;
 import google.registry.config.RegistryConfig;
+import google.registry.persistence.JpaRetries;
 import google.registry.persistence.VKey;
 import google.registry.util.Clock;
+import google.registry.util.Retrier;
+import google.registry.util.SystemSleeper;
 import java.lang.reflect.Field;
 import java.util.Map.Entry;
 import java.util.NoSuchElementException;
@@ -49,6 +52,7 @@ import org.joda.time.DateTime;
 public class JpaTransactionManagerImpl implements JpaTransactionManager {
 
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+  private static final Retrier retrier = new Retrier(new SystemSleeper(), 3);
 
   // EntityManagerFactory is thread safe.
   private final EntityManagerFactory emf;
@@ -94,6 +98,40 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   public <T> T transact(Supplier<T> work) {
     // TODO(shicong): Investigate removing transactNew functionality after migration as it may
     //  be same as this one.
+    return retrier.callWithRetry(
+        () -> {
+          if (inTransaction()) {
+            return work.get();
+          }
+          TransactionInfo txnInfo = transactionInfo.get();
+          txnInfo.entityManager = emf.createEntityManager();
+          EntityTransaction txn = txnInfo.entityManager.getTransaction();
+          try {
+            txn.begin();
+            txnInfo.start(clock);
+            T result = work.get();
+            txnInfo.recordTransaction();
+            txn.commit();
+            return result;
+          } catch (RuntimeException | Error e) {
+            // Error is unchecked!
+            try {
+              txn.rollback();
+              logger.atWarning().log("Error during transaction; transaction rolled back");
+            } catch (Throwable rollbackException) {
+              logger.atSevere().withCause(rollbackException).log(
+                  "Rollback failed; suppressing error");
+            }
+            throw e;
+          } finally {
+            txnInfo.clear();
+          }
+        },
+        JpaRetries::isFailedTxnRetriable);
+  }
+
+  @Override
+  public <T> T transactNoRetry(Supplier<T> work) {
     if (inTransaction()) {
       return work.get();
     }
@@ -130,6 +168,15 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
         });
   }
 
+  @Override
+  public void transactNoRetry(Runnable work) {
+    transactNoRetry(
+        () -> {
+          work.run();
+          return null;
+        });
+  }
+
   @Override
   public <T> T transactNew(Supplier<T> work) {
     return transact(work);
@@ -142,11 +189,14 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
 
   @Override
   public <T> T transactNewReadOnly(Supplier<T> work) {
-    return transact(
+    return retrier.callWithRetry(
+        () ->
+            transact(
                 () -> {
                   getEntityManager().createNativeQuery("SET TRANSACTION READ ONLY").executeUpdate();
                   return work.get();
-        });
+                }),
+        JpaRetries::isFailedQueryRetriable);
   }
 
   @Override
@@ -160,7 +210,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
 
   @Override
   public <T> T doTransactionless(Supplier<T> work) {
-    return transact(work);
+    return retrier.callWithRetry(() -> transact(work), JpaRetries::isFailedQueryRetriable);
   }
 
   @Override

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

@@ -18,6 +18,11 @@ import static com.google.common.truth.Truth.assertThat;
 import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
 import static google.registry.testing.TestDataHelper.fileClassPath;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 import com.google.common.collect.ImmutableList;
 import google.registry.model.ImmutableObject;
@@ -26,12 +31,16 @@ import google.registry.persistence.transaction.JpaTestRules.JpaUnitTestExtension
 import google.registry.testing.FakeClock;
 import java.io.Serializable;
 import java.math.BigInteger;
+import java.sql.SQLException;
 import java.util.NoSuchElementException;
+import java.util.function.Supplier;
 import javax.persistence.Entity;
 import javax.persistence.EntityManager;
 import javax.persistence.Id;
 import javax.persistence.IdClass;
+import javax.persistence.OptimisticLockException;
 import javax.persistence.RollbackException;
+import org.hibernate.exception.JDBCConnectionException;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
@@ -130,6 +139,117 @@ class JpaTransactionManagerImplTest {
     assertThat(jpaTm().transact(() -> jpaTm().load(theEntityKey))).isEqualTo(theEntity);
   }
 
+  @Test
+  public void transact_retriesOptimisticLockExceptions() {
+    JpaTransactionManager spyJpaTm = spy(jpaTm());
+    doThrow(OptimisticLockException.class).when(spyJpaTm).delete(any(VKey.class));
+    spyJpaTm.transact(() -> spyJpaTm.saveNew(theEntity));
+    assertThrows(
+        OptimisticLockException.class,
+        () -> spyJpaTm.transact(() -> spyJpaTm.delete(theEntityKey)));
+    verify(spyJpaTm, times(3)).delete(theEntityKey);
+    Supplier<Runnable> supplier =
+        () -> {
+          Runnable work = () -> spyJpaTm.delete(theEntityKey);
+          work.run();
+          return null;
+        };
+    assertThrows(OptimisticLockException.class, () -> spyJpaTm.transact(supplier));
+    verify(spyJpaTm, times(6)).delete(theEntityKey);
+  }
+
+  @Test
+  public void transactNoRetry_doesNotRetryOptimisticLockException() {
+    JpaTransactionManager spyJpaTm = spy(jpaTm());
+    doThrow(OptimisticLockException.class).when(spyJpaTm).delete(any(VKey.class));
+    spyJpaTm.transactNoRetry(() -> spyJpaTm.saveNew(theEntity));
+    assertThrows(
+        OptimisticLockException.class,
+        () -> spyJpaTm.transactNoRetry(() -> spyJpaTm.delete(theEntityKey)));
+    verify(spyJpaTm, times(1)).delete(theEntityKey);
+    Supplier<Runnable> supplier =
+        () -> {
+          Runnable work = () -> spyJpaTm.delete(theEntityKey);
+          work.run();
+          return null;
+        };
+    assertThrows(OptimisticLockException.class, () -> spyJpaTm.transactNoRetry(supplier));
+    verify(spyJpaTm, times(2)).delete(theEntityKey);
+  }
+
+  @Test
+  public void transact_retriesNestedOptimisticLockExceptions() {
+    JpaTransactionManager spyJpaTm = spy(jpaTm());
+    doThrow(new RuntimeException().initCause(new OptimisticLockException()))
+        .when(spyJpaTm)
+        .delete(any(VKey.class));
+    spyJpaTm.transact(() -> spyJpaTm.saveNew(theEntity));
+    assertThrows(
+        RuntimeException.class, () -> spyJpaTm.transact(() -> spyJpaTm.delete(theEntityKey)));
+    verify(spyJpaTm, times(3)).delete(theEntityKey);
+    Supplier<Runnable> supplier =
+        () -> {
+          Runnable work = () -> spyJpaTm.delete(theEntityKey);
+          work.run();
+          return null;
+        };
+    assertThrows(RuntimeException.class, () -> spyJpaTm.transact(supplier));
+    verify(spyJpaTm, times(6)).delete(theEntityKey);
+  }
+
+  @Test
+  public void transactNewReadOnly_retriesJdbcConnectionExceptions() {
+    JpaTransactionManager spyJpaTm = spy(jpaTm());
+    doThrow(JDBCConnectionException.class).when(spyJpaTm).load(any(VKey.class));
+    spyJpaTm.transact(() -> spyJpaTm.saveNew(theEntity));
+    assertThrows(
+        JDBCConnectionException.class,
+        () -> spyJpaTm.transactNewReadOnly(() -> spyJpaTm.load(theEntityKey)));
+    verify(spyJpaTm, times(3)).load(theEntityKey);
+    Supplier<Runnable> supplier =
+        () -> {
+          Runnable work = () -> spyJpaTm.load(theEntityKey);
+          work.run();
+          return null;
+        };
+    assertThrows(JDBCConnectionException.class, () -> spyJpaTm.transactNewReadOnly(supplier));
+    verify(spyJpaTm, times(6)).load(theEntityKey);
+  }
+
+  @Test
+  public void transactNewReadOnly_retriesNestedJdbcConnectionExceptions() {
+    JpaTransactionManager spyJpaTm = spy(jpaTm());
+    doThrow(
+            new RuntimeException()
+                .initCause(new JDBCConnectionException("connection exception", new SQLException())))
+        .when(spyJpaTm)
+        .load(any(VKey.class));
+    spyJpaTm.transact(() -> spyJpaTm.saveNew(theEntity));
+    assertThrows(
+        RuntimeException.class,
+        () -> spyJpaTm.transactNewReadOnly(() -> spyJpaTm.load(theEntityKey)));
+    verify(spyJpaTm, times(3)).load(theEntityKey);
+    Supplier<Runnable> supplier =
+        () -> {
+          Runnable work = () -> spyJpaTm.load(theEntityKey);
+          work.run();
+          return null;
+        };
+    assertThrows(RuntimeException.class, () -> spyJpaTm.transactNewReadOnly(supplier));
+    verify(spyJpaTm, times(6)).load(theEntityKey);
+  }
+
+  @Test
+  public void doTransactionless_retriesJdbcConnectionExceptions() {
+    JpaTransactionManager spyJpaTm = spy(jpaTm());
+    doThrow(JDBCConnectionException.class).when(spyJpaTm).load(any(VKey.class));
+    spyJpaTm.transact(() -> spyJpaTm.saveNew(theEntity));
+    assertThrows(
+        RuntimeException.class,
+        () -> spyJpaTm.doTransactionless(() -> spyJpaTm.load(theEntityKey)));
+    verify(spyJpaTm, times(3)).load(theEntityKey);
+  }
+
   @Test
   void saveNew_throwsExceptionIfEntityExists() {
     assertThat(jpaTm().transact(() -> jpaTm().checkExists(theEntity))).isFalse();