Added a setting to turn the reactive Pools into thread-local pools

Author: FroMage

extensions/reactive-datasource/runtime/src/main/java/io/quarkus/reactive/datasource/runtime/DataSourceReactiveRuntimeConfig.java

@@ -88,4 +88,10 @@ public class DataSourceReactiveRuntimeConfig {
      */
     @ConfigItem
     public PfxConfiguration keyCertificatePfx;
+
+    /**
+     * Experimental: use one connection pool per thread.
+     */
+    @ConfigItem
+    public Optional<Boolean> threadLocal;
 }

extensions/reactive-datasource/runtime/src/main/java/io/quarkus/reactive/datasource/runtime/ThreadLocalPool.java

@@ -0,0 +1,91 @@
+package io.quarkus.reactive.datasource.runtime;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.jboss.logging.Logger;
+
+import io.vertx.core.AsyncResult;
+import io.vertx.core.Handler;
+import io.vertx.core.Vertx;
+import io.vertx.sqlclient.Pool;
+import io.vertx.sqlclient.PoolOptions;
+import io.vertx.sqlclient.PreparedQuery;
+import io.vertx.sqlclient.Query;
+import io.vertx.sqlclient.Row;
+import io.vertx.sqlclient.RowSet;
+import io.vertx.sqlclient.SqlConnection;
+import io.vertx.sqlclient.Transaction;
+
+public abstract class ThreadLocalPool<PoolType extends Pool> implements Pool {
+
+    private static final Logger log = Logger.getLogger(ThreadLocalPool.class);
+
+    private final AtomicReference<ThreadLocal<PoolType>> pool = new AtomicReference<>(new ThreadLocal<>());
+    private static final List<Pool> threadLocalPools = new ArrayList<>();
+
+    protected final PoolOptions poolOptions;
+    protected final Vertx vertx;
+
+    public ThreadLocalPool(Vertx vertx, PoolOptions poolOptions) {
+        this.vertx = vertx;
+        this.poolOptions = poolOptions;
+    }
+
+    private PoolType pool() {
+        ThreadLocal<PoolType> poolThreadLocal = pool.get();
+        PoolType ret = poolThreadLocal.get();
+        if (ret == null) {
+            log.debugf("Making pool for thread: %s", Thread.currentThread());
+            ret = createThreadLocalPool();
+            synchronized (threadLocalPools) {
+                threadLocalPools.add(ret);
+            }
+            poolThreadLocal.set(ret);
+        }
+        return ret;
+    }
+
+    protected abstract PoolType createThreadLocalPool();
+
+    @Override
+    public void getConnection(Handler<AsyncResult<SqlConnection>> handler) {
+        pool().getConnection(handler);
+    }
+
+    @Override
+    public Query<RowSet<Row>> query(String sql) {
+        return pool().query(sql);
+    }
+
+    @Override
+    public PreparedQuery<RowSet<Row>> preparedQuery(String sql) {
+        return pool().preparedQuery(sql);
+    }
+
+    @Override
+    public void begin(Handler<AsyncResult<Transaction>> handler) {
+        pool().begin(handler);
+    }
+
+    /**
+     * This is a bit weird because it works on all ThreadLocal pools, but it's only
+     * called from a single thread, when doing shutdown, and needs to close all the
+     * pools and reinitialise the thread local so that all newly created pools after
+     * the restart will start with an empty thread local instead of a closed one.
+     */
+    @Override
+    public void close() {
+        // close all the thread-local pools
+        synchronized (threadLocalPools) {
+            for (Pool pool : threadLocalPools) {
+                log.debugf("Closing pool: %s", pool);
+                pool.close();
+            }
+            threadLocalPools.clear();
+        }
+        // discard the TL to clear them all
+        pool.set(new ThreadLocal<PoolType>());
+    }
+}

extensions/reactive-db2-client/runtime/src/main/java/io/quarkus/reactive/db2/client/runtime/DB2PoolRecorder.java

@@ -46,6 +46,10 @@ private DB2Pool initialize(Vertx vertx, DataSourceRuntimeConfig dataSourceRuntim
                 dataSourceReactiveDB2Config);
         DB2ConnectOptions connectOptions = toConnectOptions(dataSourceRuntimeConfig, dataSourceReactiveRuntimeConfig,
                 dataSourceReactiveDB2Config);
