use pipelining of r2dbc to improve perf

inemtsev · inemtsev · commit 6f3c2c9dd32e · 2025-12-05T20:01:02.000+07:00
diff --git a/frameworks/Kotlin/ktor/ktor-r2dbc/build.gradle.kts b/frameworks/Kotlin/ktor/ktor-r2dbc/build.gradle.kts
@@ -11,13 +11,13 @@ plugins {
 group = "org.jetbrains.ktor"
 version = "1.0-SNAPSHOT"
 
-val ktorVersion = "3.1.3"
+val ktorVersion = "3.3.3"
 val serializationVersion = "1.8.1"
 val kotlinxHtmlVersion = "0.12.0"
 val coroutinesVersion = "1.10.1"
 val logbackVersion = "1.5.13"
-val reactorVersion = "3.7.1"
-val r2dbcVersion = "1.0.7.RELEASE"
+val reactorVersion = "3.8.0"
+val r2dbcPstgrsVersion = "1.1.1.RELEASE"
 val r2dbcPoolVersion = "1.0.2.RELEASE"
 val postgresqlVersion = "42.7.5"
 
@@ -40,7 +40,7 @@ dependencies {
     implementation("io.ktor:ktor-server-html-builder-jvm:$ktorVersion")
     implementation("io.ktor:ktor-server-netty-jvm:$ktorVersion")
 
-    implementation("org.postgresql:r2dbc-postgresql:$r2dbcVersion")
+    implementation("org.postgresql:r2dbc-postgresql:$r2dbcPstgrsVersion")
     implementation("io.r2dbc:r2dbc-pool:$r2dbcPoolVersion")
     implementation("io.projectreactor:reactor-core:$reactorVersion")
 
diff --git a/frameworks/Kotlin/ktor/ktor-r2dbc/src/main/kotlin/org/jetbrains/ktor/benchmarks/Hello.kt b/frameworks/Kotlin/ktor/ktor-r2dbc/src/main/kotlin/org/jetbrains/ktor/benchmarks/Hello.kt
@@ -15,11 +15,6 @@ import io.r2dbc.postgresql.PostgresqlConnectionFactory
 import io.r2dbc.postgresql.client.SSLMode
 import io.r2dbc.spi.Connection
 import io.r2dbc.spi.ConnectionFactory
-import kotlinx.coroutines.async
-import kotlinx.coroutines.awaitAll
-import kotlinx.coroutines.coroutineScope
-import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.reactive.awaitFirst
 import kotlinx.coroutines.reactive.awaitFirstOrNull
 import kotlinx.coroutines.reactor.awaitSingle
@@ -28,7 +23,7 @@ import reactor.core.publisher.Flux
 import reactor.core.publisher.Mono
 import java.time.Duration
 import java.util.concurrent.ThreadLocalRandom
-import kotlin.random.Random
+import kotlin.math.min
 
 const val HELLO_WORLD = "Hello, World!"
 const val WORLD_QUERY = "SELECT id, randomnumber FROM world WHERE id = $1"
@@ -54,34 +49,18 @@ fun Application.main() {
         }
 
         get("/db") {
-            val request = getWorld(dbConnFactory)
-            val result = request.awaitFirstOrNull()
-
-            call.respondJson(result)
+            val world = dbConnFactory.fetchWorld()
+            call.respondJson(world)
         }
 
         get("/queries") {
             val queries = call.queries()
-
-            val result = fetchWorldsConcurrently(dbConnFactory, queries)
-
-            call.respondJson(result)
+            val worlds = dbConnFactory.fetchWorlds(queries)
+            call.respondJson(worlds)
         }
 
         get("/fortunes") {
-            val result = mutableListOf<Fortune>()
-
-            val request = Flux.usingWhen(dbConnFactory.create(), { connection ->
-                Flux.from(connection.createStatement(FORTUNES_QUERY).execute()).flatMap { r ->
-                    Flux.from(r.map { row, _ ->
-                        Fortune(
-                            row.get(0, Int::class.java)!!, row.get(1, String::class.java)!!
-                        )
-                    })
-                }
-            }, { connection -> connection.close() })
-
-            request.collectList().awaitFirstOrNull()?.let { result.addAll(it) }
+            val result = dbConnFactory.fetchFortunes().toMutableList()
 
             result.add(Fortune(0, "Additional fortune added at request time."))
             result.sortBy { it.message }
@@ -107,7 +86,7 @@ fun Application.main() {
         get("/updates") {
             val queries = call.queries()
 
-            val worlds = fetchWorldsConcurrently(dbConnFactory, queries)
+            val worlds = dbConnFactory.fetchWorlds(queries)
             val updatedWorlds = worlds.map {
                 it.copy(randomNumber = ThreadLocalRandom.current().nextInt(1, DB_ROWS + 1))
             }.sortedBy { it.id }
@@ -137,30 +116,55 @@ fun Application.main() {
     }
 }
 
-private fun getWorld(
-    dbConnFactory: ConnectionFactory, random: ThreadLocalRandom = ThreadLocalRandom.current()
-): Mono<World> = Mono.usingWhen(dbConnFactory.create(), { connection ->
-    Mono.from(connection.createStatement(WORLD_QUERY)
-        .bind("$1", random.nextInt(DB_ROWS) + 1)
-        .execute())
-        .flatMap { result ->
-            Mono.from(result.map { row, _ ->
-                World(
-                    row.get(0, Int::class.java)
-                        ?: error("id is null"),
-                    row.get(1, Int::class.java)
-                        ?: error("randomNumber is null")
-                )
-            })
-        }
-}, Connection::close)
+private suspend fun ConnectionFactory.fetchWorld(): World =
+    Mono.usingWhen(create(), { connection ->
+        selectWorld(connection)
+    }, Connection::close).awaitSingle()
+
+private suspend fun ConnectionFactory.fetchWorlds(
+    count: Int
+): List<World> {
+    if (count <= 0) return emptyList()
+    val concurrency = min(count, 32)
+    return Mono.usingWhen(create(), { connection ->
+        Flux.range(0, count)
+            .flatMap({ selectWorldPublisher(connection) }, concurrency)
+            .collectList()
+    }, Connection::close).awaitSingle()
+}
 
-suspend fun fetchWorldsConcurrently(factory: ConnectionFactory, count: Int): List<World> =
-    coroutineScope {
-        (0 until count).map {
-            async { getWorld(factory).awaitSingle() }
-        }.awaitAll()
+private fun selectWorld(connection: Connection): Mono<World> =
+    selectWorldPublisher(connection)
+
+private fun selectWorldPublisher(connection: Connection): Mono<World> {
+    val worldId = ThreadLocalRandom.current().nextInt(1, DB_ROWS + 1)
+    return Mono.from(
+        connection.createStatement(WORLD_QUERY)
+            .bind("$1", worldId)
+            .execute()
+    ).flatMap { result ->
+        Mono.from(result.map { row, _ ->
+            World(
+                row.get(0, Int::class.java) ?: error("id is null"),
+                row.get(1, Int::class.java) ?: error("randomNumber is null")
+            )
+        })
     }
+}
+
+private suspend fun ConnectionFactory.fetchFortunes(): List<Fortune> =
+    Mono.usingWhen(create(), { connection ->
+        Flux.from(connection.createStatement(FORTUNES_QUERY).execute())
+            .flatMap { result ->
+                Flux.from(result.map { row, _ ->
+                    Fortune(
+                        row.get(0, Int::class.java) ?: error("id is null"),
+                        row.get(1, String::class.java) ?: error("message is null")
+                    )
+                })
+            }
+            .collectList()
+    }, Connection::close).awaitSingle()
 
 private fun configurePostgresR2DBC(config: ApplicationConfig): ConnectionFactory {
     val cfo = PostgresqlConnectionConfiguration.builder()
diff --git a/frameworks/Kotlin/ktor/ktor/build.gradle.kts b/frameworks/Kotlin/ktor/ktor/build.gradle.kts
@@ -11,7 +11,7 @@ plugins {
 group = "org.jetbrains.ktor"
 version = "1.0-SNAPSHOT"
 
-val ktorVersion = "3.1.3"
+val ktorVersion = "3.3.3"
 val serializationVersion = "1.8.1"
 val kotlinxHtmlVersion = "0.12.0"
 val hikariVersion = "5.1.0"
diff --git a/frameworks/Kotlin/ktor/ktor/src/main/kotlin/org/jetbrains/ktor/benchmarks/main.kt b/frameworks/Kotlin/ktor/ktor/src/main/kotlin/org/jetbrains/ktor/benchmarks/main.kt
@@ -13,9 +13,9 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.withContext
 import kotlinx.html.*
-import java.sql.Connection
 import java.util.StringJoiner
-import kotlin.random.Random
+import java.util.concurrent.ThreadLocalRandom
+import kotlinx.coroutines.CoroutineDispatcher
 
 const val HELLO_WORLD = "Hello, World!"
 const val WORLD_QUERY = "SELECT id, randomNumber FROM World WHERE id = ?"
@@ -30,7 +30,6 @@ fun Application.main() {
     // Create a dedicated dispatcher for database operations
     val databaseDispatcher = Dispatchers.IO.limitedParallelism(poolSize)
     val helloWorldContent = TextContent(HELLO_WORLD, ContentType.Text.Plain)
-    val random = Random.Default
 
     install(DefaultHeaders)
 
@@ -44,36 +43,13 @@ fun Application.main() {
         }
 
         get("/db") {
-            val world = withContext(databaseDispatcher) {
-                pool.connection.use { connection ->
-                    connection.prepareStatement(WORLD_QUERY).use { statement ->
-                        statement.setInt(1, random.nextInt(DB_ROWS) + 1)
-                        statement.executeQuery().use { rs ->
-                            rs.next()
-                            World(rs.getInt(1), rs.getInt(2))
-                        }
-                    }
-                }
-            }
+            val world = fetchWorld(pool, databaseDispatcher)
             call.respondJson(world)
         }
 
-        fun Connection.selectWorlds(queries: Int): Array<World> =
-            prepareStatement(WORLD_QUERY).use { statement ->
-                Array<World>(queries) { i ->
-                    statement.setInt(1, random.nextInt(DB_ROWS) + 1)
-                    statement.executeQuery().use { rs ->
-                        rs.next()
-                        World(rs.getInt(1), rs.getInt(2))
-                    }
-                }
-            }
-
         get("/queries") {
             val queries = call.queries()
-            val result = withContext(databaseDispatcher) {
-                pool.connection.use { it.selectWorlds(queries) }
-            }
+            val result = fetchWorlds(pool, queries, databaseDispatcher)
             call.respondJson(result)
         }
 
@@ -113,20 +89,18 @@ fun Application.main() {
 
         get("/updates") {
             val queries = call.queries()
-            val result: Array<World>
+            val result = fetchWorlds(pool, queries, databaseDispatcher)
 
             withContext(databaseDispatcher) {
                 pool.connection.use { connection ->
-                    result = connection.selectWorlds(queries)
-
                     val updateSql = StringJoiner(
                         ", ",
                         "UPDATE World SET randomNumber = temp.randomNumber FROM (VALUES ",
                         " ORDER BY 1) AS temp(id, randomNumber) WHERE temp.id = World.id"
                     )
 
-                    for (i in result.indices) {
-                        result[i].randomNumber = random.nextInt(DB_ROWS) + 1
+                    for (world in result) {
+                        world.randomNumber = ThreadLocalRandom.current().nextInt(1, DB_ROWS + 1)
                         updateSql.add("(?, ?)")
                     }
 
@@ -146,6 +120,46 @@ fun Application.main() {
     }
 }
 
+suspend fun fetchWorld(
+    pool: HikariDataSource,
+    dispatcher: CoroutineDispatcher
+): World = withContext(dispatcher) {
+    pool.connection.use { connection ->
+        fetchWorld(connection)
+    }
+}
+
+private fun fetchWorld(connection: java.sql.Connection): World =
+    connection.prepareStatement(WORLD_QUERY).use { statement ->
+        statement.setInt(1, ThreadLocalRandom.current().nextInt(1, DB_ROWS + 1))
+        statement.executeQuery().use { rs ->
+            rs.next()
+            World(rs.getInt(1), rs.getInt(2))
+        }
+    }
+
+suspend fun fetchWorlds(
+    pool: HikariDataSource,
+    queries: Int,
+    dispatcher: CoroutineDispatcher
+): Array<World> = withContext(dispatcher) {
+    if (queries <= 0) {
+        emptyArray()
+    } else {
+        pool.connection.use { connection ->
+            connection.prepareStatement(WORLD_QUERY).use { statement ->
+                Array(queries) {
+                    statement.setInt(1, ThreadLocalRandom.current().nextInt(1, DB_ROWS + 1))
+                    statement.executeQuery().use { rs ->
+                        rs.next()
+                        World(rs.getInt(1), rs.getInt(2))
+                    }
+                }
+            }
+        }
+    }
+}
+
 
 fun ApplicationCall.queries() =
     request.queryParameters["queries"]?.toIntOrNull()?.coerceIn(1, 500) ?: 1
