add optional db connection pool support

Author: arturmkr

gateway-ha/pom.xml

@@ -318,6 +318,12 @@
             <scope>runtime</scope>
         </dependency>
 
+        <dependency>
+            <groupId>com.zaxxer</groupId>
+            <artifactId>HikariCP</artifactId>
+            <version>7.0.2</version>
+        </dependency>
+
         <!-- Test deps -->
         <dependency>
             <groupId>com.h2database</groupId>

gateway-ha/src/main/java/io/trino/gateway/ha/config/DataStoreConfiguration.java

@@ -21,15 +21,28 @@ public class DataStoreConfiguration
     private String driver;
     private Integer queryHistoryHoursRetention = 4;
     private boolean runMigrationsEnabled = true;
+    private Integer maxPoolSize;
 
-    public DataStoreConfiguration(String jdbcUrl, String user, String password, String driver, Integer queryHistoryHoursRetention, boolean runMigrationsEnabled)
-    {
+    public DataStoreConfiguration(String jdbcUrl, String user, String password, String driver, Integer queryHistoryHoursRetention, boolean runMigrationsEnabled, Integer maxPoolSize) {
         this.jdbcUrl = jdbcUrl;
         this.user = user;
         this.password = password;
         this.driver = driver;
         this.queryHistoryHoursRetention = queryHistoryHoursRetention;
         this.runMigrationsEnabled = runMigrationsEnabled;
+        this.maxPoolSize = maxPoolSize;
+    }
+
+    public DataStoreConfiguration(String jdbcUrl, String user, String password, String driver, Integer queryHistoryHoursRetention, boolean runMigrationsEnabled) {
+        this(jdbcUrl, user, password, driver, queryHistoryHoursRetention, runMigrationsEnabled, null);
+    }
+
+    public Integer getMaxPoolSize() {
+        return this.maxPoolSize;
+    }
+
+    public void setMaxPoolSize(Integer maxPoolSize) {
+        this.maxPoolSize = maxPoolSize;
     }
 
     public DataStoreConfiguration() {}

gateway-ha/src/main/java/io/trino/gateway/ha/persistence/JdbcConnectionManager.java

@@ -14,6 +14,8 @@
 package io.trino.gateway.ha.persistence;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.zaxxer.hikari.HikariConfig;
+import com.zaxxer.hikari.HikariDataSource;
 import io.airlift.log.Logger;
 import io.trino.gateway.ha.config.DataStoreConfiguration;
 import io.trino.gateway.ha.persistence.dao.QueryHistoryDao;
@@ -59,6 +61,20 @@ public Jdbi getJdbi(@Nullable String routingGroupDatabase)
             return jdbi;
         }
 
+        if (configuration.getMaxPoolSize() != null && configuration.getMaxPoolSize() > 0) {
+            HikariConfig hikariConfig = new HikariConfig();
+            hikariConfig.setJdbcUrl(buildJdbcUrl(routingGroupDatabase));
+            hikariConfig.setUsername(configuration.getUser());
+            hikariConfig.setPassword(configuration.getPassword());
+            if (configuration.getDriver() != null) {
+                hikariConfig.setDriverClassName(configuration.getDriver());
+            }
+            hikariConfig.setMaximumPoolSize(configuration.getMaxPoolSize());
+            return Jdbi.create(new HikariDataSource(hikariConfig))
+                    .installPlugin(new SqlObjectPlugin())
+                    .registerRowMapper(new RecordAndAnnotatedConstructorMapper());
+        }
+
         return Jdbi.create(buildJdbcUrl(routingGroupDatabase), configuration.getUser(), configuration.getPassword())
                 .installPlugin(new SqlObjectPlugin())
                 .registerRowMapper(new RecordAndAnnotatedConstructorMapper());

gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestJdbcConnectionManagerPool.java

