From b00f1d837b655b126748e739fcba40f88b17041f Mon Sep 17 00:00:00 2001
From: bog-walk <82039410+bog-walk@users.noreply.github.com>
Date: Tue, 22 Aug 2023 08:22:56 -0400
Subject: [PATCH] fix: EXPOSED-133 Suspend transactions blocking Hikari
 connection pool (#1837)

* fix: EXPOSED-133 Suspend transactions blocking Hikari connection pool

If the maximumPoolSize of a HikariDataSource is reached, the active connections
made through newSuspendedTransaction() are not being released back into the pool,
leading to a connection timeout exception when attempting to get more connections.

resetIfClosed(), in suspendedTransactionAsyncInternal(), calls getConnection() to
check if the transaction is closed (may happen after a repetition attempt), so
that it can be properly reset for a new attempt.

This blocks the connections being released and should not be necessary for the
first loop, as the created transaction has an open connection.

resetIfClosed() is now only accessed on subsequent repetition attempt loops.
---
 .../transactions/experimental/Suspended.kt    |  3 +-
 exposed-tests/build.gradle.kts                |  1 +
 .../sql/tests/h2/ConnectionPoolTests.kt       | 70 +++++++++++++++++++
 3 files changed, 73 insertions(+), 1 deletion(-)
 create mode 100644 exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/h2/ConnectionPoolTests.kt

diff --git a/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/transactions/experimental/Suspended.kt b/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/transactions/experimental/Suspended.kt
index 6c7dc45de6..c52b09f7f4 100644
--- a/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/transactions/experimental/Suspended.kt
+++ b/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/transactions/experimental/Suspended.kt
@@ -162,6 +162,7 @@ private fun Transaction.resetIfClosed(): Transaction {
     }
 }
 
+@Suppress("CyclomaticComplexMethod")
 private fun <T> TransactionScope.suspendedTransactionAsyncInternal(
     shouldCommit: Boolean,
     statement: suspend Transaction.() -> T
@@ -172,7 +173,7 @@ private fun <T> TransactionScope.suspendedTransactionAsyncInternal(
 
     var answer: T
     while (true) {
-        val transaction = tx.value.resetIfClosed()
+        val transaction = if (repetitions == 0) tx.value else tx.value.resetIfClosed()
 
         @Suppress("TooGenericExceptionCaught")
         try {
diff --git a/exposed-tests/build.gradle.kts b/exposed-tests/build.gradle.kts
index 8081974f1d..3841af3a7c 100644
--- a/exposed-tests/build.gradle.kts
+++ b/exposed-tests/build.gradle.kts
@@ -24,6 +24,7 @@ dependencies {
     implementation("org.apache.logging.log4j", "log4j-core", Versions.log4j2)
     implementation("junit", "junit", "4.12")
     implementation("org.hamcrest", "hamcrest-library", "1.3")
+    implementation("com.zaxxer", "HikariCP", "5.0.1")
     implementation("org.jetbrains.kotlinx", "kotlinx-coroutines-debug", Versions.kotlinCoroutines)
 
     implementation("org.testcontainers", "mysql", Versions.testContainers)
diff --git a/exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/h2/ConnectionPoolTests.kt b/exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/h2/ConnectionPoolTests.kt
new file mode 100644
index 0000000000..5f68e15bc8
--- /dev/null
+++ b/exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/h2/ConnectionPoolTests.kt
@@ -0,0 +1,70 @@
+package org.jetbrains.exposed.sql.tests.h2
+
+import com.zaxxer.hikari.HikariConfig
+import com.zaxxer.hikari.HikariDataSource
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.runBlocking
+import org.jetbrains.exposed.dao.IntEntity
+import org.jetbrains.exposed.dao.IntEntityClass
+import org.jetbrains.exposed.dao.id.EntityID
+import org.jetbrains.exposed.dao.id.IntIdTable
+import org.jetbrains.exposed.sql.Database
+import org.jetbrains.exposed.sql.SchemaUtils
+import org.jetbrains.exposed.sql.tests.TestDB
+import org.jetbrains.exposed.sql.tests.shared.assertEquals
+import org.jetbrains.exposed.sql.transactions.experimental.newSuspendedTransaction
+import org.jetbrains.exposed.sql.transactions.transaction
+import org.junit.Assume
+import org.junit.Test
+
+class ConnectionPoolTests {
+    private val hikariDataSource1 by lazy {
+        HikariDataSource(
+            HikariConfig().apply {
+                jdbcUrl = "jdbc:h2:mem:hikariDB1"
+                maximumPoolSize = 10
+            }
+        )
+    }
+
+    private val hikariDB1 by lazy {
+        Database.connect(hikariDataSource1)
+    }
+
+    @Test
+    fun testSuspendTransactionsExceedingPoolSize() {
+        Assume.assumeTrue(TestDB.H2 in TestDB.enabledInTests())
+        transaction(db = hikariDB1) {
+            SchemaUtils.create(TestTable)
+        }
+
+        val exceedsPoolSize = (hikariDataSource1.maximumPoolSize * 2 + 1).coerceAtMost(50)
+        runBlocking {
+            repeat(exceedsPoolSize) {
+                launch {
+                    newSuspendedTransaction {
+                        delay(100)
+                        TestEntity.new { testValue = "test$it" }
+                    }
+                }
+            }
+        }
+
+        transaction(db = hikariDB1) {
+            assertEquals(exceedsPoolSize, TestEntity.all().toList().count())
+
+            SchemaUtils.drop(TestTable)
+        }
+    }
+
+    object TestTable : IntIdTable("HIKARI_TESTER") {
+        val testValue = varchar("test_value", 32)
+    }
+
+    class TestEntity(id: EntityID<Int>) : IntEntity(id) {
+        companion object : IntEntityClass<TestEntity>(TestTable)
+
+        var testValue by TestTable.testValue
+    }
+}