YARN-10415. Create a group matcher which checks ALL groups of the use…

…r. Contributed by Gergely Pollak.

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/placement/CSMappingPlacementRule.java

@@ -156,23 +156,24 @@ public boolean initialize(ResourceScheduler scheduler) throws IOException {
     return mappingRules.size() > 0;
   }
 
-  private String getPrimaryGroup(String user) throws IOException {
-    return groups.getGroupsSet(user).iterator().next();
-  }
-
   /**
-   * Traverse all secondary groups (as there could be more than one
-   * and position is not guaranteed) and ensure there is queue with
-   * the same name.
+   * Sets group related data for the provided variable context.
+   * Primary group is the first group returned by getGroups.
+   * To determine secondary group we traverse all groups
+   * (as there could be more than one and position is not guaranteed) and
+   * ensure there is queue with the same name.
+   * This method also sets the groups set for the variable context for group
+   * matching.
+   * @param vctx Variable context to be updated
    * @param user Name of the user
-   * @return Name of the secondary group if found, null otherwise
    * @throws IOException
    */
-  private String getSecondaryGroup(String user) throws IOException {
+  private void setupGroupsForVariableContext(VariableContext vctx, String user)
+      throws IOException {
     Set<String> groupsSet = groups.getGroupsSet(user);
     String secondaryGroup = null;
     Iterator<String> it = groupsSet.iterator();
-    it.next();
+    String primaryGroup = it.next();
     while (it.hasNext()) {
       String group = it.next();
       if (this.queueManager.getQueue(group) != null) {
@@ -185,7 +186,10 @@ private String getSecondaryGroup(String user) throws IOException {
       LOG.debug("User {} is not associated with any Secondary " +
           "Group. Hence it may use the 'default' queue", user);
     }
-    return secondaryGroup;
+
+    vctx.put("%primary_group", primaryGroup);
+    vctx.put("%secondary_group", secondaryGroup);
+    vctx.putExtraDataset("groups", groupsSet);
   }
 
   private VariableContext createVariableContext(
@@ -195,9 +199,8 @@ private VariableContext createVariableContext(
     vctx.put("%user", user);
     vctx.put("%specified", asc.getQueue());
     vctx.put("%application", asc.getApplicationName());
-    vctx.put("%primary_group", getPrimaryGroup(user));
-    vctx.put("%secondary_group", getSecondaryGroup(user));
     vctx.put("%default", "root.default");
+    setupGroupsForVariableContext(vctx, user);
 
     vctx.setImmutables(immutableVariables);
     return vctx;

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/placement/MappingRule.java

@@ -109,7 +109,7 @@ public static MappingRule createLegacyRule(
       matcher = MappingRuleMatchers.createUserMatcher(source);
       break;
     case GROUP_MAPPING:
-      matcher = MappingRuleMatchers.createGroupMatcher(source);
+      matcher = MappingRuleMatchers.createUserGroupMatcher(source);
       break;
     case APPLICATION_MAPPING:
       matcher = MappingRuleMatchers.createApplicationNameMatcher(source);

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/placement/MappingRuleMatchers.java

@@ -19,6 +19,7 @@
 package org.apache.hadoop.yarn.server.resourcemanager.placement;
 
 import java.util.Arrays;
+import java.util.Set;
 
 /**
  * This class contains all the matcher and some helper methods to generate them.
@@ -96,6 +97,48 @@ public String toString() {
     }
   }
 
+  /**
+   * The GroupMatcher will check if any of the user's groups match the provided
+   * group name. It does not care if it's primary or secondary group, it just
+   * checks if the user is member of the expected group.
+   */
+  public static class UserGroupMatcher implements MappingRuleMatcher {
+    /**
+     * The group which should match the users's groups.
+     */
+    private String group;
+
+    UserGroupMatcher(String value) {
+      this.group = value;
+    }
+
+    /**
+     * The method will match (return true) if the user is in the provided group.
+     * This matcher expect an extraVariableSet to be present in the variable
+     * context, if it's not present, we return false.
+     * If the expected group is null we always return false.
+     * @param variables The variable context, which contains all the variables
+     * @return true if user is member of the group
+     */
+    @Override
+    public boolean match(VariableContext variables) {
+      Set<String> groups = variables.getExtraDataset("groups");
+
+      if (group == null || groups == null) {
+        return false;
+      }
+
+      String substituted = variables.replaceVariables(group);
+      return groups.contains(substituted);
+    }
+
+    @Override
+    public String toString() {
+      return "GroupMatcher{" +
+          "group='" + group + '\'' +
+          '}';
+    }
+  }
   /**
    * AndMatcher is a basic boolean matcher which takes multiple other
    * matcher as it's arguments, and on match it checks if all of them are true.
@@ -193,13 +236,13 @@ public static MappingRuleMatcher createUserMatcher(String userName) {
   }
 
   /**
-   * Convenience method to create a variable matcher which matches against the
-   * user's primary group.
+   * Convenience method to create a group matcher which matches against the
+   * groups of the user.
    * @param groupName The groupName to be matched
-   * @return VariableMatcher with %primary_group as the variable
+   * @return UserGroupMatcher
    */
-  public static MappingRuleMatcher createGroupMatcher(String groupName) {
-    return new VariableMatcher("%primary_group", groupName);
+  public static MappingRuleMatcher createUserGroupMatcher(String groupName) {
+    return new UserGroupMatcher(groupName);
   }
 
   /**
@@ -215,7 +258,7 @@ public static MappingRuleMatcher createUserGroupMatcher(
       String userName, String groupName) {
     return new AndMatcher(
         createUserMatcher(userName),
-        createGroupMatcher(groupName));
+        createUserGroupMatcher(groupName));
   }
 
   /**

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/placement/VariableContext.java

@@ -44,6 +44,12 @@ public class VariableContext {
    */
   private Set<String> immutableNames;
 
+  /**
+   * Some matchers may need to find a data in a set, which is not usable
+   * as a variable in substitutions, this store is for those sets.
+   */
+  private Map<String, Set<String>> extraDataset = new HashMap<>();
+
   /**
    * Checks if the provided variable is immutable.
    * @param name Name of the variable to check
@@ -114,6 +120,31 @@ public String get(String name) {
     return ret == null ? "" : ret;
   }
 
+  /**
+   * Adds a set to the context, each name can only be added once. The extra
+   * dataset is different from the regular variables because it cannot be
+   * referenced via tokens in the paths or any other input. However matchers
+   * and actions can explicitly access these datasets and can make decisions
+   * based on them.
+   * @param name Name which can be used to reference the collection
+   * @param set The dataset to be stored
+   */
+  public void putExtraDataset(String name, Set<String> set) {
+    if (extraDataset.containsKey(name)) {
+      throw new IllegalStateException(
+          "Dataset '" + name + "' is already set!");
+    }
+    extraDataset.put(name, set);
+  }
+
+  /**
+   * Returns the dataset referenced by the name.
+   * @param name Name of the set to be returned.
+   */
+  public Set<String> getExtraDataset(String name) {
+    return extraDataset.get(name);
+  }
+
   /**
    * Check if a variable is part of the context.
    * @param name Name of the variable to be checked
@@ -195,5 +226,4 @@ public String replacePathVariables(String input) {
 
     return String.join(".", parts);
   }
-
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/placement/TestCSMappingPlacementRule.java

@@ -411,4 +411,37 @@ public void testAllowCreateFlag() throws IOException {
         engine, app, "charlie", "root.man.create");
 
   }
+
+  private MappingRule createGroupMapping(String group, String queue) {
+    MappingRuleMatcher matcher = MappingRuleMatchers.createUserGroupMatcher(group);
+    MappingRuleAction action =
+        (new MappingRuleActions.PlaceToQueueAction(queue, true))
+        .setFallbackReject();
+    return new MappingRule(matcher, action);
+  }
+
+  @Test
+  public void testGroupMatching() throws IOException {
+    ArrayList<MappingRule> rules = new ArrayList<>();
+
+    rules.add(createGroupMapping("p_alice", "root.man.p_alice"));
+    rules.add(createGroupMapping("developer", "root.man.developer"));
+
+    //everybody is in the user group, this should catch all
+    rules.add(createGroupMapping("user", "root.man.user"));
+
+    CSMappingPlacementRule engine = setupEngine(true, rules);
+    ApplicationSubmissionContext app = createApp("app");
+
+    assertPlace(
+        "Alice should be placed to root.man.p_alice based on her primary group",
+        engine, app, "alice", "root.man.p_alice");
+    assertPlace(
+        "Bob should be placed to root.man.developer based on his developer " +
+        "group", engine, app, "bob", "root.man.developer");
+    assertPlace(
+        "Charlie should be placed to root.man.user because he is not a " +
+        "developer nor in the p_alice group", engine, app, "charlie",
+        "root.man.user");
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/placement/TestMappingRule.java

@@ -21,6 +21,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
+import com.google.common.collect.Sets;
 import org.apache.hadoop.util.StringUtils;
 import org.junit.Test;
 
@@ -107,6 +108,7 @@ void evaluateLegacyStringTestcase(
   public void testLegacyEvaluation() {
     VariableContext matching = setupVariables(
         "bob", "developer", "users", "MR");
+    matching.putExtraDataset("groups", Sets.newHashSet("developer"));
     VariableContext mismatching = setupVariables(
         "joe", "tester", "admins", "Spark");
 

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/placement/TestMappingRuleMatchers.java

@@ -18,6 +18,7 @@
 
 package org.apache.hadoop.yarn.server.resourcemanager.placement;
 
+import com.google.common.collect.Sets;
 import junit.framework.TestCase;
 import org.junit.Test;
 
@@ -41,19 +42,21 @@ public void testVariableMatcher() {
     matchingContext.put("%primary_group", "developers");
     matchingContext.put("%application", "TurboMR");
     matchingContext.put("%custom", "Matching string");
+    matchingContext.putExtraDataset("groups", Sets.newHashSet("developers"));
 
     VariableContext mismatchingContext = new VariableContext();
     mismatchingContext.put("%user", "dave");
     mismatchingContext.put("%primary_group", "testers");
     mismatchingContext.put("%application", "Tester APP");
     mismatchingContext.put("%custom", "Not matching string");
+    mismatchingContext.putExtraDataset("groups", Sets.newHashSet("testers"));
 
     VariableContext emptyContext = new VariableContext();
 
     Map<String, MappingRuleMatcher> matchers = new HashMap<>();
     matchers.put("User matcher", MappingRuleMatchers.createUserMatcher("bob"));
     matchers.put("Group matcher",
-        MappingRuleMatchers.createGroupMatcher("developers"));
+        MappingRuleMatchers.createUserGroupMatcher("developers"));
     matchers.put("Application name matcher",
         MappingRuleMatchers.createApplicationNameMatcher("TurboMR"));
     matchers.put("Custom matcher",
@@ -184,16 +187,17 @@ public void testBoolOperatorMatchers() {
     VariableContext developerBob = new VariableContext();
     developerBob.put("%user", "bob");
     developerBob.put("%primary_group", "developers");
-
+    developerBob.putExtraDataset("groups", Sets.newHashSet("developers"));
 
     VariableContext testerBob = new VariableContext();
     testerBob.put("%user", "bob");
     testerBob.put("%primary_group", "testers");
+    testerBob.putExtraDataset("groups", Sets.newHashSet("testers"));
 
     VariableContext testerDave = new VariableContext();
     testerDave.put("%user", "dave");
     testerDave.put("%primary_group", "testers");
-
+    testerDave.putExtraDataset("groups", Sets.newHashSet("testers"));
 
     VariableContext accountantDave = new VariableContext();
     accountantDave.put("%user", "dave");
@@ -252,4 +256,56 @@ public void testToStrings() {
         ", " + var.toString() + "]}", or.toString());
   }
 
+  @Test
+  public void testGroupMatching() {
+    VariableContext letterGroups = new VariableContext();
+    letterGroups.putExtraDataset("groups", Sets.newHashSet("a", "b", "c"));
+
+    VariableContext numberGroups = new VariableContext();
+    numberGroups.putExtraDataset("groups", Sets.newHashSet("1", "2", "3"));
+
+    VariableContext noGroups = new VariableContext();
+
+    MappingRuleMatcher matchA =
+        MappingRuleMatchers.createUserGroupMatcher("a");
+    MappingRuleMatcher matchB =
+        MappingRuleMatchers.createUserGroupMatcher("b");
+    MappingRuleMatcher matchC =
+        MappingRuleMatchers.createUserGroupMatcher("c");
+    MappingRuleMatcher match1 =
+        MappingRuleMatchers.createUserGroupMatcher("1");
+    MappingRuleMatcher match2 =
+        MappingRuleMatchers.createUserGroupMatcher("2");
+    MappingRuleMatcher match3 =
+        MappingRuleMatchers.createUserGroupMatcher("3");
+    MappingRuleMatcher matchNull =
+        MappingRuleMatchers.createUserGroupMatcher(null);
+
+    //letter groups submission should match only the letters
+    assertTrue(matchA.match(letterGroups));
+    assertTrue(matchB.match(letterGroups));
+    assertTrue(matchC.match(letterGroups));
+    assertFalse(match1.match(letterGroups));
+    assertFalse(match2.match(letterGroups));
+    assertFalse(match3.match(letterGroups));
+    assertFalse(matchNull.match(letterGroups));
+
+    //numeric groups submission should match only the numbers
+    assertFalse(matchA.match(numberGroups));
+    assertFalse(matchB.match(numberGroups));
+    assertFalse(matchC.match(numberGroups));
+    assertTrue(match1.match(numberGroups));
+    assertTrue(match2.match(numberGroups));
+    assertTrue(match3.match(numberGroups));
+    assertFalse(matchNull.match(numberGroups));
+
+    //noGroups submission should not match anything
+    assertFalse(matchA.match(noGroups));
+    assertFalse(matchB.match(noGroups));
+    assertFalse(matchC.match(noGroups));
+    assertFalse(match1.match(noGroups));
+    assertFalse(match2.match(noGroups));
+    assertFalse(match3.match(noGroups));
+    assertFalse(matchNull.match(noGroups));
+  }
 }