HIVE-26794: Explore retiring TxnHandler#connPoolMutex idle connections

standalone-metastore/metastore-server/src/main/java/org/apache/hadoop/hive/metastore/datasource/DbCPDataSourceProvider.java

@@ -18,6 +18,7 @@
 package org.apache.hadoop.hive.metastore.datasource;
 
 import java.sql.SQLException;
+import java.util.Collections;
 import java.util.Map;
 
 import javax.sql.DataSource;
@@ -112,9 +113,22 @@ public DataSource create(Configuration hdpConfig, int maxPoolSize) throws SQLExc
     objectPool.setSoftMinEvictableIdleTimeMillis(softMinEvictableIdleTimeMillis);
     objectPool.setLifo(lifo);
 
+    // Enable TxnHandler#connPoolMutex to release the idle connection if possible,
+    // TxnHandler#connPoolMutex is mostly used for MutexAPI that is primarily designed to
+    // provide coarse-grained mutex support to maintenance tasks running inside the Metastore,
+    // this will make Metastore more scalable especially if there is a leader in the warehouse.
+    if ("mutex".equalsIgnoreCase(poolName)) {
+      if (timeBetweenEvictionRuns < 0) {
+        // When timeBetweenEvictionRunsMillis non-positive, no idle object evictor thread runs
+        objectPool.setTimeBetweenEvictionRunsMillis(30 * 1000);
+      }
+      if (softMinEvictableIdleTimeMillis < 0) {
+        objectPool.setSoftMinEvictableIdleTimeMillis(600 * 1000);
+      }
+    }
     String stmt = dbProduct.getPrepareTxnStmt();
     if (stmt != null) {
-      poolableConnFactory.setValidationQuery(stmt);
+      poolableConnFactory.setConnectionInitSql(Collections.singletonList(stmt));
     }
     return new PoolingDataSource(objectPool);
   }

standalone-metastore/metastore-server/src/main/java/org/apache/hadoop/hive/metastore/datasource/HikariCPDataSourceProvider.java

@@ -72,6 +72,17 @@ public DataSource create(Configuration hdpConfig, int maxPoolSize) throws SQLExc
       config.setPoolName(poolName);
     }
 
+    // It's kind of a waste to create a fixed size connection pool as same as the TxnHandler#connPool,
+    // TxnHandler#connPoolMutex is mostly used for MutexAPI that is primarily designed to
+    // provide coarse-grained mutex support to maintenance tasks running inside the Metastore,
+    // add minimumIdle=2 and idleTimeout=10min(default, can be set by hikaricp.idleTimeout) to the pool,
+    // so that the connection pool can retire the idle connection aggressively,
+    // this will make Metastore more scalable especially if there is a leader in the warehouse.
+    if ("mutex".equals(poolName)) {
+      int minimumIdle = Integer.valueOf(hdpConfig.get(HIKARI + ".minimumIdle", "2"));
+      config.setMinimumIdle(Math.min(maxPoolSize, minimumIdle));
+    }
+
     //https://github.com/brettwooldridge/HikariCP
     config.setConnectionTimeout(connectionTimeout);
 

standalone-metastore/metastore-server/src/main/java/org/apache/hadoop/hive/metastore/txn/TxnHandler.java

@@ -6188,6 +6188,12 @@ public Connection getConnection(String username, String password) throws SQLExce
         connectionProps.setProperty("user", username);
         connectionProps.setProperty("password", password);
         Connection conn = driver.connect(connString, connectionProps);
