Move hibernate to new exception handling in misk-jdbc

Author: vitaly-block

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
@@ -59,6 +59,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,
@@ -217,46 +222,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()
     }
   }
 
@@ -304,7 +276,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, "")
       }
@@ -423,68 +395,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

misk-jdbc/src/main/kotlin/misk/jdbc/TransactionRetryHandler.kt

@@ -148,3 +148,10 @@ open class DefaultExceptionClassifier : ExceptionClassifier {
   protected fun isCauseRetryable(th: Throwable): Boolean = 
     th.cause?.let { isRetryable(it) } ?: false
 }
+
+/**
+ * Exception that can be thrown by application code to force a transaction retry.
+ * This is commonly used in Hibernate-based applications.
+ */
+class RetryTransactionException(message: String? = null, cause: Throwable? = null) : 
+  RuntimeException(message, cause)