YARN-11340. [Federation] Improve SQLFederationStateStore DataSource Config. (#5403)

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/conf/YarnConfiguration.java

@@ -4075,6 +4075,43 @@ public class YarnConfiguration extends Configuration {
 
   public static final int DEFAULT_FEDERATION_STATESTORE_SQL_MAXCONNECTIONS = 1;
 
+  /** Database connection pool minimum number of connections. **/
+  public static final String FEDERATION_STATESTORE_SQL_MINIMUMIDLE =
+      FEDERATION_STATESTORE_SQL_PREFIX + "minimum-idle";
+
+  /** The default value of the minimum number of connections in the database connection pool. **/
+  public static final int DEFAULT_FEDERATION_STATESTORE_SQL_MINIMUMIDLE = 1;
+
+  /** The name of the database connection pool. **/
+  public static final String FEDERATION_STATESTORE_POOL_NAME =
+      FEDERATION_STATESTORE_SQL_PREFIX + "pool-name";
+
+  /** The default name of the database connection pool. **/
+  public static final String DEFAULT_FEDERATION_STATESTORE_POOL_NAME =
+      "YARN-Federation-DataBasePool";
+
+  /** The maximum lifetime of a database connection. **/
+  public static final String FEDERATION_STATESTORE_CONN_MAX_LIFE_TIME =
+      FEDERATION_STATESTORE_SQL_PREFIX + "max-life-time";
+
+  /** Database connection maximum lifetime. **/
+  public static final long DEFAULT_FEDERATION_STATESTORE_CONN_MAX_LIFE_TIME =
+      TimeUnit.MINUTES.toMillis(30);
+
+  /** Database connection idle timeout time. **/
+  public static final String FEDERATION_STATESTORE_CONN_IDLE_TIMEOUT_TIME =
+      FEDERATION_STATESTORE_SQL_PREFIX + "idle-time-out";
+
+  public static final long DEFAULT_FEDERATION_STATESTORE_CONN_IDLE_TIMEOUT_TIME =
+      TimeUnit.MINUTES.toMillis(10);
+
+  /** Database connection timeout time. **/
+  public static final String FEDERATION_STATESTORE_CONNECTION_TIMEOUT =
+      FEDERATION_STATESTORE_SQL_PREFIX + "conn-time-out";
+
+  public static final long DEFAULT_FEDERATION_STATESTORE_CONNECTION_TIMEOUT_TIME =
+      TimeUnit.SECONDS.toMillis(10);
+
   public static final String FEDERATION_STATESTORE_MAX_APPLICATIONS =
       FEDERATION_PREFIX + "state-store.max-applications";
 

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/resources/yarn-default.xml

@@ -5193,6 +5193,45 @@
     <value>1000</value>
   </property>
 
+  <property>
+    <description>The property controls the minimum number of idle connections that
+      HikariCP tries to maintain in the pool.</description>
+    <name>yarn.federation.state-store.sql.minimum-idle</name>
+    <value>1</value>
+  </property>
+
+  <property>
+    <description>
+      Specifies the name of the connection pool used by the FederationSQLStateStore.
+    </description>
+    <name>yarn.federation.state-store.sql.pool-name</name>
+    <value>YARN-Federation-DataBasePool</value>
+  </property>
+
+  <property>
+    <description>
+      This property controls the maximum lifetime of a connection in the pool.
+    </description>
+    <name>yarn.federation.state-store.sql.max-life-time</name>
+    <value>30m</value>
+  </property>
+
+  <property>
+    <description>
+      This property controls the maximum amount of time
+      that a connection is allowed to sit idle in the pool.
+    </description>
+    <name>yarn.federation.state-store.sql.idle-time-out</name>
+    <value>10m</value>
+  </property>
+
+  <property>
+    <description>Set the maximum amount of time
+      that a client will wait for a connection from the pool.</description>
+    <name>yarn.federation.state-store.sql.conn-time-out</name>
+    <value>10s</value>
+  </property>
+
   <property>
     <description>
       Specifies the class name of the cache implementation in YARN FederationCache.

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/impl/SQLFederationStateStore.java

@@ -29,6 +29,7 @@ import java.util.ArrayList;
 import java.util.Calendar;
 import java.util.List;
 import java.util.TimeZone;
+import java.util.concurrent.TimeUnit;
 
 import org.apache.commons.lang3.StringUtils;
 import org.apache.hadoop.conf.Configuration;
@@ -224,17 +225,35 @@ public class SQLFederationStateStore implements FederationStateStore {
   @VisibleForTesting
   private Connection conn = null;
   private int maxAppsInStateStore;
+  private int minimumIdle;
+  private String dataSourcePoolName;
+  private long maxLifeTime;
+  private long idleTimeout;
+  private long connectionTimeout;
 
   protected static final Version CURRENT_VERSION_INFO = Version.newInstance(1, 1);
 
   @Override
   public void init(Configuration conf) throws YarnException {
-    driverClass =
+    // Database connection configuration
-        conf.get(YarnConfiguration.FEDERATION_STATESTORE_SQL_JDBC_CLASS,
+    driverClass = conf.get(YarnConfiguration.FEDERATION_STATESTORE_SQL_JDBC_CLASS,
-            YarnConfiguration.DEFAULT_FEDERATION_STATESTORE_SQL_JDBC_CLASS);
+        YarnConfiguration.DEFAULT_FEDERATION_STATESTORE_SQL_JDBC_CLASS);
-    maximumPoolSize =
+    maximumPoolSize = conf.getInt(YarnConfiguration.FEDERATION_STATESTORE_SQL_MAXCONNECTIONS,
-        conf.getInt(YarnConfiguration.FEDERATION_STATESTORE_SQL_MAXCONNECTIONS,
+        YarnConfiguration.DEFAULT_FEDERATION_STATESTORE_SQL_MAXCONNECTIONS);
-            YarnConfiguration.DEFAULT_FEDERATION_STATESTORE_SQL_MAXCONNECTIONS);
+    minimumIdle = conf.getInt(YarnConfiguration.FEDERATION_STATESTORE_SQL_MINIMUMIDLE,
+        YarnConfiguration.DEFAULT_FEDERATION_STATESTORE_SQL_MINIMUMIDLE);
+    dataSourcePoolName = conf.get(YarnConfiguration.FEDERATION_STATESTORE_POOL_NAME,
+        YarnConfiguration.DEFAULT_FEDERATION_STATESTORE_POOL_NAME);
+    maxLifeTime = conf.getTimeDuration(YarnConfiguration.FEDERATION_STATESTORE_CONN_MAX_LIFE_TIME,
+        YarnConfiguration.DEFAULT_FEDERATION_STATESTORE_CONN_MAX_LIFE_TIME, TimeUnit.MILLISECONDS);
+    idleTimeout = conf.getTimeDuration(
+        YarnConfiguration.FEDERATION_STATESTORE_CONN_IDLE_TIMEOUT_TIME,
+        YarnConfiguration.DEFAULT_FEDERATION_STATESTORE_CONN_IDLE_TIMEOUT_TIME,
+        TimeUnit.MILLISECONDS);
+    connectionTimeout = conf.getTimeDuration(
+        YarnConfiguration.FEDERATION_STATESTORE_CONNECTION_TIMEOUT,
+        YarnConfiguration.DEFAULT_FEDERATION_STATESTORE_CONNECTION_TIMEOUT_TIME,
+        TimeUnit.MILLISECONDS);
 
     // An helper method avoids to assign a null value to these property
     userName = conf.get(YarnConfiguration.FEDERATION_STATESTORE_SQL_USERNAME);
@@ -254,7 +273,14 @@ public class SQLFederationStateStore implements FederationStateStore {
     FederationStateStoreUtils.setPassword(dataSource, password);
     FederationStateStoreUtils.setProperty(dataSource,
         FederationStateStoreUtils.FEDERATION_STORE_URL, url);
+
     dataSource.setMaximumPoolSize(maximumPoolSize);
+    dataSource.setPoolName(dataSourcePoolName);
+    dataSource.setMinimumIdle(minimumIdle);
+    dataSource.setMaxLifetime(maxLifeTime);
+    dataSource.setIdleTimeout(idleTimeout);
+    dataSource.setConnectionTimeout(connectionTimeout);
+
     LOG.info("Initialized connection pool to the Federation StateStore database at address: {}.",
         url);
 
@@ -2010,4 +2036,9 @@ public class SQLFederationStateStore implements FederationStateStore {
       }
     }
   }
+
+  @VisibleForTesting
+  public HikariDataSource getDataSource() {
+    return dataSource;
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/test/java/org/apache/hadoop/yarn/server/federation/store/impl/TestSQLFederationStateStore.java

@@ -17,6 +17,7 @@
 
 package org.apache.hadoop.yarn.server.federation.store.impl;
 
+import com.zaxxer.hikari.HikariDataSource;
 import org.apache.hadoop.security.token.delegation.DelegationKey;
 import org.apache.hadoop.test.LambdaTestUtils;
 import org.apache.hadoop.util.Time;
@@ -627,4 +628,25 @@ public class TestSQLFederationStateStore extends FederationStateStoreBaseTest {
     String expectUpdateSQL = "select sequenceName, nextVal from sequenceTable with (updlock)";
     assertEquals(expectUpdateSQL, sqlServerDBSQL);
   }
+
+  @Test
+  public void testCheckHikariDataSourceParam() throws SQLException {
+    HikariDataSource dataSource = sqlFederationStateStore.getDataSource();
+    long maxLifeTime = dataSource.getMaxLifetime();
+    long idleTimeOut = dataSource.getIdleTimeout();
+    long connTimeOut = dataSource.getConnectionTimeout();
+    String poolName = dataSource.getPoolName();
+    int minimumIdle = dataSource.getMinimumIdle();
+    int maximumPoolSize = dataSource.getMaximumPoolSize();
+
+    // maxLifeTime 30 minute, 1800000 ms
+    assertEquals(1800000, maxLifeTime);
+    // idleTimeOut 10 minute, 600000 ms
+    assertEquals(600000, idleTimeOut);
+    // connTimeOut 10 second, 10000 ms
+    assertEquals(10000, connTimeOut);
+    assertEquals("YARN-Federation-DataBasePool", poolName);
+    assertEquals(1, minimumIdle);
+    assertEquals(1, maximumPoolSize);
+  }
 }