add optional db connection pool support

Author: arturmkr

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

@@ -23,7 +23,8 @@ public class DataStoreConfiguration
     private boolean runMigrationsEnabled = true;
     private Integer maxPoolSize;
 
-    public DataStoreConfiguration(String jdbcUrl, String user, String password, String driver, Integer queryHistoryHoursRetention, boolean runMigrationsEnabled, Integer maxPoolSize) {
+    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;
@@ -33,15 +34,18 @@ public DataStoreConfiguration(String jdbcUrl, String user, String password, Stri
         this.maxPoolSize = 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)
+    {
         this(jdbcUrl, user, password, driver, queryHistoryHoursRetention, runMigrationsEnabled, null);
     }
 
-    public Integer getMaxPoolSize() {
+    public Integer getMaxPoolSize()
+    {
         return this.maxPoolSize;
     }
 
-    public void setMaxPoolSize(Integer maxPoolSize) {
+    public void setMaxPoolSize(Integer maxPoolSize)
+    {
         this.maxPoolSize = maxPoolSize;
     }
 

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

@@ -26,6 +26,8 @@
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.nio.file.Path;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
@@ -41,6 +43,8 @@ public class JdbcConnectionManager
     private final ScheduledExecutorService executorService =
             Executors.newSingleThreadScheduledExecutor();
 
+    private final Map<String, HikariDataSource> pools = new ConcurrentHashMap<>();
+
     public JdbcConnectionManager(Jdbi jdbi, DataStoreConfiguration configuration)
     {
         this.jdbi = requireNonNull(jdbi, "jdbi is null")
@@ -61,21 +65,18 @@ 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))
+        Integer maxPoolSize = configuration.getMaxPoolSize();
+        if (maxPoolSize != null && maxPoolSize > 0) {
+            HikariDataSource ds = getOrCreateDataSource(routingGroupDatabase, maxPoolSize);
+            return Jdbi.create(ds)
                     .installPlugin(new SqlObjectPlugin())
                     .registerRowMapper(new RecordAndAnnotatedConstructorMapper());
         }
 
-        return Jdbi.create(buildJdbcUrl(routingGroupDatabase), configuration.getUser(), configuration.getPassword())
+        return Jdbi.create(
+                        buildJdbcUrl(routingGroupDatabase),
+                        configuration.getUser(),
+                        configuration.getPassword())
                 .installPlugin(new SqlObjectPlugin())
                 .registerRowMapper(new RecordAndAnnotatedConstructorMapper());
     }
@@ -123,11 +124,51 @@ private void startCleanUps()
         executorService.scheduleWithFixedDelay(
                 () -> {
                     log.info("Performing query history cleanup task");
-                    long created = System.currentTimeMillis() - TimeUnit.HOURS.toMillis(this.configuration.getQueryHistoryHoursRetention());
+                    long created = System.currentTimeMillis()
+                            - TimeUnit.HOURS.toMillis(this.configuration.getQueryHistoryHoursRetention());
                     jdbi.onDemand(QueryHistoryDao.class).deleteOldHistory(created);
                 },
                 1,
                 120,
                 TimeUnit.MINUTES);
     }
+
+    private HikariDataSource getOrCreateDataSource(String routingGroupDatabase, int maxPoolSize)
+    {
+        return pools.compute(routingGroupDatabase, (key, existing) -> {
+            if (existing != null && !existing.isClosed()) {
+                return existing;
+            }
+
+            HikariConfig cfg = new HikariConfig();
+            cfg.setJdbcUrl(buildJdbcUrl(key));
+            cfg.setUsername(configuration.getUser());
+            cfg.setPassword(configuration.getPassword());
+            if (configuration.getDriver() != null) {
+                cfg.setDriverClassName(configuration.getDriver());
+            }
+            cfg.setMaximumPoolSize(maxPoolSize);
+            cfg.setPoolName("gateway-ha-" + key);
+
+            return new HikariDataSource(cfg);
+        });
+    }
+
+    public void close()
+    {
+        for (Map.Entry<String, HikariDataSource> e : pools.entrySet()) {
+            HikariDataSource ds = e.getValue();
+            if (ds != null && !ds.isClosed()) {
+                try {
+                    ds.close();
+                }
+                catch (RuntimeException ex) {
+                    log.warn(ex, "Failed to close datasource for key: %s", e.getKey());
+                }
+            }
+        }
+        pools.clear();
+
+        executorService.shutdownNow();
+    }
 }

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

@@ -21,21 +21,28 @@
 import java.sql.Connection;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.concurrent.*;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
-final class TestJdbcConnectionManagerPool {
+final class TestJdbcConnectionManagerPool
+{
     @Test
-    void blocksWhenExceedingMaxPoolSize() throws Exception {
+    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
-        );
+                2);
 
         JdbcConnectionManager cm = new JdbcConnectionManager(Jdbi.create(jdbcUrl, "sa", "sa"), cfg);
         Jdbi jdbi = cm.getJdbi("testdb");
@@ -73,7 +80,8 @@ void blocksWhenExceedingMaxPoolSize() throws Exception {
             try {
                 third.get(200, TimeUnit.MILLISECONDS);
                 completedIn200ms = true; // if this happens when connection was not blocked, which is wrong
-            } catch (TimeoutException expected) {
+            }
+            catch (TimeoutException expected) {
                 // expected, means the request was blocked on the pool
             }
 
@@ -92,9 +100,10 @@ void blocksWhenExceedingMaxPoolSize() throws Exception {
         }
     }
 
-
     @Test
-    void doesNotBlockWhenMaxPoolSizeIsNull() throws Exception {
+    void doesNotBlockWhenMaxPoolSizeIsNull()
+            throws Exception
+    {
         String dbPath = Path.of(System.getProperty("java.io.tmpdir"), "h2db-nopool-" + System.currentTimeMillis()).toString();
         String jdbcUrl = "jdbc:h2:" + dbPath;
 
@@ -138,7 +147,8 @@ void doesNotBlockWhenMaxPoolSizeIsNull() throws Exception {
                 try {
                     third.get(200, TimeUnit.MILLISECONDS);
                     completedIn200ms = true;   // not blocked - expected behavior
-                } catch (TimeoutException ignore) {
+                }
+                catch (TimeoutException ignore) {
                     completedIn200ms = false;  // blocked - incorrect for no-pool case
                 }
 
@@ -156,7 +166,8 @@ void doesNotBlockWhenMaxPoolSizeIsNull() throws Exception {
                 // Release the first two connections
                 hold.countDown();
                 assertThat(third.get(3, TimeUnit.SECONDS)).isTrue();
-            } finally {
+            }
+            finally {
                 es.shutdownNow();
             }
         }