+        if (dataSourceReactiveRuntimeConfig.threadLocal.isPresent() &&
+                dataSourceReactiveRuntimeConfig.threadLocal.get()) {
+            return new ThreadLocalDB2Pool(vertx, connectOptions, poolOptions);
+        }
         return DB2Pool.pool(vertx, connectOptions, poolOptions);
     }
 

extensions/reactive-db2-client/runtime/src/main/java/io/quarkus/reactive/db2/client/runtime/ThreadLocalDB2Pool.java

@@ -0,0 +1,22 @@
+package io.quarkus.reactive.db2.client.runtime;
+
+import io.quarkus.reactive.datasource.runtime.ThreadLocalPool;
+import io.vertx.core.Vertx;
+import io.vertx.db2client.DB2ConnectOptions;
+import io.vertx.db2client.DB2Pool;
+import io.vertx.sqlclient.PoolOptions;
+
+public class ThreadLocalDB2Pool extends ThreadLocalPool<DB2Pool> implements DB2Pool {
+
+    private final DB2ConnectOptions db2ConnectOptions;
+
+    public ThreadLocalDB2Pool(Vertx vertx, DB2ConnectOptions db2ConnectOptions, PoolOptions poolOptions) {
+        super(vertx, poolOptions);
+        this.db2ConnectOptions = db2ConnectOptions;
+    }
+
+    @Override
+    protected DB2Pool createThreadLocalPool() {
+        return DB2Pool.pool(vertx, db2ConnectOptions, poolOptions);
+    }
+}

extensions/reactive-mysql-client/runtime/src/main/java/io/quarkus/reactive/mysql/client/runtime/MySQLPoolRecorder.java

@@ -68,6 +68,10 @@ private MySQLPool initialize(Vertx vertx, DataSourceRuntimeConfig dataSourceRunt
                 dataSourceReactiveMySQLConfig);
         MySQLConnectOptions mysqlConnectOptions = toMySQLConnectOptions(dataSourceRuntimeConfig,
                 dataSourceReactiveRuntimeConfig, dataSourceReactiveMySQLConfig);
+        if (dataSourceReactiveRuntimeConfig.threadLocal.isPresent() &&
+                dataSourceReactiveRuntimeConfig.threadLocal.get()) {
+            return new ThreadLocalMySQLPool(vertx, mysqlConnectOptions, poolOptions);
+        }
         return MySQLPool.pool(vertx, mysqlConnectOptions, poolOptions);
     }
 

extensions/reactive-mysql-client/runtime/src/main/java/io/quarkus/reactive/mysql/client/runtime/ThreadLocalMySQLPool.java

@@ -0,0 +1,22 @@
+package io.quarkus.reactive.mysql.client.runtime;
+
+import io.quarkus.reactive.datasource.runtime.ThreadLocalPool;
+import io.vertx.core.Vertx;
+import io.vertx.mysqlclient.MySQLConnectOptions;
+import io.vertx.mysqlclient.MySQLPool;
+import io.vertx.sqlclient.PoolOptions;
+
+public class ThreadLocalMySQLPool extends ThreadLocalPool<MySQLPool> implements MySQLPool {
+
+    private final MySQLConnectOptions mySQLConnectOptions;
+
+    public ThreadLocalMySQLPool(Vertx vertx, MySQLConnectOptions mySQLConnectOptions, PoolOptions poolOptions) {
+        super(vertx, poolOptions);
+        this.mySQLConnectOptions = mySQLConnectOptions;
+    }
+
+    @Override
+    protected MySQLPool createThreadLocalPool() {
+        return MySQLPool.pool(vertx, mySQLConnectOptions, poolOptions);
+    }
+}

extensions/reactive-pg-client/runtime/src/main/java/io/quarkus/reactive/pg/client/runtime/PgPoolRecorder.java

@@ -68,6 +68,10 @@ private PgPool initialize(Vertx vertx, DataSourceRuntimeConfig dataSourceRuntime
                 dataSourceReactivePostgreSQLConfig);
         PgConnectOptions pgConnectOptions = toPgConnectOptions(dataSourceRuntimeConfig, dataSourceReactiveRuntimeConfig,
                 dataSourceReactivePostgreSQLConfig);