@@ -0,0 +1,151 @@
+package io.trino.gateway.ha.persistence;
+
+import io.trino.gateway.ha.config.DataStoreConfiguration;
+import org.jdbi.v3.core.Jdbi;
+import org.junit.jupiter.api.Test;
+
+import java.nio.file.Path;
+import java.sql.Connection;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.*;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+final class TestJdbcConnectionManagerPool {
+    @Test
+    void blocksWhenExceedingMaxPoolSize() throws Exception {
+        String dbPath = Path.of(System.getProperty("java.io.tmpdir"), "h2db-pool-" + System.currentTimeMillis()).toString();
+        String jdbcUrl = "jdbc:h2:" + dbPath;
+
+        DataStoreConfiguration cfg = new DataStoreConfiguration(
+                jdbcUrl, "sa", "sa", "org.h2.Driver",
+                4, true,
+                2
+        );
+
+        JdbcConnectionManager cm = new JdbcConnectionManager(Jdbi.create(jdbcUrl, "sa", "sa"), cfg);
+        Jdbi jdbi = cm.getJdbi("testdb");
+
+        try (ExecutorService es = Executors.newFixedThreadPool(3)) {
+            List<Future<Connection>> acquired = new ArrayList<>();
+
+            CountDownLatch hold = new CountDownLatch(1);
+            CountDownLatch acquiredLatch = new CountDownLatch(2);
+
+            // Open exactly maxPoolSize connections and keep them open
+            for (int i = 0; i < 2; i++) {
+                acquired.add(es.submit(() -> {
+                    try (var h = jdbi.open()) {
+                        acquiredLatch.countDown();
+                        boolean released = hold.await(10, TimeUnit.SECONDS);
+                        assertThat(released).as("hold latch should be released by the test").isTrue();
+                    }
+                    return null;
+                }));
+            }
+
+            // Wait until both connections are actually acquired (avoid race)
+            boolean bothAcquired = acquiredLatch.await(3, TimeUnit.SECONDS);
+            assertThat(bothAcquired).as("both connections should be acquired before third attempt").isTrue();
+
+            // Third attempt should block since the pool is full
+            Future<Boolean> third = es.submit(() -> {
+                var h = jdbi.open();
+                h.close();
+                return true;
+            });
+
+            boolean completedIn200ms = false;
+            try {
+                third.get(200, TimeUnit.MILLISECONDS);
+                completedIn200ms = true; // if this happens when connection was not blocked, which is wrong
+            } catch (TimeoutException expected) {
+                // expected, means the request was blocked on the pool
+            }
+
+            assertThat(completedIn200ms)
+                    .as("third getJdbi().open() should be blocked by maxPoolSize=2")
+                    .isFalse();
+
+            // Release the first two connections, the third one should complete now
+            hold.countDown();
+            assertThat(third.get(3, TimeUnit.SECONDS)).isTrue();
+
+            // Wait for the first two to finish gracefully
+            for (Future<Connection> f : acquired) {
+                f.get(3, TimeUnit.SECONDS);
+            }
+        }
+    }
+
+
+    @Test
+    void doesNotBlockWhenMaxPoolSizeIsNull() throws Exception {
+        String dbPath = Path.of(System.getProperty("java.io.tmpdir"), "h2db-nopool-" + System.currentTimeMillis()).toString();
+        String jdbcUrl = "jdbc:h2:" + dbPath;
+
+        // maxPoolSize == null  ->  no pool path
+        DataStoreConfiguration cfg = new DataStoreConfiguration(
+                jdbcUrl, "sa", "sa", "org.h2.Driver",
+                4, true);
+
+        JdbcConnectionManager cm = new JdbcConnectionManager(Jdbi.create(jdbcUrl, "sa", "sa"), cfg);
+        Jdbi jdbi = cm.getJdbi("testdb");
+
+        try (ExecutorService es = Executors.newFixedThreadPool(3)) {
+            try {
+                CountDownLatch hold = new CountDownLatch(1);
+                CountDownLatch acquiredLatch = new CountDownLatch(2);
+
+                // Open two connections and keep them open
+                for (int i = 0; i < 2; i++) {
+                    es.submit(() -> {
+                        try (var h = jdbi.open()) {
+                            acquiredLatch.countDown();
+                            boolean released = hold.await(10, TimeUnit.SECONDS);
+                            assertThat(released).isTrue();
+                        }
+                        return null;
+                    });
+                }
+
+                // Wait until both connections are really open (avoid race conditions)
+                boolean bothAcquired = acquiredLatch.await(3, TimeUnit.SECONDS);
+                assertThat(bothAcquired).isTrue();
+
+                // Third connection attempt should NOT block since no pool is used
+                Future<Boolean> third = es.submit(() -> {
+                    var h = jdbi.open();
+                    h.close();
+                    return true;
+                });
+
+                boolean completedIn200ms;
+                try {
+                    third.get(200, TimeUnit.MILLISECONDS);
+                    completedIn200ms = true;   // not blocked - expected behavior
+                } catch (TimeoutException ignore) {
+                    completedIn200ms = false;  // blocked - incorrect for no-pool case
+                }
+
+                assertThat(completedIn200ms)
+                        .as("third getJdbi().open() should NOT block when no pool is configured")
+                        .isTrue();
+
+                // check H2 session count to confirm multiple physical connections were opened
+                int sessions = jdbi.withHandle(h ->
+                        h.createQuery("SELECT COUNT(*) FROM INFORMATION_SCHEMA.SESSIONS")
+                                .mapTo(int.class)
+                                .one());
+                assertThat(sessions).isGreaterThanOrEqualTo(3);
+
+                // Release the first two connections
+                hold.countDown();
+                assertThat(third.get(3, TimeUnit.SECONDS)).isTrue();
+            } finally {
+                es.shutdownNow();
+            }
+        }
+    }
+}