From 811fd23f234d07f0c23c7a6ee9d7e1917011d626 Mon Sep 17 00:00:00 2001
From: 9uapaw <gyorandr@gmail.com>
Date: Fri, 10 Sep 2021 16:59:46 +0200
Subject: [PATCH] YARN-10852. Optimise CSConfiguration
 getAllUserWeightsForQueue (#3392)

---
 .../org/apache/hadoop/conf/Configuration.java |  2 +-
 .../CapacitySchedulerConfiguration.java       | 34 +++++++++++++------
 .../scheduler/capacity/TestLeafQueue.java     |  1 +
 3 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/conf/Configuration.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/conf/Configuration.java
index e4e36a24a6..1d69a34eee 100755
--- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/conf/Configuration.java
+++ b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/conf/Configuration.java
@@ -1126,7 +1126,7 @@ private static int[] findSubVariable(String eval) {
    * @throws IllegalArgumentException when more than
    * {@link Configuration#MAX_SUBST} replacements are required
    */
-  private String substituteVars(String expr) {
+  protected String substituteVars(String expr) {
     if (expr == null) {
       return null;
     }
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/CapacitySchedulerConfiguration.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/CapacitySchedulerConfiguration.java
index 84302e57d7..2bc4077741 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/CapacitySchedulerConfiguration.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/CapacitySchedulerConfiguration.java
@@ -127,6 +127,13 @@ public class CapacitySchedulerConfiguration extends ReservationSchedulerConfigur
   @Private
   public static final String USER_SETTINGS = "user-settings";
 
+  @Private
+  public static final String USER_WEIGHT_REGEX = "\\S+\\." + USER_WEIGHT;
+
+  @Private
+  public static final Pattern USER_WEIGHT_PATTERN = Pattern.compile(
+      USER_WEIGHT_REGEX);
+
   @Private
   public static final float DEFAULT_USER_WEIGHT = 1.0f;
 
@@ -2031,18 +2038,23 @@ public void setPUOrderingPolicyUnderUtilizedPreemptionMoveReservation(
    * @return map of user weights, if they exists. Otherwise, return empty map.
    */
   public Map<String, Float> getAllUserWeightsForQueue(String queuePath) {
-    Map <String, Float> userWeights = new HashMap <String, Float>();
-    String qPathPlusPrefix =
-        getQueuePrefix(queuePath).replaceAll("\\.", "\\\\.")
-        + USER_SETTINGS + "\\.";
-    String weightKeyRegex =
-        qPathPlusPrefix + "\\S+\\." + USER_WEIGHT;
-    Map<String, String> props = getValByRegex(weightKeyRegex);
-    for (Entry<String, String> e : props.entrySet()) {
+    Map <String, Float> userWeights = new HashMap <>();
+    String qPathPlusPrefix = getQueuePrefix(queuePath) + USER_SETTINGS;
+    Map<String, String> props = getConfigurationProperties()
+        .getPropertiesWithPrefix(qPathPlusPrefix);
+
+    Map<String, String> result = new HashMap<>();
+    for(Map.Entry<String, String> item: props.entrySet()) {
+      Matcher m = USER_WEIGHT_PATTERN.matcher(item.getKey());
+      if(m.find()) {
+        result.put(item.getKey(), substituteVars(item.getValue()));
+      }
+    }
+
+    for (Entry<String, String> e : result.entrySet()) {
       String userName =
-          e.getKey().replaceFirst(qPathPlusPrefix, "")
-          .replaceFirst("\\." + USER_WEIGHT, "");
-      if (userName != null && !userName.isEmpty()) {
+          e.getKey().replaceFirst("\\." + USER_WEIGHT, "");
+      if (!userName.isEmpty()) {
         userWeights.put(userName, new Float(e.getValue()));
       }
     }
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestLeafQueue.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestLeafQueue.java
index 5bb12a4434..854eeb6e7b 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestLeafQueue.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestLeafQueue.java
@@ -1853,6 +1853,7 @@ public void testUserSpecificUserLimits() throws Exception {
         + ".user-settings.firstname.lastname."
         + CapacitySchedulerConfiguration.USER_WEIGHT,
         0.7f);
+    csConf.reinitializeConfigurationProperties();
 
     when(csContext.getClusterResource())
         .thenReturn(Resources.createResource(16 * GB, 32));