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(
-        () -> {
-          getEntityManager().createNativeQuery("SET TRANSACTION READ ONLY").executeUpdate();
-          return work.get();
-        });
+    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