YARN-10371. Create variable context class for CS queue mapping rules. 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/VariableContext.java

@@ -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);
+  }
+
+}

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

@@ -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)));
+  }
+
+}