1 /**
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements.  See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership.  The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License.  You may obtain a copy of the License at
9  *
10  *     http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing, software
13  * distributed under the License is distributed on an "AS IS" BASIS,
14  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15  * See the License for the specific language governing permissions and
16  * limitations under the License.
17  */
18 package org.apache.hadoop.yarn.server.resourcemanager.scheduler.fair;
19 
20 import static org.junit.Assert.*;
21 
22 import java.util.HashMap;
23 import java.util.HashSet;
24 import java.util.Map;
25 import java.util.Set;
26 
27 import javax.xml.parsers.DocumentBuilder;
28 import javax.xml.parsers.DocumentBuilderFactory;
29 
30 import org.apache.commons.io.IOUtils;
31 import org.apache.hadoop.conf.Configuration;
32 import org.apache.hadoop.fs.CommonConfigurationKeys;
33 import org.apache.hadoop.security.GroupMappingServiceProvider;
34 import org.junit.Before;
35 import org.junit.BeforeClass;
36 import org.junit.Test;
37 import org.w3c.dom.Document;
38 import org.w3c.dom.Element;
39 
40 public class TestQueuePlacementPolicy {
41   private final static Configuration conf = new Configuration();
42   private Map<FSQueueType, Set<String>> configuredQueues;
43 
44   @BeforeClass
setup()45   public static void setup() {
46     conf.setClass(CommonConfigurationKeys.HADOOP_SECURITY_GROUP_MAPPING,
47         SimpleGroupsMapping.class, GroupMappingServiceProvider.class);
48   }
49 
50   @Before
initTest()51   public void initTest() {
52     configuredQueues = new HashMap<FSQueueType, Set<String>>();
53     for (FSQueueType type : FSQueueType.values()) {
54       configuredQueues.put(type, new HashSet<String>());
55     }
56   }
57 
58   @Test
testSpecifiedUserPolicy()59   public void testSpecifiedUserPolicy() throws Exception {
60     StringBuffer sb = new StringBuffer();
61     sb.append("<queuePlacementPolicy>");
62     sb.append("  <rule name='specified' />");
63     sb.append("  <rule name='user' />");
64     sb.append("</queuePlacementPolicy>");
65     QueuePlacementPolicy policy = parse(sb.toString());
66     assertEquals("root.specifiedq",
67         policy.assignAppToQueue("specifiedq", "someuser"));
68     assertEquals("root.someuser",
69         policy.assignAppToQueue("default", "someuser"));
70     assertEquals("root.otheruser",
71         policy.assignAppToQueue("default", "otheruser"));
72   }
73 
74   @Test
testNoCreate()75   public void testNoCreate() throws Exception {
76     StringBuffer sb = new StringBuffer();
77     sb.append("<queuePlacementPolicy>");
78     sb.append("  <rule name='specified' />");
79     sb.append("  <rule name='user' create=\"false\" />");
80     sb.append("  <rule name='default' />");
81     sb.append("</queuePlacementPolicy>");
82 
83     configuredQueues.get(FSQueueType.LEAF).add("root.someuser");
84     QueuePlacementPolicy policy = parse(sb.toString());
85     assertEquals("root.specifiedq", policy.assignAppToQueue("specifiedq", "someuser"));
86     assertEquals("root.someuser", policy.assignAppToQueue("default", "someuser"));
87     assertEquals("root.specifiedq", policy.assignAppToQueue("specifiedq", "otheruser"));
88     assertEquals("root.default", policy.assignAppToQueue("default", "otheruser"));
89   }
90 
91   @Test
testSpecifiedThenReject()92   public void testSpecifiedThenReject() throws Exception {
93     StringBuffer sb = new StringBuffer();
94     sb.append("<queuePlacementPolicy>");
95     sb.append("  <rule name='specified' />");
96     sb.append("  <rule name='reject' />");
97     sb.append("</queuePlacementPolicy>");
98     QueuePlacementPolicy policy = parse(sb.toString());
99     assertEquals("root.specifiedq",
100         policy.assignAppToQueue("specifiedq", "someuser"));
101     assertEquals(null, policy.assignAppToQueue("default", "someuser"));
102   }
103 
104   @Test (expected = AllocationConfigurationException.class)
testOmittedTerminalRule()105   public void testOmittedTerminalRule() throws Exception {
106     StringBuffer sb = new StringBuffer();
107     sb.append("<queuePlacementPolicy>");
108     sb.append("  <rule name='specified' />");
109     sb.append("  <rule name='user' create=\"false\" />");
110     sb.append("</queuePlacementPolicy>");
111     parse(sb.toString());
112   }
113 
114   @Test (expected = AllocationConfigurationException.class)
testTerminalRuleInMiddle()115   public void testTerminalRuleInMiddle() throws Exception {
116     StringBuffer sb = new StringBuffer();
117     sb.append("<queuePlacementPolicy>");
118     sb.append("  <rule name='specified' />");
119     sb.append("  <rule name='default' />");
120     sb.append("  <rule name='user' />");
121     sb.append("</queuePlacementPolicy>");
122     parse(sb.toString());
123   }
124 
125   @Test
testTerminals()126   public void testTerminals() throws Exception {
127     // Should make it through without an exception
128     StringBuffer sb = new StringBuffer();
129     sb.append("<queuePlacementPolicy>");
130     sb.append("  <rule name='secondaryGroupExistingQueue' create='true'/>");
131     sb.append("  <rule name='default' create='false'/>");
132     sb.append("</queuePlacementPolicy>");
133     parse(sb.toString());
134   }
135 
136   @Test
testDefaultRuleWithQueueAttribute()137   public void testDefaultRuleWithQueueAttribute() throws Exception {
138     // This test covers the use case where we would like default rule
139     // to point to a different queue by default rather than root.default
140     configuredQueues.get(FSQueueType.LEAF).add("root.someDefaultQueue");
141     StringBuffer sb = new StringBuffer();
142     sb.append("<queuePlacementPolicy>");
143     sb.append("  <rule name='specified' create='false' />");
144     sb.append("  <rule name='default' queue='root.someDefaultQueue'/>");
145     sb.append("</queuePlacementPolicy>");
146 
147     QueuePlacementPolicy policy = parse(sb.toString());
148     assertEquals("root.someDefaultQueue",
149         policy.assignAppToQueue("root.default", "user1"));
150   }
151 
152   @Test
testNestedUserQueueParsingErrors()153   public void testNestedUserQueueParsingErrors() {
154     // No nested rule specified in hierarchical user queue
155     StringBuffer sb = new StringBuffer();
156     sb.append("<queuePlacementPolicy>");
157     sb.append("  <rule name='specified' />");
158     sb.append("  <rule name='nestedUserQueue'/>");
159     sb.append("  <rule name='default' />");
160     sb.append("</queuePlacementPolicy>");
161 
162     assertIfExceptionThrown(sb);
163 
164     // Specified nested rule is not a QueuePlacementRule
165     sb = new StringBuffer();
166     sb.append("<queuePlacementPolicy>");
167     sb.append("  <rule name='specified' />");
168     sb.append("  <rule name='nestedUserQueue'>");
169     sb.append("       <rule name='unknownRule'/>");
170     sb.append("  </rule>");
171     sb.append("  <rule name='default' />");
172     sb.append("</queuePlacementPolicy>");
173 
174     assertIfExceptionThrown(sb);
175   }
176 
assertIfExceptionThrown(StringBuffer sb)177   private void assertIfExceptionThrown(StringBuffer sb) {
178     Throwable th = null;
179     try {
180       parse(sb.toString());
181     } catch (Exception e) {
182       th = e;
183     }
184 
185     assertTrue(th instanceof AllocationConfigurationException);
186   }
187 
188   @Test
testNestedUserQueueParsing()189   public void testNestedUserQueueParsing() throws Exception {
190     StringBuffer sb = new StringBuffer();
191     sb.append("<queuePlacementPolicy>");
192     sb.append("  <rule name='specified' />");
193     sb.append("  <rule name='nestedUserQueue'>");
194     sb.append("       <rule name='primaryGroup'/>");
195     sb.append("  </rule>");
196     sb.append("  <rule name='default' />");
197     sb.append("</queuePlacementPolicy>");
198 
199     Throwable th = null;
200     try {
201       parse(sb.toString());
202     } catch (Exception e) {
203       th = e;
204     }
205 
206     assertNull(th);
207   }
208 
209   @Test
testNestedUserQueuePrimaryGroup()210   public void testNestedUserQueuePrimaryGroup() throws Exception {
211     StringBuffer sb = new StringBuffer();
212     sb.append("<queuePlacementPolicy>");
213     sb.append("  <rule name='specified' create='false' />");
214     sb.append("  <rule name='nestedUserQueue'>");
215     sb.append("       <rule name='primaryGroup'/>");
216     sb.append("  </rule>");
217     sb.append("  <rule name='default' />");
218     sb.append("</queuePlacementPolicy>");
219 
220     // User queue would be created under primary group queue
221     QueuePlacementPolicy policy = parse(sb.toString());
222     assertEquals("root.user1group.user1",
223         policy.assignAppToQueue("root.default", "user1"));
224     // Other rules above and below hierarchical user queue rule should work as
225     // usual
226     configuredQueues.get(FSQueueType.LEAF).add("root.specifiedq");
227     // test if specified rule(above nestedUserQueue rule) works ok
228     assertEquals("root.specifiedq",
229         policy.assignAppToQueue("root.specifiedq", "user2"));
230 
231     // test if default rule(below nestedUserQueue rule) works
232     configuredQueues.get(FSQueueType.LEAF).add("root.user3group");
233     assertEquals("root.default",
234         policy.assignAppToQueue("root.default", "user3"));
235   }
236 
237   @Test
testNestedUserQueuePrimaryGroupNoCreate()238   public void testNestedUserQueuePrimaryGroupNoCreate() throws Exception {
239     // Primary group rule has create='false'
240     StringBuffer sb = new StringBuffer();
241     sb.append("<queuePlacementPolicy>");
242     sb.append("  <rule name='nestedUserQueue'>");
243     sb.append("       <rule name='primaryGroup' create='false'/>");
244     sb.append("  </rule>");
245     sb.append("  <rule name='default' />");
246     sb.append("</queuePlacementPolicy>");
247 
248     QueuePlacementPolicy policy = parse(sb.toString());
249 
250     // Should return root.default since primary group 'root.user1group' is not
251     // configured
252     assertEquals("root.default",
253         policy.assignAppToQueue("root.default", "user1"));
254 
255     // Let's configure primary group and check if user queue is created
256     configuredQueues.get(FSQueueType.PARENT).add("root.user1group");
257     policy = parse(sb.toString());
258     assertEquals("root.user1group.user1",
259         policy.assignAppToQueue("root.default", "user1"));
260 
261     // Both Primary group and nestedUserQueue rule has create='false'
262     sb = new StringBuffer();
263     sb.append("<queuePlacementPolicy>");
264     sb.append("  <rule name='nestedUserQueue' create='false'>");
265     sb.append("       <rule name='primaryGroup' create='false'/>");
266     sb.append("  </rule>");
267     sb.append("  <rule name='default' />");
268     sb.append("</queuePlacementPolicy>");
269 
270     // Should return root.default since primary group and user queue for user 2
271     // are not configured.
272     assertEquals("root.default",
273         policy.assignAppToQueue("root.default", "user2"));
274 
275     // Now configure both primary group and the user queue for user2
276     configuredQueues.get(FSQueueType.PARENT).add("root.user2group");
277     configuredQueues.get(FSQueueType.LEAF).add("root.user2group.user2");
278     policy = parse(sb.toString());
279 
280     assertEquals("root.user2group.user2",
281         policy.assignAppToQueue("root.default", "user2"));
282   }
283 
284   @Test
testNestedUserQueueSecondaryGroup()285   public void testNestedUserQueueSecondaryGroup() throws Exception {
286     StringBuffer sb = new StringBuffer();
287     sb.append("<queuePlacementPolicy>");
288     sb.append("  <rule name='nestedUserQueue'>");
289     sb.append("       <rule name='secondaryGroupExistingQueue'/>");
290     sb.append("  </rule>");
291     sb.append("  <rule name='default' />");
292     sb.append("</queuePlacementPolicy>");
293 
294     QueuePlacementPolicy policy = parse(sb.toString());
295     // Should return root.default since secondary groups are not configured
296     assertEquals("root.default",
297         policy.assignAppToQueue("root.default", "user1"));
298 
299     // configure secondary group for user1
300     configuredQueues.get(FSQueueType.PARENT).add("root.user1subgroup1");
301     policy = parse(sb.toString());
302     // user queue created should be created under secondary group
303     assertEquals("root.user1subgroup1.user1",
304         policy.assignAppToQueue("root.default", "user1"));
305   }
306 
307   @Test
testNestedUserQueueSpecificRule()308   public void testNestedUserQueueSpecificRule() throws Exception {
309     // This test covers the use case where users can specify different parent
310     // queues and want user queues under those.
311     StringBuffer sb = new StringBuffer();
312     sb.append("<queuePlacementPolicy>");
313     sb.append("  <rule name='nestedUserQueue'>");
314     sb.append("       <rule name='specified' create='false'/>");
315     sb.append("  </rule>");
316     sb.append("  <rule name='default' />");
317     sb.append("</queuePlacementPolicy>");
318 
319     // Let's create couple of parent queues
320     configuredQueues.get(FSQueueType.PARENT).add("root.parent1");
321     configuredQueues.get(FSQueueType.PARENT).add("root.parent2");
322 
323     QueuePlacementPolicy policy = parse(sb.toString());
324     assertEquals("root.parent1.user1",
325         policy.assignAppToQueue("root.parent1", "user1"));
326     assertEquals("root.parent2.user2",
327         policy.assignAppToQueue("root.parent2", "user2"));
328   }
329 
330   @Test
testNestedUserQueueDefaultRule()331   public void testNestedUserQueueDefaultRule() throws Exception {
332     // This test covers the use case where we would like user queues to be
333     // created under a default parent queue
334     configuredQueues.get(FSQueueType.PARENT).add("root.parentq");
335     StringBuffer sb = new StringBuffer();
336     sb.append("<queuePlacementPolicy>");
337     sb.append("  <rule name='specified' create='false' />");
338     sb.append("  <rule name='nestedUserQueue'>");
339     sb.append("       <rule name='default' queue='root.parentq'/>");
340     sb.append("  </rule>");
341     sb.append("  <rule name='default' />");
342     sb.append("</queuePlacementPolicy>");
343 
344     QueuePlacementPolicy policy = parse(sb.toString());
345     assertEquals("root.parentq.user1",
346         policy.assignAppToQueue("root.default", "user1"));
347   }
348 
349   @Test
testUserContainsPeriod()350   public void testUserContainsPeriod() throws Exception {
351     // This test covers the user case where the username contains periods.
352     StringBuffer sb = new StringBuffer();
353     sb.append("<queuePlacementPolicy>");
354     sb.append("  <rule name='user' />");
355     sb.append("</queuePlacementPolicy>");
356     QueuePlacementPolicy policy = parse(sb.toString());
357     assertEquals("root.first_dot_last",
358         policy.assignAppToQueue("default", "first.last"));
359 
360     sb = new StringBuffer();
361     sb.append("<queuePlacementPolicy>");
362     sb.append("  <rule name='specified' create='false' />");
363     sb.append("  <rule name='nestedUserQueue'>");
364     sb.append("       <rule name='default'/>");
365     sb.append("  </rule>");
366     sb.append("  <rule name='default' />");
367     sb.append("</queuePlacementPolicy>");
368     policy = parse(sb.toString());
369     assertEquals("root.default.first_dot_last",
370         policy.assignAppToQueue("root.default", "first.last"));
371   }
372 
373   @Test
testGroupContainsPeriod()374   public void testGroupContainsPeriod() throws Exception {
375     StringBuffer sb = new StringBuffer();
376     sb.append("<queuePlacementPolicy>");
377     sb.append("  <rule name='specified' create='false' />");
378     sb.append("  <rule name='nestedUserQueue'>");
379     sb.append("       <rule name='primaryGroup'/>");
380     sb.append("  </rule>");
381     sb.append("  <rule name='default' />");
382     sb.append("</queuePlacementPolicy>");
383 
384     conf.setClass(CommonConfigurationKeys.HADOOP_SECURITY_GROUP_MAPPING,
385         PeriodGroupsMapping.class, GroupMappingServiceProvider.class);
386     // User queue would be created under primary group queue, and the period
387     // in the group name should be converted into _dot_
388     QueuePlacementPolicy policy = parse(sb.toString());
389     assertEquals("root.user1_dot_group.user1",
390         policy.assignAppToQueue("root.default", "user1"));
391 
392     conf.setClass(CommonConfigurationKeys.HADOOP_SECURITY_GROUP_MAPPING,
393         SimpleGroupsMapping.class, GroupMappingServiceProvider.class);
394   }
395 
parse(String str)396   private QueuePlacementPolicy parse(String str) throws Exception {
397     // Read and parse the allocations file.
398     DocumentBuilderFactory docBuilderFactory = DocumentBuilderFactory
399         .newInstance();
400     docBuilderFactory.setIgnoringComments(true);
401     DocumentBuilder builder = docBuilderFactory.newDocumentBuilder();
402     Document doc = builder.parse(IOUtils.toInputStream(str));
403     Element root = doc.getDocumentElement();
404     return QueuePlacementPolicy.fromXml(root, configuredQueues, conf);
405   }
406 }
407