Skip to content

Commit f6cdea1

Browse files
vitaly-blockvitaly-block
committed
Upgrade retry handler to DB type specific
1 parent cb1a6e7 commit f6cdea1

File tree

3 files changed

+118
-9
lines changed

3 files changed

+118
-9
lines changed

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

Lines changed: 37 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,15 @@ import java.time.Duration
1515
*/
1616
class TransactionRetryHandler(
1717
private val qualifierName: String = "database",
18-
private val exceptionClassifier: ExceptionClassifier = DefaultExceptionClassifier()
18+
private val exceptionClassifier: ExceptionClassifier = DefaultExceptionClassifier(),
19+
databaseType: DataSourceType? = null
1920
) {
2021

22+
constructor(
23+
qualifierName: String = "database",
24+
databaseType: DataSourceType? = null
25+
) : this(qualifierName, DefaultExceptionClassifier(databaseType), databaseType)
26+
2127
/**
2228
* Executes a block with retry logic for transient database failures.
2329
*/
@@ -85,7 +91,9 @@ interface ExceptionClassifier {
8591
/**
8692
* Default exception classifier that handles common database retry scenarios.
8793
*/
88-
open class DefaultExceptionClassifier : ExceptionClassifier {
94+
open class DefaultExceptionClassifier(
95+
private val databaseType: DataSourceType? = null
96+
) : ExceptionClassifier {
8997

9098
override fun isRetryable(th: Throwable): Boolean {
9199
return when (th) {
@@ -98,9 +106,9 @@ open class DefaultExceptionClassifier : ExceptionClassifier {
98106

99107
protected fun isMessageRetryable(th: SQLException): Boolean =
100108
isConnectionClosed(th) ||
101-
isVitessTransactionNotFound(th) ||
102-
isCockroachRestartTransaction(th) ||
103-
isTidbWriteConflict(th)
109+
(databaseType == DataSourceType.VITESS_MYSQL && isRecoverableVitessException(th)) ||
110+
(databaseType == DataSourceType.COCKROACHDB && isRecoverableCockroachException(th)) ||
111+
(databaseType == DataSourceType.TIDB && isRecoverableTidbException(th))
104112

105113
/**
106114
* This is thrown as a raw SQLException from Hikari even though it is most certainly a recoverable exception.
@@ -109,6 +117,27 @@ open class DefaultExceptionClassifier : ExceptionClassifier {
109117
protected fun isConnectionClosed(th: SQLException): Boolean =
110118
th.message.equals("Connection is closed")
111119

120+
/**
121+
* Vitess-specific recoverable exceptions.
122+
*/
123+
protected fun isRecoverableVitessException(th: SQLException): Boolean {
124+
return isVitessTransactionNotFound(th)
125+
}
126+
127+
/**
128+
* CockroachDB-specific recoverable exceptions.
129+
*/
130+
protected fun isRecoverableCockroachException(th: SQLException): Boolean {
131+
return isCockroachRestartTransaction(th)
132+
}
133+
134+
/**
135+
* TiDB-specific recoverable exceptions.
136+
*/
137+
protected fun isRecoverableTidbException(th: SQLException): Boolean {
138+
return isTidbWriteConflict(th)
139+
}
140+
112141
/**
113142
* We get this error as a MySQLQueryInterruptedException when a tablet gracefully terminates, we just need to retry
114143
* the transaction and the new primary should handle it.
@@ -118,7 +147,7 @@ open class DefaultExceptionClassifier : ExceptionClassifier {
118147
* not found (CallerID: )
119148
* ```
120149
*/
121-
protected fun isVitessTransactionNotFound(th: SQLException): Boolean {
150+
private fun isVitessTransactionNotFound(th: SQLException): Boolean {
122151
val message = th.message
123152
return message != null &&
124153
message.contains("vttablet: rpc error") &&
@@ -132,7 +161,7 @@ open class DefaultExceptionClassifier : ExceptionClassifier {
132161
* it conflicted with another concurrent or recent transaction accessing the same data. The transaction needs to be
133162
* retried by the client." https://www.cockroachlabs.com/docs/stable/common-errors.html#restart-transaction
134163
*/
135-
protected fun isCockroachRestartTransaction(th: SQLException): Boolean {
164+
private fun isCockroachRestartTransaction(th: SQLException): Boolean {
136165
val message = th.message
137166
return th.errorCode == 40001 && message != null && message.contains("restart transaction")
138167
}
@@ -141,7 +170,7 @@ open class DefaultExceptionClassifier : ExceptionClassifier {
141170
* "Transactions in TiKV encounter write conflicts". This can happen when optimistic transaction mode is on. Conflicts
142171
* are detected during transaction commit https://docs.pingcap.com/tidb/dev/tidb-faq#error-9007-hy000-write-conflict
143172
*/
144-
protected fun isTidbWriteConflict(th: SQLException): Boolean {
173+
private fun isTidbWriteConflict(th: SQLException): Boolean {
145174
return th.errorCode == 9007
146175
}
147176

‎misk-jdbc/src/test/kotlin/misk/jdbc/TransactionRetryHandlerTest.kt‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
package misk.jdbc
2+
3+
import org.junit.jupiter.api.Test
4+
import org.junit.jupiter.api.assertThrows
5+
import java.sql.SQLException
6+
import kotlin.test.assertFalse
7+
import kotlin.test.assertTrue
8+
9+
class TransactionRetryHandlerTest {
10+
11+
@Test
12+
fun `connection closed exceptions are retryable for all database types`() {
13+
val mysqlClassifier = DefaultExceptionClassifier(DataSourceType.MYSQL)
14+
val vitessClassifier = DefaultExceptionClassifier(DataSourceType.VITESS_MYSQL)
15+
val tidbClassifier = DefaultExceptionClassifier(DataSourceType.TIDB)
16+
val cockroachClassifier = DefaultExceptionClassifier(DataSourceType.COCKROACHDB)
17+
18+
val connectionClosedException = SQLException("Connection is closed")
19+
20+
assertTrue(mysqlClassifier.isRetryable(connectionClosedException))
21+
assertTrue(vitessClassifier.isRetryable(connectionClosedException))
22+
assertTrue(tidbClassifier.isRetryable(connectionClosedException))
23+
assertTrue(cockroachClassifier.isRetryable(connectionClosedException))
24+
}
25+
26+
@Test
27+
fun `vitess transaction not found is only retryable for vitess`() {
28+
val mysqlClassifier = DefaultExceptionClassifier(DataSourceType.MYSQL)
29+
val vitessClassifier = DefaultExceptionClassifier(DataSourceType.VITESS_MYSQL)
30+
val tidbClassifier = DefaultExceptionClassifier(DataSourceType.TIDB)
31+
val cockroachClassifier = DefaultExceptionClassifier(DataSourceType.COCKROACHDB)
32+
33+
val vitessException = SQLException("vttablet: rpc error: code = Aborted desc = transaction 1572922696317821557: not found (CallerID: )")
34+
35+
assertFalse(mysqlClassifier.isRetryable(vitessException))
36+
assertTrue(vitessClassifier.isRetryable(vitessException))
37+
assertFalse(tidbClassifier.isRetryable(vitessException))
38+
assertFalse(cockroachClassifier.isRetryable(vitessException))
39+
}
40+
41+
@Test
42+
fun `cockroach restart transaction is only retryable for cockroach`() {
43+
val mysqlClassifier = DefaultExceptionClassifier(DataSourceType.MYSQL)
44+
val vitessClassifier = DefaultExceptionClassifier(DataSourceType.VITESS_MYSQL)
45+
val tidbClassifier = DefaultExceptionClassifier(DataSourceType.TIDB)
46+
val cockroachClassifier = DefaultExceptionClassifier(DataSourceType.COCKROACHDB)
47+
48+
val cockroachException = SQLException("restart transaction", "40001", 40001)
49+
50+
assertFalse(mysqlClassifier.isRetryable(cockroachException))
51+
assertFalse(vitessClassifier.isRetryable(cockroachException))
52+
assertFalse(tidbClassifier.isRetryable(cockroachException))
53+
assertTrue(cockroachClassifier.isRetryable(cockroachException))
54+
}
55+
56+
@Test
57+
fun `tidb write conflict is only retryable for tidb`() {
58+
val mysqlClassifier = DefaultExceptionClassifier(DataSourceType.MYSQL)
59+
val vitessClassifier = DefaultExceptionClassifier(DataSourceType.VITESS_MYSQL)
60+
val tidbClassifier = DefaultExceptionClassifier(DataSourceType.TIDB)
61+
val cockroachClassifier = DefaultExceptionClassifier(DataSourceType.COCKROACHDB)
62+
63+
val tidbException = SQLException("write conflict", "HY000", 9007)
64+
65+
assertFalse(mysqlClassifier.isRetryable(tidbException))
66+
assertFalse(vitessClassifier.isRetryable(tidbException))
67+
assertTrue(tidbClassifier.isRetryable(tidbException))
68+
assertFalse(cockroachClassifier.isRetryable(tidbException))
69+
}
70+
71+
@Test
72+
fun `transaction retry handler uses database type correctly`() {
73+
val vitessHandler = TransactionRetryHandler("vitess", DataSourceType.VITESS_MYSQL)
74+
val tidbHandler = TransactionRetryHandler("tidb", DataSourceType.TIDB)
75+
76+
// Test that we can create handlers with database types
77+
// This primarily tests that the constructors work correctly
78+
assertTrue(vitessHandler.toString().isNotEmpty())
79+
assertTrue(tidbHandler.toString().isNotEmpty())
80+
}
81+
}

‎misk-jooq/src/main/kotlin/misk/jooq/JooqTransacter.kt‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,6 @@ class JooqTransacter @JvmOverloads constructor(
131131
DataSourceType.VITESS_MYSQL -> SQLDialect.MYSQL
132132
DataSourceType.POSTGRESQL -> SQLDialect.POSTGRES
133133
DataSourceType.TIDB -> SQLDialect.MYSQL
134-
DataSourceType.COCKROACHDB -> throw IllegalArgumentException("CockroachDB is not supported in JOOQ")
135134
else -> throw IllegalArgumentException("no SQLDialect for " + this.name)
136135
}
137136

0 commit comments

Comments
 (0)