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

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

@@ -122,10 +122,14 @@ BEGIN
    WHERE applicationId = applicationID_IN;
 END //
 
-CREATE PROCEDURE sp_getApplicationsHomeSubCluster()
+CREATE PROCEDURE sp_getApplicationsHomeSubCluster(IN limit_IN int)
 BEGIN
-   SELECT applicationId, homeSubCluster
-   FROM applicationsHomeSubCluster;
+   SELECT
+       applicationId, homeSubCluster
+   FROM
+   (SELECT
+        applicationId, homeSubCluster, ROW_NUMBER() OVER (ORDER BY applicationId) AS row_num
+    FROM applicationsHomeSubCluster) WHERE row_num <= limit_IN
 END //
 
 CREATE PROCEDURE sp_deleteApplicationHomeSubCluster(

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-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,14 +261,33 @@ 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");
+    }
+
+    List<ApplicationHomeSubCluster> result = new ArrayList<>();
+    List<ApplicationId> applicationIdList =
+        applications.keySet().stream().collect(Collectors.toList());
+
+    SubClusterId requestSubClusterId = request.getSubClusterId();
+    int appCount = 0;
+    for (int i = 0; i < applicationIdList.size(); i++) {
+      if (appCount >= maxAppsInStateStore) {
+        break;
+      }
+      ApplicationId applicationId = applicationIdList.get(i);
+      SubClusterId subClusterId = applications.get(applicationId);
+      // If the requestSubClusterId that needs to be filtered in the request
+      // is inconsistent with the SubClusterId in the data, continue to the next round
+      if (requestSubClusterId != null && !requestSubClusterId.equals(subClusterId)){
+        continue;
+      }
+      result.add(ApplicationHomeSubCluster.newInstance(applicationId, subClusterId));
+      appCount++;
     }
 
-    GetApplicationsHomeSubClusterResponse.newInstance(result);
+    LOG.info("requestSubClusterId = {}, appCount = {}.", requestSubClusterId, appCount);
     return GetApplicationsHomeSubClusterResponse.newInstance(result);
   }
 

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,23 @@ 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;;
+      if (request.getSubClusterId() != null) {
+        homeSubClusterIN = request.getSubClusterId().toString();
+      }
+      cstmt.setString("homeSubCluster_IN", homeSubClusterIN);
 
       // Execute the query
       long startTime = clock.getTime();

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

@@ -133,6 +133,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 +145,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,23 +260,41 @@ public GetApplicationHomeSubClusterResponse getApplicationHomeSubCluster(
   @Override
   public GetApplicationsHomeSubClusterResponse getApplicationsHomeSubCluster(
       GetApplicationsHomeSubClusterRequest request) throws YarnException {
+
+    if (request == null) {
+      throw new YarnException("Missing getApplicationsHomeSubCluster request");
+    }
+
     long start = clock.getTime();
     List<ApplicationHomeSubCluster> result = new ArrayList<>();
+    SubClusterId requestSubClusterId = request.getSubClusterId();
+    int appCount = 0;
 
     try {
-      for (String child : zkManager.getChildren(appsZNode)) {
+      List<String> childrens = zkManager.getChildren(appsZNode);
+      for (String child : childrens) {
+        if (appCount >= maxAppsInStateStore) {
+          break;
+        }
         ApplicationId appId = ApplicationId.fromString(child);
         SubClusterId homeSubCluster = getApp(appId);
-        ApplicationHomeSubCluster app =
-            ApplicationHomeSubCluster.newInstance(appId, homeSubCluster);
+        // If the requestSubClusterId that needs to be filtered in the request
+        // is inconsistent with the SubClusterId in the data, continue to the next round
+        if (requestSubClusterId != null && !requestSubClusterId.equals(homeSubCluster)) {
+          continue;
+        }
+        ApplicationHomeSubCluster app = ApplicationHomeSubCluster.newInstance(appId, homeSubCluster);
         result.add(app);
+        appCount ++;
       }
     } catch (Exception e) {
       String errMsg = "Cannot get apps: " + e.getMessage();
       FederationStateStoreUtils.logAndThrowStoreException(LOG, errMsg);
     }
     long end = clock.getTime();
     opDurations.addGetAppsHomeSubClusterDuration(start, end);
+
+    LOG.info("requestSubClusterId = {}, appCount = {}.", requestSubClusterId, appCount);
     return GetApplicationsHomeSubClusterResponse.newInstance(result);
   }
 

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);
 }

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

@@ -19,10 +19,13 @@
 
 import org.apache.hadoop.classification.InterfaceAudience.Private;
 import org.apache.hadoop.classification.InterfaceStability.Unstable;
+import org.apache.hadoop.yarn.federation.proto.YarnServerFederationProtos;
 import org.apache.hadoop.yarn.federation.proto.YarnServerFederationProtos.GetApplicationsHomeSubClusterRequestProto;
+import org.apache.hadoop.yarn.federation.proto.YarnServerFederationProtos.GetApplicationsHomeSubClusterRequestProtoOrBuilder;
 import org.apache.hadoop.yarn.server.federation.store.records.GetApplicationsHomeSubClusterRequest;
 
 import org.apache.hadoop.thirdparty.protobuf.TextFormat;
+import org.apache.hadoop.yarn.server.federation.store.records.SubClusterId;
 
 /**
  * Protocol buffer based implementation of
@@ -75,4 +78,37 @@ public String toString() {
     return TextFormat.shortDebugString(getProto());
   }
 
+  private void maybeInitBuilder() {
+    if (viaProto || builder == null) {
+      builder = GetApplicationsHomeSubClusterRequestProto.newBuilder(proto);
+    }
+    viaProto = false;
+  }
+
+  @Override
+  public SubClusterId getSubClusterId() {
+    GetApplicationsHomeSubClusterRequestProtoOrBuilder p = viaProto ? proto : builder;
+    if (!p.hasSubClusterId()) {
+      return null;
+    }
+    return convertFromProtoFormat(p.getSubClusterId());
+  }
+
+  @Override
+  public void setSubClusterId(SubClusterId subClusterId) {
+    maybeInitBuilder();
+    if (subClusterId == null) {
+      builder.clearSubClusterId();
+      return;
+    }
+    builder.setSubClusterId(convertToProtoFormat(subClusterId));
+  }
+
+  private SubClusterId convertFromProtoFormat(YarnServerFederationProtos.SubClusterIdProto sc) {
+    return new SubClusterIdPBImpl(sc);
+  }
+
+  private YarnServerFederationProtos.SubClusterIdProto convertToProtoFormat(SubClusterId sc) {
+    return ((SubClusterIdPBImpl) sc).getProto();
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-common/src/main/proto/yarn_server_federation_protos.proto

@@ -122,7 +122,7 @@ message GetApplicationHomeSubClusterResponseProto {
 }
 
 message GetApplicationsHomeSubClusterRequestProto {
-
+  optional SubClusterIdProto sub_cluster_id = 1;
 }
 
 message GetApplicationsHomeSubClusterResponseProto {