YARN-9866. u:user2:%primary_group is not working as expected. Contributed by Manikandan R

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

@@ -220,9 +220,13 @@ private ApplicationPlacementContext getPlacementForUser(String user)
           }
         }
         if (user.equals(mapping.source)) {
+          if (mapping.queue.equals(PRIMARY_GROUP_MAPPING)) {
+            return getPlacementContext(mapping, groups.getGroups(user).get(0));
+          } else {
             return getPlacementContext(mapping);
           }
         }
+      }
       if (mapping.type == MappingType.GROUP) {
         for (String userGroups : groups.getGroups(user)) {
           if (userGroups.equals(mapping.source)) {

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestCapacitySchedulerQueueMappingFactory.java

@@ -51,8 +51,6 @@ public class TestCapacitySchedulerQueueMappingFactory {
   public static final String USER = "user_";
   public static final String PARENT_QUEUE = "c";
 
-  private MockRM mockRM = null;
-
   public static CapacitySchedulerConfiguration setupQueueMappingsForRules(
       CapacitySchedulerConfiguration conf, String parentQueue,
       boolean overrideWithQueueMappings, int[] sourceIds) {
@@ -114,6 +112,8 @@ public void testUpdatePlacementRulesFactory() throws Exception {
     // init queue mapping for UserGroupMappingRule and AppNameMappingRule
     setupQueueMappingsForRules(conf, PARENT_QUEUE, true, new int[] {1, 2, 3});
 
+    MockRM mockRM = null;
+    try {
       mockRM = new MockRM(conf);
       CapacityScheduler cs = (CapacityScheduler) mockRM.getResourceScheduler();
       cs.updatePlacementRules();
@@ -131,6 +131,11 @@ public void testUpdatePlacementRulesFactory() throws Exception {
       // verify both placement rules were added successfully
       assertThat(placementRuleNames, hasItems(QUEUE_MAPPING_RULE_USER_GROUP));
       assertThat(placementRuleNames, hasItems(QUEUE_MAPPING_RULE_APP_NAME));
+    } finally {
+      if(mockRM != null) {
+        mockRM.close();
+      }
+    }
   }
 
   @Test
@@ -173,6 +178,8 @@ public void testNestedUserQueueWithStaticParentQueue() throws Exception {
     // override with queue mappings
     conf.setOverrideWithQueueMappings(true);
 
+    MockRM mockRM = null;
+    try {
       mockRM = new MockRM(conf);
       CapacityScheduler cs = (CapacityScheduler) mockRM.getResourceScheduler();
       cs.updatePlacementRules();
@@ -195,6 +202,11 @@ public void testNestedUserQueueWithStaticParentQueue() throws Exception {
       ApplicationPlacementContext ctx2 = r.getPlacementForApp(asc, "user2");
       assertEquals("Queue", "user2", ctx2.getQueue());
       assertEquals("Queue", "c", ctx2.getParentQueue());
+    } finally {
+      if(mockRM != null) {
+        mockRM.close();
+      }
+    }
   }
 
   @Test
@@ -303,6 +315,8 @@ private void testNestedUserQueueWithDynamicParentQueue(
     // override with queue mappings
     conf.setOverrideWithQueueMappings(true);
 
+    MockRM mockRM = null;
+    try {
       mockRM = new MockRM(conf);
       CapacityScheduler cs = (CapacityScheduler) mockRM.getResourceScheduler();
       cs.updatePlacementRules();
@@ -324,8 +338,82 @@ private void testNestedUserQueueWithDynamicParentQueue(
       if (primary) {
         assertEquals("Primary Group", user + "group", ctx.getParentQueue());
       } else {
-      assertEquals("Secondary Group", user + "subgroup1", ctx.getParentQueue());
+        assertEquals("Secondary Group", user + "subgroup1",
+            ctx.getParentQueue());
       }
+    } finally {
+      if (mockRM != null) {
         mockRM.close();
       }
     }
+  }
+
+  @Test
+  public void testDynamicPrimaryGroupQueue() throws Exception {
+    CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
+    setupQueueConfiguration(conf);
+    conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class,
+        ResourceScheduler.class);
+    conf.setClass(CommonConfigurationKeys.HADOOP_SECURITY_GROUP_MAPPING,
+        SimpleGroupsMapping.class, GroupMappingServiceProvider.class);
+
+    List<String> queuePlacementRules = new ArrayList<>();
+    queuePlacementRules.add(QUEUE_MAPPING_RULE_USER_GROUP);
+    conf.setQueuePlacementRules(queuePlacementRules);
+
+    List<UserGroupMappingPlacementRule.QueueMapping> existingMappingsForUG =
+        conf.getQueueMappings();
+
+    // set queue mapping
+    List<UserGroupMappingPlacementRule.QueueMapping> queueMappingsForUG =
+        new ArrayList<>();
+
+    // u:user1:b1
+    UserGroupMappingPlacementRule.QueueMapping userQueueMapping1 =
+        new UserGroupMappingPlacementRule.QueueMapping(
+            UserGroupMappingPlacementRule.QueueMapping.MappingType.USER,
+            "user1", "b1");
+
+    // u:user2:%primary_group
+    UserGroupMappingPlacementRule.QueueMapping userQueueMapping2 =
+        new UserGroupMappingPlacementRule.QueueMapping(
+            UserGroupMappingPlacementRule.QueueMapping.MappingType.USER,
+            "user2", "%primary_group");
+
+    queueMappingsForUG.add(userQueueMapping1);
+    queueMappingsForUG.add(userQueueMapping2);
+    existingMappingsForUG.addAll(queueMappingsForUG);
+    conf.setQueueMappings(existingMappingsForUG);
+
+    // override with queue mappings
+    conf.setOverrideWithQueueMappings(true);
+
+    MockRM mockRM = null;
+    try {
+      mockRM = new MockRM(conf);
+      CapacityScheduler cs = (CapacityScheduler) mockRM.getResourceScheduler();
+      cs.updatePlacementRules();
+      mockRM.start();
+      cs.start();
+
+      ApplicationSubmissionContext asc =
+          Records.newRecord(ApplicationSubmissionContext.class);
+      asc.setQueue("default");
+
+      List<PlacementRule> rules =
+          cs.getRMContext().getQueuePlacementManager().getPlacementRules();
+      UserGroupMappingPlacementRule r =
+          (UserGroupMappingPlacementRule) rules.get(0);
+
+      ApplicationPlacementContext ctx = r.getPlacementForApp(asc, "user1");
+      assertEquals("Queue", "b1", ctx.getQueue());
+
+      ApplicationPlacementContext ctx1 = r.getPlacementForApp(asc, "user2");
+      assertEquals("Queue", "user2group", ctx1.getQueue());
+    } finally {
+      if (mockRM != null) {
+        mockRM.close();
+      }
+    }
+  }
+}