SRE-627: match hibernate exception handling in jooq

.gitignore

@@ -1,3 +1,4 @@
+org
 .gradle
 .gradletasknamecache
 **/out/

build.gradle.kts

@@ -330,6 +330,8 @@ subprojects {
   }
 
   tasks.withType<Detekt>().configureEach {
+    // Keep Detekt CLI happy regardless of host JDK; 1.23.x supports up to 22
+    jvmTarget = "21"
     dependsOn(":detektive:assemble")
     exclude { it.file.absolutePath.contains("/generated/source/") || it.file.absolutePath.contains("SampledLogger") }
   }

misk-hibernate/api/misk-hibernate.api

@@ -108,6 +108,11 @@ public abstract class misk/hibernate/HibernateEntityModule : misk/inject/KAbstra
 	protected final fun installHibernateAdminDashboardWebActions ()V
 }
 
+public final class misk/hibernate/HibernateExceptionClassifier : misk/jdbc/DefaultExceptionClassifier {
+	public fun <init> ()V
+	public fun isRetryable (Ljava/lang/Throwable;)Z
+}
+
 public final class misk/hibernate/HibernateExceptionLogLevelConfig : misk/config/Config {
 	public fun <init> ()V
 	public fun <init> (Lorg/slf4j/event/Level;)V

misk-hibernate/src/main/kotlin/misk/hibernate/HibernateExceptionClassifier.kt

@@ -0,0 +1,29 @@
+package misk.hibernate
+
+import misk.jdbc.DefaultExceptionClassifier
+import misk.jdbc.RetryTransactionException
+import javax.persistence.OptimisticLockException
+import org.hibernate.StaleObjectStateException
+import org.hibernate.exception.LockAcquisitionException
+
+/**
+ * Exception classifier for Hibernate-specific exceptions.
+ * 
+ * This extends the default classifier to handle Hibernate-specific retryable exceptions
+ * like OptimisticLockException, StaleObjectStateException, etc.
+ */
+class HibernateExceptionClassifier : DefaultExceptionClassifier() {
+
+  override fun isRetryable(th: Throwable): Boolean {
+    return when (th) {
+      // Hibernate-specific retryable exceptions
+      is OptimisticLockException,
+      is StaleObjectStateException,
+      is LockAcquisitionException -> true
+      // Custom retry exception for application-level retries
+      is RetryTransactionException -> true
+      // Fall back to default classification
+      else -> super.isRetryable(th)
+    }
+  }
+}

misk-hibernate/src/main/kotlin/misk/hibernate/RealTransacter.kt

@@ -17,8 +17,8 @@ import java.util.concurrent.Future
 import java.util.concurrent.TimeUnit
 import javax.persistence.OptimisticLockException
 import kotlin.reflect.KClass
-import misk.backoff.ExponentialBackoff
 import misk.concurrent.ExecutorServiceFactory
+import misk.jdbc.TransactionRetryHandler
 import misk.hibernate.advisorylocks.tryAcquireLock
 import misk.hibernate.advisorylocks.tryReleaseLock
 import misk.jdbc.CheckDisabler
@@ -58,6 +58,11 @@ private constructor(
   private val shardListFetcher: ShardListFetcher,
   private val hibernateEntities: Set<HibernateEntity>,
 ) : Transacter {
+
+  private val retryHandler = TransactionRetryHandler(
+    qualifierName = qualifier.simpleName ?: "hibernate",
+    exceptionClassifier = HibernateExceptionClassifier()
+  )
   constructor(
     qualifier: KClass<out Annotation>,
     sessionFactoryService: SessionFactoryService,
@@ -250,46 +255,13 @@ private constructor(
   }
 
   private fun <T> transactionWithRetriesInternal(block: () -> T): T {
-    require(options.maxAttempts > 0)
-
-    val backoff =
-      ExponentialBackoff(
-        Duration.ofMillis(options.minRetryDelayMillis),
-        Duration.ofMillis(options.maxRetryDelayMillis),
-        Duration.ofMillis(options.retryJitterMillis),
-      )
-    var attempt = 0
-
-    while (true) {
-      try {
-        attempt++
-        val result = block()
-
-        if (attempt > 1) {
-          logger.info { "retried ${qualifier.simpleName} transaction succeeded (attempt $attempt)" }
-        }
-
-        return result
-      } catch (e: Exception) {
-        if (!isRetryable(e)) throw e
-
-        if (attempt >= options.maxAttempts) {
-          logger.info {
-            "${qualifier.simpleName} recoverable transaction exception " + "(attempt $attempt), no more attempts"
-          }
-          throw e
-        }
-
-        val sleepDuration = backoff.nextRetry()
-        logger.info(e) {
-          "${qualifier.simpleName} recoverable transaction exception " +
-            "(attempt $attempt), will retry after a $sleepDuration delay"
-        }
-
-        if (!sleepDuration.isZero) {
-          Thread.sleep(sleepDuration.toMillis())
-        }
-      }
+    return retryHandler.executeWithRetries(
+      maxAttempts = options.maxAttempts,
+      minRetryDelayMillis = options.minRetryDelayMillis,
+      maxRetryDelayMillis = options.maxRetryDelayMillis,
+      retryJitterMillis = options.retryJitterMillis
+    ) {
+      block()
     }
   }
 
@@ -335,7 +307,7 @@ private constructor(
         return ConstraintViolationException(sqlException.message, sqlException, "")
       }
       // write-write conflicts fail at COMMIT rather than waiting on a lock.
-      e.cause is SQLException && isTidbWriteConflict(e.cause as SQLException) -> {
+      e.cause is SQLException && (e.cause as SQLException).errorCode == 9007 -> {
         val sqlException = e.cause as SQLException
         return ConstraintViolationException(sqlException.message, sqlException, "")
       }
@@ -454,68 +426,7 @@ private constructor(
     }
   }
 
-  private fun isRetryable(th: Throwable): Boolean {
-    return when (th) {
-      is RetryTransactionException,
-      is StaleObjectStateException,
-      is LockAcquisitionException,
-      is SQLRecoverableException,
-      is SQLTransientException,
-      is OptimisticLockException -> true
-      is SQLException -> if (isMessageRetryable(th)) true else isCauseRetryable(th)
-      else -> isCauseRetryable(th)
-    }
-  }
-
-  private fun isMessageRetryable(th: SQLException) =
-    isConnectionClosed(th) ||
-      isVitessTransactionNotFound(th) ||
-      isCockroachRestartTransaction(th) ||
-      isTidbWriteConflict(th)
-
-  /**
-   * This is thrown as a raw SQLException from Hikari even though it is most certainly a recoverable exception. See
-   * com/zaxxer/hikari/pool/ProxyConnection.java:493
-   */
-  private fun isConnectionClosed(th: SQLException) = th.message.equals("Connection is closed")
-
-  /**
-   * We get this error as a MySQLQueryInterruptedException when a tablet gracefully terminates, we just need to retry
-   * the transaction and the new primary should handle it.
-   *
-   * ```
-   * vttablet: rpc error: code = Aborted desc = transaction 1572922696317821557:
-   * not found (CallerID: )
-   * ```
-   */
-  private fun isVitessTransactionNotFound(th: SQLException): Boolean {
-    val message = th.message
-    return message != null &&
-      message.contains("vttablet: rpc error") &&
-      message.contains("code = Aborted") &&
-      message.contains("transaction") &&
-      message.contains("not found")
-  }
-
-  /**
-   * "Messages with the error code 40001 and the string restart transaction indicate that a transaction failed because
-   * it conflicted with another concurrent or recent transaction accessing the same data. The transaction needs to be
-   * retried by the client." https://www.cockroachlabs.com/docs/stable/common-errors.html#restart-transaction
-   */
-  private fun isCockroachRestartTransaction(th: SQLException): Boolean {
-    val message = th.message
-    return th.errorCode == 40001 && message != null && message.contains("restart transaction")
-  }
-
-  /**
-   * "Transactions in TiKV encounter write conflicts". This can happen when optimistic transaction mode is on. Conflicts
-   * are detected during transaction commit https://docs.pingcap.com/tidb/dev/tidb-faq#error-9007-hy000-write-conflict
-   */
-  private fun isTidbWriteConflict(th: SQLException): Boolean {
-    return th.errorCode == 9007
-  }
 
-  private fun isCauseRetryable(th: Throwable) = th.cause?.let { isRetryable(it) } ?: false
 
   // NB: all options should be immutable types as copy() is shallow.
   internal data class TransacterOptions(

misk-hibernate/src/test/kotlin/misk/hibernate/TransacterTest.kt

@@ -6,6 +6,7 @@ import misk.hibernate.VitessTestExtensions.createInSeparateShard
 import misk.hibernate.VitessTestExtensions.save
 import misk.hibernate.VitessTestExtensions.shard
 import misk.jdbc.DataSourceType
+import misk.jdbc.RetryTransactionException
 import misk.jdbc.uniqueString
 import misk.testing.MiskExternalDependency
 import misk.testing.MiskTest
@@ -720,7 +721,7 @@ abstract class TransacterTest {
       if (callCount.getAndIncrement() < 2) throw RetryTransactionException()
     }
 
-    val logs = logCollector.takeMessages(RealTransacter::class)
+    val logs = logCollector.takeMessages(misk.jdbc.TransactionRetryHandler::class)
     assertThat(logs).hasSize(3)
     assertThat(logs[0]).matches(
       "Movies recoverable transaction exception " +

misk-jdbc/api/misk-jdbc.api

@@ -398,6 +398,23 @@ public final class misk/jdbc/DeclarativeSchemaConfig {
 	public fun toString ()Ljava/lang/String;
 }
 
+public class misk/jdbc/DefaultExceptionClassifier : misk/jdbc/ExceptionClassifier {
+	public fun <init> ()V
+	public fun <init> (Lmisk/jdbc/DataSourceType;)V
+	public synthetic fun <init> (Lmisk/jdbc/DataSourceType;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	protected final fun isCauseRetryable (Ljava/lang/Throwable;)Z
+	protected final fun isConnectionClosed (Ljava/sql/SQLException;)Z
+	protected final fun isMessageRetryable (Ljava/sql/SQLException;)Z
+	protected final fun isRecoverableCockroachException (Ljava/sql/SQLException;)Z
+	protected final fun isRecoverableTidbException (Ljava/sql/SQLException;)Z
+	protected final fun isRecoverableVitessException (Ljava/sql/SQLException;)Z
+	public fun isRetryable (Ljava/lang/Throwable;)Z
+}
+
+public abstract interface class misk/jdbc/ExceptionClassifier {
+	public abstract fun isRetryable (Ljava/lang/Throwable;)Z
+}
+
 public class misk/jdbc/ExtendedQueryExecutionListener : net/ttddyy/dsproxy/listener/MethodExecutionListener, net/ttddyy/dsproxy/listener/QueryExecutionListener {
 	public static final field Companion Lmisk/jdbc/ExtendedQueryExecutionListener$Companion;
 	public fun <init> ()V
@@ -561,6 +578,13 @@ public final class misk/jdbc/RealTransacter : misk/jdbc/Transacter {
 	public fun transactionWithSession (Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 }
 
+public final class misk/jdbc/RetryTransactionException : java/lang/RuntimeException {
+	public fun <init> ()V
+	public fun <init> (Ljava/lang/String;)V
+	public fun <init> (Ljava/lang/String;Ljava/lang/Throwable;)V
+	public synthetic fun <init> (Ljava/lang/String;Ljava/lang/Throwable;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+}
+
 public final class misk/jdbc/ScaleSafetyChecks {
 	public static final field INSTANCE Lmisk/jdbc/ScaleSafetyChecks;
 	public final fun checkQueryForTableScan (Ljava/sql/Connection;Ljava/lang/String;Ljava/lang/String;)V
@@ -677,6 +701,26 @@ public abstract interface class misk/jdbc/Transacter {
 	public abstract fun transactionWithSession (Lkotlin/jvm/functions/Function1;)Ljava/lang/Object;
 }
 
+public final class misk/jdbc/TransactionRetryHandler {
+	public static final field Companion Lmisk/jdbc/TransactionRetryHandler$Companion;
+	public fun <init> ()V
+	public fun <init> (Ljava/lang/String;)V
+	public fun <init> (Ljava/lang/String;Lmisk/jdbc/DataSourceType;)V
+	public synthetic fun <init> (Ljava/lang/String;Lmisk/jdbc/DataSourceType;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public fun <init> (Ljava/lang/String;Lmisk/jdbc/ExceptionClassifier;)V
+	public fun <init> (Ljava/lang/String;Lmisk/jdbc/ExceptionClassifier;Lmisk/jdbc/DataSourceType;)V
+	public synthetic fun <init> (Ljava/lang/String;Lmisk/jdbc/ExceptionClassifier;Lmisk/jdbc/DataSourceType;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public final fun executeWithRetries (IJJJLkotlin/jvm/functions/Function0;)Ljava/lang/Object;
+	public final fun executeWithRetries (IJJLkotlin/jvm/functions/Function0;)Ljava/lang/Object;
+	public final fun executeWithRetries (IJLkotlin/jvm/functions/Function0;)Ljava/lang/Object;
+	public final fun executeWithRetries (ILkotlin/jvm/functions/Function0;)Ljava/lang/Object;
+	public final fun executeWithRetries (Lkotlin/jvm/functions/Function0;)Ljava/lang/Object;
+	public static synthetic fun executeWithRetries$default (Lmisk/jdbc/TransactionRetryHandler;IJJJLkotlin/jvm/functions/Function0;ILjava/lang/Object;)Ljava/lang/Object;
+}
+
+public final class misk/jdbc/TransactionRetryHandler$Companion {
+}
+
 public final class misk/jdbc/TruncateTablesService : com/google/common/util/concurrent/AbstractIdleService {
 	public fun <init> (Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceService;Lcom/google/inject/Provider;)V
 	public fun <init> (Lkotlin/reflect/KClass;Lmisk/jdbc/DataSourceService;Lcom/google/inject/Provider;Ljava/util/List;)V