+        String prepareStmt = dbProduct != null ? dbProduct.getPrepareTxnStmt() : null;
+        if (prepareStmt != null) {
+          try (Statement stmt = conn.createStatement()) {
+            stmt.execute(prepareStmt);
+          }
+        }
         conn.setAutoCommit(false);
         return conn;
       } catch (SQLException e) {

standalone-metastore/metastore-server/src/test/java/org/apache/hadoop/hive/metastore/datasource/TestDataSourceProviderFactory.java

@@ -19,7 +19,10 @@
 
 import com.zaxxer.hikari.HikariDataSource;
 
+import com.zaxxer.hikari.HikariPoolMXBean;
 import org.apache.commons.dbcp2.PoolingDataSource;
+import org.apache.commons.lang3.reflect.MethodUtils;
+import org.apache.commons.pool2.impl.GenericObjectPool;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hive.metastore.PersistenceManagerProvider;
 import org.apache.hadoop.hive.metastore.annotation.MetastoreUnitTest;
@@ -34,6 +37,8 @@
 import javax.sql.DataSource;
 import java.sql.Connection;
 import java.sql.SQLException;
+import java.util.ArrayList;
+import java.util.List;
 
 @Category(MetastoreUnitTest.class)
 public class TestDataSourceProviderFactory {
@@ -143,6 +148,60 @@ public void testCreateDbCpDataSource() throws SQLException {
     Assert.assertTrue(ds instanceof PoolingDataSource);
   }
 
+  @Test
+  public void testEvictIdleConnection() throws Exception {
+    String[] dataSourceType = {HikariCPDataSourceProvider.HIKARI, DbCPDataSourceProvider.DBCP};
+    try (DataSourceProvider.DataSourceNameConfigurator configurator =
+             new DataSourceProvider.DataSourceNameConfigurator(conf, "mutex")) {
+      for (final String type: dataSourceType) {
+        MetastoreConf.setVar(conf, ConfVars.CONNECTION_POOLING_TYPE, type);
+        boolean isHikari = HikariCPDataSourceProvider.HIKARI.equals(type);
+        if (isHikari) {
+          conf.unset("hikaricp.connectionInitSql");
+          // The minimum of idleTimeout is 10s
+          conf.set("hikaricp.idleTimeout", "10000");
+          System.setProperty("com.zaxxer.hikari.housekeeping.periodMs", "1000");
+        } else {
+          conf.set("dbcp.timeBetweenEvictionRunsMillis", "1000");
+          conf.set("dbcp.softMinEvictableIdleTimeMillis", "3000");
+          conf.set("dbcp.maxIdle", "0");
+        }
+        DataSourceProvider dsp = DataSourceProviderFactory.tryGetDataSourceProviderOrNull(conf);
+        DataSource ds = dsp.create(conf, 5);
+        List<Connection> connections = new ArrayList<>();
+        for (int i = 0; i < 5; i++) {
+          connections.add(ds.getConnection());
+        }
+        HikariPoolMXBean poolMXBean = null;
+        GenericObjectPool objectPool = null;
+        if (isHikari) {
+          poolMXBean = ((HikariDataSource) ds).getHikariPoolMXBean();
+          Assert.assertEquals(type, 5, poolMXBean.getTotalConnections());
+          Assert.assertEquals(type, 5, poolMXBean.getActiveConnections());
+        } else {
+          objectPool = (GenericObjectPool) MethodUtils.invokeMethod(ds, true, "getPool");
+          Assert.assertEquals(type, 5, objectPool.getNumActive());
+          Assert.assertEquals(type, 5, objectPool.getMaxTotal());
+        }
+        connections.forEach(connection -> {
+          try {
+            connection.close();
+          } catch (SQLException e) {
+            throw new RuntimeException(e);
+          }
+        });
+        Thread.sleep(isHikari ? 15000 : 7000);
+        if (isHikari) {
+          Assert.assertEquals(type, 2, poolMXBean.getTotalConnections());
+          Assert.assertEquals(type, 2, poolMXBean.getIdleConnections());
+        } else {
+          Assert.assertEquals(type, 0, objectPool.getNumActive());
+          Assert.assertEquals(type, 0, objectPool.getNumIdle());
+        }
+      }
+    }
+  }
+
   @Test
   public void testClosePersistenceManagerProvider() throws Exception {
     String[] dataSourceType = {HikariCPDataSourceProvider.HIKARI, DbCPDataSourceProvider.DBCP};