diff --git a/core/src/main/java/google/registry/persistence/JpaRetries.java b/core/src/main/java/google/registry/persistence/JpaRetries.java index e9e210548..d3f2b0596 100644 --- a/core/src/main/java/google/registry/persistence/JpaRetries.java +++ b/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 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; } } diff --git a/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManager.java b/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManager.java index 281482a8a..f723ac41d 100644 --- a/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManager.java +++ b/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 transactNoRetry(Supplier 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 void assertDelete(VKey key); diff --git a/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java b/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java index 7766aad5d..ff17843bf 100644 --- a/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java +++ b/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 transact(Supplier 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 transactNoRetry(Supplier 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 transactNew(Supplier work) { return transact(work); @@ -142,11 +189,14 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager { @Override public T transactNewReadOnly(Supplier 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 doTransactionless(Supplier work) { - return transact(work); + return retrier.callWithRetry(() -> transact(work), JpaRetries::isFailedQueryRetriable); } @Override diff --git a/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerImplTest.java b/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerImplTest.java index 586867709..360773061 100644 --- a/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerImplTest.java +++ b/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 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 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 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 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 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();