+        if (dataSourceReactiveRuntimeConfig.threadLocal.isPresent() &&
+                dataSourceReactiveRuntimeConfig.threadLocal.get()) {
+            return new ThreadLocalPgPool(vertx, pgConnectOptions, poolOptions);
+        }
         return PgPool.pool(vertx, pgConnectOptions, poolOptions);
     }
 

extensions/reactive-pg-client/runtime/src/main/java/io/quarkus/reactive/pg/client/runtime/ThreadLocalPgPool.java

@@ -0,0 +1,22 @@
+package io.quarkus.reactive.pg.client.runtime;
+
+import io.quarkus.reactive.datasource.runtime.ThreadLocalPool;
+import io.vertx.core.Vertx;
+import io.vertx.pgclient.PgConnectOptions;
+import io.vertx.pgclient.PgPool;
+import io.vertx.sqlclient.PoolOptions;
+
+public class ThreadLocalPgPool extends ThreadLocalPool<PgPool> implements PgPool {
+
+    private final PgConnectOptions pgConnectOptions;
+
+    public ThreadLocalPgPool(Vertx vertx, PgConnectOptions pgConnectOptions, PoolOptions poolOptions) {
+        super(vertx, poolOptions);
+        this.pgConnectOptions = pgConnectOptions;
+    }
+
+    @Override
+    protected PgPool createThreadLocalPool() {
+        return PgPool.pool(vertx, pgConnectOptions, poolOptions);
+    }
+}

integration-tests/reactive-pg-client/pom.xml

@@ -38,6 +38,11 @@
             <artifactId>quarkus-junit5</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>io.quarkus</groupId>
+            <artifactId>quarkus-junit5-internal</artifactId>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <groupId>io.rest-assured</groupId>
             <artifactId>rest-assured</artifactId>
@@ -51,6 +56,11 @@
                 <directory>src/main/resources</directory>
                 <filtering>true</filtering>
             </resource>
+            <resource>
+                <directory>src/test/resources</directory>
+                <filtering>true</filtering>
+                <targetPath>../test-classes</targetPath>
+            </resource>
         </resources>
         <plugins>
             <plugin>

integration-tests/reactive-pg-client/src/test/java/io/quarkus/it/reactive/pg/client/HotReloadFruitResource.java

@@ -0,0 +1,53 @@
+package io.quarkus.it.reactive.pg.client;
+
+import java.util.concurrent.CompletionStage;
+
+import javax.annotation.PostConstruct;
+import javax.inject.Inject;
+import javax.ws.rs.GET;
+import javax.ws.rs.Path;
+import javax.ws.rs.Produces;
+import javax.ws.rs.core.MediaType;
+
+import io.vertx.core.json.JsonArray;
+import io.vertx.core.json.JsonObject;
+import io.vertx.mutiny.pgclient.PgPool;
+import io.vertx.mutiny.sqlclient.Row;
+
+@Path("/hot-fruits")
+public class HotReloadFruitResource {
+
+    @Inject
+    PgPool client;
+
+    @PostConstruct
+    void setupDb() {
+        client.query("DROP TABLE IF EXISTS fruits").execute()
+                .flatMap(r -> client.query("CREATE TABLE fruits (id SERIAL PRIMARY KEY, name TEXT NOT NULL)").execute())
+                .flatMap(r -> client.query("INSERT INTO fruits (name) VALUES ('Orange')").execute())
+                .flatMap(r -> client.query("INSERT INTO fruits (name) VALUES ('Pear')").execute())
+                .flatMap(r -> client.query("INSERT INTO fruits (name) VALUES ('Apple')").execute())
+                .await().indefinitely();
+    }
+
+    @GET
+    @Produces(MediaType.APPLICATION_JSON)
+    public CompletionStage<JsonArray> listFruits() {
+        return client.query("SELECT * FROM fruits").execute()
+                .map(pgRowSet -> {
+                    JsonArray jsonArray = new JsonArray();
+                    for (Row row : pgRowSet) {
+                        jsonArray.add(toJson(row));
+                    }
+                    return jsonArray;
+                })
+                .subscribeAsCompletionStage();
+    }
+
+    private JsonObject toJson(Row row) {
+        return new JsonObject()
+                .put("id", row.getLong("id"))
+                .put("name", row.getString("name"));
+    }
+
+}