YARN-11290. Improve Query Condition of FederationStateStore#getApplicationsHomeSubCluster.

hadoop-yarn-project/hadoop-yarn/bin/FederationStateStore/MySQL/FederationStateStoreStoredProcs.sql

@@ -122,10 +122,17 @@ BEGIN
    WHERE applicationId = applicationID_IN;
 END //
 
-CREATE PROCEDURE sp_getApplicationsHomeSubCluster()
+CREATE PROCEDURE sp_getApplicationsHomeSubCluster(IN limit_IN int, IN homeSubCluster_IN varchar(256))
 BEGIN
-   SELECT applicationId, homeSubCluster
-   FROM applicationsHomeSubCluster;
+   SET @row_num = 0;
+   SELECT
+       applicationId, homeSubCluster
+   FROM (SELECT
+         applicationId, homeSubCluster,  (@rownum := @rownum + 1) AS row_num
+        FROM applicationsHomeSubCluster) AS t
+        WHERE row_num <= limit_IN
+    AND (CASE WHEN homeSubCluster_IN IS NULL THEN 1 = 1
+              WHEN homeSubCluster_IN IS NOT NULL THEN homeSubCluster = homeSubCluster_IN END);
 END //
 
 CREATE PROCEDURE sp_deleteApplicationHomeSubCluster(

hadoop-yarn-project/hadoop-yarn/bin/FederationStateStore/SQLServer/FederationStateStoreStoreProcs.sql

@@ -111,12 +111,27 @@ IF OBJECT_ID ( '[sp_getApplicationsHomeSubCluster]', 'P' ) IS NOT NULL
 GO
 
 CREATE PROCEDURE [dbo].[sp_getApplicationsHomeSubCluster]
+    @limit int,
+    @homeSubCluster VARCHAR(256)
 AS BEGIN
     DECLARE @errorMessage nvarchar(4000)
 
     BEGIN TRY
-        SELECT [applicationId], [homeSubCluster], [createTime]
-        FROM [dbo].[applicationsHomeSubCluster]
+        SELECT
+            [applicationId],
+            [homeSubCluster],
+            [createTime]
+        FROM
+        (SELECT
+            [applicationId],
+            [homeSubCluster],
+            [createTime],
+            row_number() over(partition by [homeSubCluster] order by [createTime] desc) as row_num
+        FROM [dbo].[applicationsHomeSubCluster]) AS t
+        WHERE row_num <= @limit
+          AND (CASE WHEN @homeSubCluster IS NULL THEN 1
+                    WHEN @homeSubCluster IS NOT NULL AND [homeSubCluster] = @homeSubCluster THEN 1
+               ELSE 0 END) = 1
     END TRY
 
     BEGIN CATCH

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

@@ -4004,6 +4004,11 @@ public static boolean isAclEnabled(Configuration conf) {
 
   public static final int DEFAULT_FEDERATION_STATESTORE_SQL_MAXCONNECTIONS = 1;
 
+  public static final String FEDERATION_STATESTORE_MAX_APPLICATIONS =
+      FEDERATION_PREFIX + "state-store.max-applications";
+
+  public static final int DEFAULT_FEDERATION_STATESTORE_MAX_APPLICATIONS = 1000;
+
   public static final String ROUTER_PREFIX = YARN_PREFIX + "router.";
 
   public static final String ROUTER_BIND_HOST = ROUTER_PREFIX + "bind-host";

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

@@ -4963,4 +4963,14 @@
     </description>
   </property>
 
+  <property>
+    <name>yarn.federation.state-store.max-applications</name>
+    <value>1000</value>
+    <description>
+      yarn federation state-store supports
+      querying the maximum number of apps
+      Default is 1000
+    </description>
+  </property>
+
 </configuration>

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

@@ -24,10 +24,12 @@
 import java.util.Map.Entry;
 import java.util.TimeZone;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
 
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.hadoop.yarn.api.records.ReservationId;
+import org.apache.hadoop.yarn.conf.YarnConfiguration;
 import org.apache.hadoop.yarn.exceptions.YarnException;
 import org.apache.hadoop.yarn.server.federation.store.FederationStateStore;
 import org.apache.hadoop.yarn.server.federation.store.records.AddApplicationHomeSubClusterRequest;
@@ -90,6 +92,7 @@ public class MemoryFederationStateStore implements FederationStateStore {
   private Map<ApplicationId, SubClusterId> applications;
   private Map<ReservationId, SubClusterId> reservations;
   private Map<String, SubClusterPolicyConfiguration> policies;
+  private int maxAppsInStateStore;
 
   private final MonotonicClock clock = new MonotonicClock();
 
@@ -102,6 +105,9 @@ public void init(Configuration conf) {
     applications = new ConcurrentHashMap<ApplicationId, SubClusterId>();
     reservations = new ConcurrentHashMap<ReservationId, SubClusterId>();
     policies = new ConcurrentHashMap<String, SubClusterPolicyConfiguration>();
+    maxAppsInStateStore = conf.getInt(
+        YarnConfiguration.FEDERATION_STATESTORE_MAX_APPLICATIONS,
+        YarnConfiguration.DEFAULT_FEDERATION_STATESTORE_MAX_APPLICATIONS);
   }
 
   @Override
@@ -255,17 +261,36 @@ public GetApplicationHomeSubClusterResponse getApplicationHomeSubCluster(
   @Override
   public GetApplicationsHomeSubClusterResponse getApplicationsHomeSubCluster(
       GetApplicationsHomeSubClusterRequest request) throws YarnException {
-    List<ApplicationHomeSubCluster> result =
-        new ArrayList<ApplicationHomeSubCluster>();
-    for (Entry<ApplicationId, SubClusterId> e : applications.entrySet()) {
-      result
-          .add(ApplicationHomeSubCluster.newInstance(e.getKey(), e.getValue()));
+
+    if (request == null) {
+      throw new YarnException("Missing getApplicationsHomeSubCluster request");
     }
 
-    GetApplicationsHomeSubClusterResponse.newInstance(result);
+    SubClusterId requestSC = request.getSubClusterId();
+    List<ApplicationHomeSubCluster> result = applications.keySet().stream()
+        .map(applicationId -> generateAppHomeSC(applicationId))
+        .filter(appHomeSC -> judgeAdd(requestSC, appHomeSC.getHomeSubCluster()))
+        .limit(maxAppsInStateStore)
+        .collect(Collectors.toList());
+
+    LOG.info("filterSubClusterId = {}, appCount = {}.", requestSC, result.size());
     return GetApplicationsHomeSubClusterResponse.newInstance(result);
   }
 
+  private ApplicationHomeSubCluster generateAppHomeSC(ApplicationId applicationId) {
+    SubClusterId subClusterId = applications.get(applicationId);
+    return ApplicationHomeSubCluster.newInstance(applicationId, subClusterId);
+  }
+
+  private boolean judgeAdd(SubClusterId filterSubCluster, SubClusterId homeSubCluster) {
+    if (filterSubCluster == null) {
+      return true;
+    } else if (filterSubCluster.equals(homeSubCluster)) {
+      return true;
+    }
+    return false;
+  }
+
   @Override
   public DeleteApplicationHomeSubClusterResponse deleteApplicationHomeSubCluster(
       DeleteApplicationHomeSubClusterRequest request) throws YarnException {

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

@@ -129,7 +129,7 @@ public class SQLFederationStateStore implements FederationStateStore {
       "{call sp_getApplicationHomeSubCluster(?, ?)}";
 
   private static final String CALL_SP_GET_APPLICATIONS_HOME_SUBCLUSTER =
-      "{call sp_getApplicationsHomeSubCluster()}";
+      "{call sp_getApplicationsHomeSubCluster(?, ?)}";
 
   private static final String CALL_SP_SET_POLICY_CONFIGURATION =
       "{call sp_setPolicyConfiguration(?, ?, ?, ?)}";
@@ -154,6 +154,7 @@ public class SQLFederationStateStore implements FederationStateStore {
   private final Clock clock = new MonotonicClock();
   @VisibleForTesting
   Connection conn = null;
+  private int maxAppsInStateStore;
 
   @Override
   public void init(Configuration conf) throws YarnException {
@@ -193,6 +194,10 @@ public void init(Configuration conf) throws YarnException {
       FederationStateStoreUtils.logAndThrowRetriableException(LOG,
           "Not able to get Connection", e);
     }
+
+    maxAppsInStateStore = conf.getInt(
+        YarnConfiguration.FEDERATION_STATESTORE_MAX_APPLICATIONS,
+        YarnConfiguration.DEFAULT_FEDERATION_STATESTORE_MAX_APPLICATIONS);
   }
 
   @Override
@@ -726,13 +731,24 @@ public GetApplicationHomeSubClusterResponse getApplicationHomeSubCluster(
   @Override
   public GetApplicationsHomeSubClusterResponse getApplicationsHomeSubCluster(
       GetApplicationsHomeSubClusterRequest request) throws YarnException {
+
+    if (request == null) {
+      throw new YarnException("Missing getApplicationsHomeSubCluster request");
+    }
+
     CallableStatement cstmt = null;
     ResultSet rs = null;
-    List<ApplicationHomeSubCluster> appsHomeSubClusters =
-        new ArrayList<ApplicationHomeSubCluster>();
+    List<ApplicationHomeSubCluster> appsHomeSubClusters = new ArrayList<>();
 
     try {
       cstmt = getCallableStatement(CALL_SP_GET_APPLICATIONS_HOME_SUBCLUSTER);
+      cstmt.setInt("limit_IN", maxAppsInStateStore);
+      String homeSubClusterIN = null;
+      SubClusterId subClusterId = request.getSubClusterId();
+      if (subClusterId != null) {
+        homeSubClusterIN = subClusterId.toString();
+      }
+      cstmt.setString("homeSubCluster_IN", homeSubClusterIN);
 
       // Execute the query
       long startTime = clock.getTime();
@@ -742,8 +758,8 @@ public GetApplicationsHomeSubClusterResponse getApplicationsHomeSubCluster(
       while (rs.next()) {
 
         // Extract the output for each tuple
-        String applicationId = rs.getString(1);
-        String homeSubCluster = rs.getString(2);
+        String applicationId = rs.getString("applicationId");
+        String homeSubCluster = rs.getString("homeSubCluster");
 
         appsHomeSubClusters.add(ApplicationHomeSubCluster.newInstance(
             ApplicationId.fromString(applicationId),

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

@@ -24,6 +24,7 @@
 import java.util.Calendar;
 import java.util.List;
 import java.util.TimeZone;
+import java.util.stream.Collectors;
 
 import org.apache.hadoop.classification.VisibleForTesting;
 import org.apache.hadoop.conf.Configuration;
@@ -133,6 +134,7 @@ public class ZookeeperFederationStateStore implements FederationStateStore {
   private String membershipZNode;
   private String policiesZNode;
   private String reservationsZNode;
+  private int maxAppsInStateStore;
 
   private volatile Clock clock = SystemClock.getInstance();
 
@@ -144,6 +146,10 @@ public class ZookeeperFederationStateStore implements FederationStateStore {
   public void init(Configuration conf) throws YarnException {
     LOG.info("Initializing ZooKeeper connection");
 
+    maxAppsInStateStore = conf.getInt(
+       YarnConfiguration.FEDERATION_STATESTORE_MAX_APPLICATIONS,
+       YarnConfiguration.DEFAULT_FEDERATION_STATESTORE_MAX_APPLICATIONS);
+
     baseZNode = conf.get(
         YarnConfiguration.FEDERATION_STATESTORE_ZK_PARENT_PATH,
         YarnConfiguration.DEFAULT_FEDERATION_STATESTORE_ZK_PARENT_PATH);
@@ -255,24 +261,52 @@ public GetApplicationHomeSubClusterResponse getApplicationHomeSubCluster(
   @Override
   public GetApplicationsHomeSubClusterResponse getApplicationsHomeSubCluster(
       GetApplicationsHomeSubClusterRequest request) throws YarnException {
-    long start = clock.getTime();
-    List<ApplicationHomeSubCluster> result = new ArrayList<>();
+
+    if (request == null) {
+      throw new YarnException("Missing getApplicationsHomeSubCluster request");
+    }
 
     try {
-      for (String child : zkManager.getChildren(appsZNode)) {
-        ApplicationId appId = ApplicationId.fromString(child);
-        SubClusterId homeSubCluster = getApp(appId);
-        ApplicationHomeSubCluster app =
-            ApplicationHomeSubCluster.newInstance(appId, homeSubCluster);
-        result.add(app);
-      }
+      long start = clock.getTime();
+      SubClusterId requestSC = request.getSubClusterId();
+      List<String> children = zkManager.getChildren(appsZNode);
+      List<ApplicationHomeSubCluster> result =
+          children.stream().map(child -> generateAppHomeSC(child))
+          .filter(appHomeSC -> judgeAdd(requestSC, appHomeSC.getHomeSubCluster()))
+          .limit(maxAppsInStateStore)
+          .collect(Collectors.toList());
+      long end = clock.getTime();
+      opDurations.addGetAppsHomeSubClusterDuration(start, end);
+      LOG.info("filterSubClusterId = {}, appCount = {}.", requestSC, result.size());
+      return GetApplicationsHomeSubClusterResponse.newInstance(result);
     } catch (Exception e) {
       String errMsg = "Cannot get apps: " + e.getMessage();
       FederationStateStoreUtils.logAndThrowStoreException(LOG, errMsg);
     }
-    long end = clock.getTime();
-    opDurations.addGetAppsHomeSubClusterDuration(start, end);
-    return GetApplicationsHomeSubClusterResponse.newInstance(result);
+
+    throw new YarnException("Cannot get app by request");
+  }
+
+  private ApplicationHomeSubCluster generateAppHomeSC(String appId) {
+    try {
+      ApplicationId applicationId = ApplicationId.fromString(appId);
+      SubClusterId homeSubCluster = getApp(applicationId);
+      ApplicationHomeSubCluster app =
+          ApplicationHomeSubCluster.newInstance(applicationId, homeSubCluster);
+      return app;
+    } catch (Exception ex) {
+      LOG.error("get homeSubCluster by appId = {}.", appId);
+    }
+    return null;
+  }
+
+  private boolean judgeAdd(SubClusterId filterSubCluster, SubClusterId homeSubCluster) {
+    if (filterSubCluster == null) {
+      return true;
+    } else if (filterSubCluster.equals(homeSubCluster)) {
+      return true;
+    }
+    return false;
   }
 
   @Override

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/java/org/apache/hadoop/yarn/server/federation/store/records/GetApplicationsHomeSubClusterRequest.java

@@ -17,6 +17,7 @@
 
 package org.apache.hadoop.yarn.server.federation.store.records;
 
+import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.classification.InterfaceAudience.Private;
 import org.apache.hadoop.classification.InterfaceStability.Unstable;
 import org.apache.hadoop.yarn.util.Records;
@@ -37,4 +38,33 @@ public static GetApplicationsHomeSubClusterRequest newInstance() {
     return request;
   }
 
+  @Private
+  @Unstable
+  public static GetApplicationsHomeSubClusterRequest
+      newInstance(SubClusterId subClusterId) {
+    GetApplicationsHomeSubClusterRequest request =
+        Records.newRecord(GetApplicationsHomeSubClusterRequest.class);
+    request.setSubClusterId(subClusterId);
+    return request;
+  }
+
+  /**
+   * Get the {@link SubClusterId} representing the unique identifier of the
+   * subcluster.
+   *
+   * @return the subcluster identifier
+   */
+  @InterfaceAudience.Public
+  @Unstable
+  public abstract SubClusterId getSubClusterId();
+
+  /**
+   * Set the {@link SubClusterId} representing the unique identifier of the
+   * subcluster.
+   *
+   * @param subClusterId the subcluster identifier
+   */
+  @InterfaceAudience.Public
+  @Unstable
+  public abstract void setSubClusterId(SubClusterId subClusterId);
 }