YARN-10371. Create variable context class for CS queue mapping rules. Contributed by Gergely Pollak
This commit is contained in:
Szilard Nemeth
parent
a888d580d8
commit
f4f872b778
@ -0,0 +1,199 @@
|
||||
/**
|
||||
* Licensed to the Apache Software Foundation (ASF) under one
|
||||
* or more contributor license agreements. See the NOTICE file
|
||||
* distributed with this work for additional information
|
||||
* regarding copyright ownership. The ASF licenses this file
|
||||
* to you under the Apache License, Version 2.0 (the
|
||||
* "License"); you may not use this file except in compliance
|
||||
* with the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package org.apache.hadoop.yarn.server.resourcemanager.placement;
|
||||
|
||||
import com.google.common.collect.ImmutableSet;
|
||||
|
||||
import java.util.Arrays;
|
||||
import java.util.HashMap;
|
||||
import java.util.Map;
|
||||
import java.util.Set;
|
||||
|
||||
/**
|
||||
* This class is a key-value store for the variables and their respective values
|
||||
* during an application placement. The class gives support for immutable
|
||||
* variables, which can be set only once, and has helper methods for replacing
|
||||
* the variables with their respective values in provided strings.
|
||||
* We don't extend the map interface, because we don't need all the features
|
||||
* a map provides, this class tries to be as simple as possible.
|
||||
*/
|
||||
public class VariableContext {
|
||||
/**
|
||||
* This is our actual variable store.
|
||||
*/
|
||||
private Map<String, String> variables = new HashMap<>();
|
||||
/**
|
||||
* This set contains the names of the immutable variables if null it is
|
||||
* ignored.
|
||||
*/
|
||||
private Set<String> immutableNames;
|
||||
|
||||
/**
|
||||
* Checks if the provided variable is immutable.
|
||||
* @param name Name of the variable to check
|
||||
* @return true if the variable is immutable
|
||||
*/
|
||||
boolean isImmutable(String name) {
|
||||
return (immutableNames != null && immutableNames.contains(name));
|
||||
}
|
||||
|
||||
/**
|
||||
* Can be used to provide a set which contains the name of the variables which
|
||||
* should be immutable.
|
||||
* @param variableNames Set containing the names of the immutable variables
|
||||
* @throws IllegalStateException if the immutable set is already provided.
|
||||
* @return same instance of VariableContext for daisy chaining.
|
||||
*/
|
||||
public VariableContext setImmutables(Set<String> variableNames) {
|
||||
if (this.immutableNames != null) {
|
||||
throw new IllegalStateException("Immutable variables are already defined,"
|
||||
+ " variable immutability cannot be changed once set!");
|
||||
}
|
||||
this.immutableNames = ImmutableSet.copyOf(variableNames);
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Can be used to provide an array of strings which contains the names of the
|
||||
* variables which should be immutable. An immutable set will be created
|
||||
* from the array.
|
||||
* @param variableNames Set containing the names of the immutable variables
|
||||
* @throws IllegalStateException if the immutable set is already provided.
|
||||
* @return same instance of VariableContext for daisy chaining.
|
||||
*/
|
||||
public VariableContext setImmutables(String... variableNames) {
|
||||
if (this.immutableNames != null) {
|
||||
throw new IllegalStateException("Immutable variables are already defined,"
|
||||
+ " variable immutability cannot be changed once set!");
|
||||
}
|
||||
this.immutableNames = ImmutableSet.copyOf(variableNames);
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Adds a variable with value to the context or overrides an already existing
|
||||
* one. If the variable is already set and immutable an IllegalStateException
|
||||
* is thrown.
|
||||
* @param name Name of the variable to be added to the context
|
||||
* @param value Value of the variable
|
||||
* @throws IllegalStateException if the variable is immutable and already set
|
||||
* @return same instance of VariableContext for daisy chaining.
|
||||
*/
|
||||
public VariableContext put(String name, String value) {
|
||||
if (variables.containsKey(name) && isImmutable(name)) {
|
||||
throw new IllegalStateException(
|
||||
"Variable '" + name + "' is immutable, cannot update it's value!");
|
||||
}
|
||||
variables.put(name, value);
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns the value of a variable, null values are replaced with "".
|
||||
* @param name Name of the variable
|
||||
* @return The value of the variable
|
||||
*/
|
||||
public String get(String name) {
|
||||
String ret = variables.get(name);
|
||||
return ret == null ? "" : ret;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a variable is part of the context.
|
||||
* @param name Name of the variable to be checked
|
||||
* @return True if the variable is added to the context, false otherwise
|
||||
*/
|
||||
public boolean containsKey(String name) {
|
||||
return variables.containsKey(name);
|
||||
}
|
||||
|
||||
/**
|
||||
* This method replaces all variables in the provided string. The variables
|
||||
* are reverse ordered by the length of their names in order to avoid partial
|
||||
* replaces when a shorter named variable is a substring of a longer named
|
||||
* variable.
|
||||
* All variables will be replaced in the string.
|
||||
* Null values will be considered as empty strings during the replace.
|
||||
* If the input is null, null will be returned.
|
||||
* @param input The string with variables
|
||||
* @return A string with all the variables substituted with their respective
|
||||
* values.
|
||||
*/
|
||||
public String replaceVariables(String input) {
|
||||
if (input == null) {
|
||||
return null;
|
||||
}
|
||||
|
||||
String[] keys = variables.keySet().toArray(new String[]{});
|
||||
//Replacing variables starting longest first, to avoid collision when a
|
||||
//shorter variable name matches the beginning of a longer one.
|
||||
//e.g. %user_something, if %user is defined it may replace the %user before
|
||||
//we would reach the %user_something variable, so we start with the longer
|
||||
//names first
|
||||
Arrays.sort(keys, (a, b) -> b.length() - a.length());
|
||||
|
||||
String ret = input;
|
||||
for (String key : keys) {
|
||||
//we cannot match for null, so we just skip if we have a variable "name"
|
||||
//with null
|
||||
if (key == null) {
|
||||
continue;
|
||||
}
|
||||
ret = ret.replace(key, get(key));
|
||||
}
|
||||
|
||||
return ret;
|
||||
}
|
||||
|
||||
/**
|
||||
* This method will consider the input as a queue path, which is a String
|
||||
* separated by dot ('.') characters. The input will be split along the dots
|
||||
* and all parts will be replaced individually. Replace only occur if a part
|
||||
* exactly matches a variable name, no composite names or additional
|
||||
* characters are supported.
|
||||
* e.g. With variables %user and %default "%user.%default" will be substituted
|
||||
* while "%user%default.something" won't.
|
||||
* Null values will be considered as empty strings during the replace.
|
||||
* If the input is null, null will be returned.
|
||||
* @param input The string with variables
|
||||
* @return A string with all the variable only path parts substituted with
|
||||
* their respective values.
|
||||
*/
|
||||
public String replacePathVariables(String input) {
|
||||
if (input == null) {
|
||||
return null;
|
||||
}
|
||||
|
||||
String[] parts = input.split("\\.");
|
||||
for (int i = 0; i < parts.length; i++) {
|
||||
//if the part is a variable it should be in the map, otherwise we keep
|
||||
//it's original value. This means undefined variables will return the
|
||||
//name of the variable, but this is working as intended.
|
||||
String newVal = variables.getOrDefault(parts[i], parts[i]);
|
||||
//if a variable's value is null, we use empty string instead
|
||||
if (newVal == null) {
|
||||
newVal = "";
|
||||
}
|
||||
parts[i] = newVal;
|
||||
}
|
||||
|
||||
return String.join(".", parts);
|
||||
}
|
||||
|
||||
}
|
||||
@ -0,0 +1,202 @@
|
||||
/**
|
||||
* Licensed to the Apache Software Foundation (ASF) under one
|
||||
* or more contributor license agreements. See the NOTICE file
|
||||
* distributed with this work for additional information
|
||||
* regarding copyright ownership. The ASF licenses this file
|
||||
* to you under the Apache License, Version 2.0 (the
|
||||
* "License"); you may not use this file except in compliance
|
||||
* with the License. You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package org.apache.hadoop.yarn.server.resourcemanager.placement;
|
||||
|
||||
import com.google.common.collect.ImmutableSet;
|
||||
import org.junit.Test;
|
||||
|
||||
import java.util.HashMap;
|
||||
|
||||
import static org.junit.Assert.*;
|
||||
|
||||
public class TestVariableContext {
|
||||
|
||||
@Test
|
||||
public void testAddAndGet() {
|
||||
VariableContext variables = new VariableContext();
|
||||
|
||||
assertEquals("", variables.get("%user"));
|
||||
assertFalse(variables.containsKey("%user"));
|
||||
|
||||
variables.put("%user", "john");
|
||||
variables.put("%group", "primary");
|
||||
variables.put("%group", "secondary");
|
||||
variables.put("%empty", null);
|
||||
assertTrue(variables.containsKey("%user"));
|
||||
assertTrue(variables.containsKey("%empty"));
|
||||
|
||||
assertEquals("john", variables.get("%user"));
|
||||
assertEquals("secondary", variables.get("%group"));
|
||||
assertEquals("", variables.get("%empty"));
|
||||
}
|
||||
|
||||
@Test(expected = IllegalStateException.class)
|
||||
public void testImmutablesCanOnlySetOnceFromSet() {
|
||||
VariableContext variables = new VariableContext();
|
||||
ImmutableSet<String> immutables =
|
||||
ImmutableSet.of("%user", "%primary_group", "%secondary_group");
|
||||
|
||||
variables.setImmutables(immutables);
|
||||
variables.setImmutables(immutables);
|
||||
}
|
||||
|
||||
@Test(expected = IllegalStateException.class)
|
||||
public void testImmutablesCanOnlySetOnceFromArray() {
|
||||
VariableContext variables = new VariableContext();
|
||||
|
||||
variables.setImmutables("%user", "%primary_group", "%secondary_group");
|
||||
variables.setImmutables("%user", "%primary_group", "%secondary_group");
|
||||
}
|
||||
|
||||
@Test(expected = IllegalStateException.class)
|
||||
public void testImmutablesCanOnlySetOnceFromSetAndArray() {
|
||||
VariableContext variables = new VariableContext();
|
||||
ImmutableSet<String> immutables =
|
||||
ImmutableSet.of("%user", "%primary_group", "%secondary_group");
|
||||
|
||||
variables.setImmutables(immutables);
|
||||
variables.setImmutables("%user", "%primary_group", "%secondary_group");
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testImmutableVariableCanBeSetOnce() {
|
||||
VariableContext variables = new VariableContext();
|
||||
ImmutableSet<String> immutables =
|
||||
ImmutableSet.of("%user", "%primary_group", "%secondary_group");
|
||||
|
||||
variables.setImmutables(immutables);
|
||||
variables.put("%user", "bob");
|
||||
}
|
||||
|
||||
@Test(expected = IllegalStateException.class)
|
||||
public void testImmutableVariableProtection() {
|
||||
VariableContext variables = new VariableContext();
|
||||
ImmutableSet<String> immutables =
|
||||
ImmutableSet.of("%user", "%primary_group", "%secondary_group");
|
||||
|
||||
variables.setImmutables(immutables);
|
||||
variables.put("%user", "bob");
|
||||
variables.put("%user", "bob");
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testAddAndGetWithImmutables() {
|
||||
VariableContext variables = new VariableContext();
|
||||
ImmutableSet<String> immutables =
|
||||
ImmutableSet.of("%user", "%primary_group", "%secondary_group");
|
||||
|
||||
assertFalse(variables.isImmutable("%user"));
|
||||
assertFalse(variables.isImmutable("%primary_group"));
|
||||
assertFalse(variables.isImmutable("%secondary_group"));
|
||||
assertFalse(variables.isImmutable("%default"));
|
||||
|
||||
variables.setImmutables(immutables);
|
||||
assertTrue(variables.isImmutable("%user"));
|
||||
assertTrue(variables.isImmutable("%primary_group"));
|
||||
assertTrue(variables.isImmutable("%secondary_group"));
|
||||
assertFalse(variables.isImmutable("%default"));
|
||||
variables.put("%user", "bob");
|
||||
variables.put("%primary_group", "primary");
|
||||
variables.put("%default", "root.default");
|
||||
|
||||
assertEquals("bob", variables.get("%user"));
|
||||
assertEquals("primary", variables.get("%primary_group"));
|
||||
assertEquals("root.default", variables.get("%default"));
|
||||
|
||||
variables.put("%default", "root.new.default");
|
||||
assertEquals("root.new.default", variables.get("%default"));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testPathPartReplace() {
|
||||
VariableContext variables = new VariableContext();
|
||||
ImmutableSet<String> immutables =
|
||||
ImmutableSet.of("%user", "%primary_group", "%secondary_group");
|
||||
|
||||
variables
|
||||
.setImmutables(immutables)
|
||||
.put("%user", "bob")
|
||||
.put("%primary_group", "developers")
|
||||
.put("%secondary_group", "yarn-dev")
|
||||
.put("%default", "default.path")
|
||||
.put("%null", null)
|
||||
.put("%empty", "");
|
||||
|
||||
HashMap<String, String> testCases = new HashMap<>();
|
||||
testCases.put("nothing_to_replace", "nothing_to_replace");
|
||||
testCases.put(null, null);
|
||||
testCases.put("", "");
|
||||
testCases.put("%empty", "");
|
||||
testCases.put("%null", "");
|
||||
testCases.put("%user", "bob");
|
||||
testCases.put("root.regular.path", "root.regular.path");
|
||||
testCases.put("root.%empty.path", "root..path");
|
||||
testCases.put("root.%empty%empty.path", "root.%empty%empty.path");
|
||||
testCases.put("root.%null.path", "root..path");
|
||||
testCases.put(
|
||||
"root.%user.%primary_group.%secondary_group.%default.%null.%empty.end",
|
||||
"root.bob.developers.yarn-dev.default.path...end");
|
||||
testCases.put(
|
||||
"%user%default.%user.%default", "%user%default.bob.default.path");
|
||||
|
||||
testCases.forEach(
|
||||
(k, v) -> assertEquals(v, variables.replacePathVariables(k)));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testVariableReplace() {
|
||||
VariableContext variables = new VariableContext();
|
||||
ImmutableSet<String> immutables =
|
||||
ImmutableSet.of("%user", "%primary_group", "%secondary_group");
|
||||
|
||||
variables
|
||||
.setImmutables(immutables)
|
||||
.put("%user", "bob")
|
||||
.put("%userPhone", "555-3221")
|
||||
.put("%primary_group", "developers")
|
||||
.put("%secondary_group", "yarn-dev")
|
||||
.put("%default", "default.path")
|
||||
.put("%null", null)
|
||||
.put("%empty", "");
|
||||
|
||||
HashMap<String, String> testCases = new HashMap<>();
|
||||
testCases.put("nothing_to_replace", "nothing_to_replace");
|
||||
testCases.put(null, null);
|
||||
testCases.put("", "");
|
||||
testCases.put("%empty", "");
|
||||
testCases.put("%null", "");
|
||||
testCases.put("%user", "bob");
|
||||
testCases.put("%userPhone", "555-3221");
|
||||
testCases.put("root.regular.path", "root.regular.path");
|
||||
testCases.put("root.%empty.path", "root..path");
|
||||
testCases.put("root.%empty%empty.path", "root..path");
|
||||
testCases.put("root.%null.path", "root..path");
|
||||
testCases.put(
|
||||
"root.%user.%primary_group.%secondary_group.%default.%null.%empty.end",
|
||||
"root.bob.developers.yarn-dev.default.path...end");
|
||||
testCases.put(
|
||||
"%user%default.%user.%default", "bobdefault.path.bob.default.path");
|
||||
testCases.put(
|
||||
"userPhoneof%useris%userPhone", "userPhoneofbobis555-3221");
|
||||
|
||||
testCases.forEach((pattern, expected) ->
|
||||
assertEquals(expected, variables.replaceVariables(pattern)));
|
||||
}
|
||||
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user