diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/nodelabels/NodeAttributesManager.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/nodelabels/NodeAttributesManager.java index 63f3dcfcc5b..effda9b6703 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/nodelabels/NodeAttributesManager.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/nodelabels/NodeAttributesManager.java @@ -69,10 +69,14 @@ public abstract void removeNodeAttributes( Map> nodeAttributeMapping) throws IOException; /** + * Returns a set of node attributes whose prefix is one of the given + * prefixes; if the prefix set is null or empty, all attributes are returned; + * if prefix set is given but no mapping could be found, an empty set + * is returned. + * * @param prefix set of prefix string's for which the attributes needs to * returned - * @return set of node Attributes objects for the specified set of prefixes, - * else return all + * @return set of node Attributes */ public abstract Set getClusterNodeAttributes( Set prefix); diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/nodelabels/NodeAttributesManagerImpl.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/nodelabels/NodeAttributesManagerImpl.java index 2e63a7ccd4a..a902ac68fdd 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/nodelabels/NodeAttributesManagerImpl.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/nodelabels/NodeAttributesManagerImpl.java @@ -27,6 +27,7 @@ import java.util.Set; import java.util.TreeMap; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentHashMap.KeySetView; import java.util.concurrent.ConcurrentMap; import java.util.concurrent.locks.ReentrantReadWriteLock; import java.util.concurrent.locks.ReentrantReadWriteLock.ReadLock; @@ -141,6 +142,7 @@ private void internalUpdateAttributesOnNodes( Host node = nodeCollections.get(nodeHost); if (node == null) { node = new Host(nodeHost); + nodeCollections.put(nodeHost, node); } switch (op) { case REMOVE: @@ -181,8 +183,16 @@ private void internalUpdateAttributesOnNodes( private void removeNodeFromAttributes(String nodeHost, Set attributeMappings) { - for (NodeAttribute attribute : attributeMappings) { - clusterAttributes.get(attribute).removeNode(nodeHost); + for (NodeAttribute rmAttribute : attributeMappings) { + RMNodeAttribute host = clusterAttributes.get(rmAttribute); + if (host != null) { + host.removeNode(nodeHost); + // If there is no other host has such attribute, + // remove it from the global mapping. + if (host.getAssociatedNodeIds().isEmpty()) { + clusterAttributes.remove(rmAttribute); + } + } } } @@ -305,19 +315,19 @@ protected String normalizeAttributeValue(String value) { @Override public Set getClusterNodeAttributes(Set prefix) { Set attributes = new HashSet<>(); - try { - readLock.lock(); - attributes.addAll(clusterAttributes.keySet()); - } finally { - readLock.unlock(); + KeySetView allAttributes = + clusterAttributes.keySet(); + // Return all if prefix is not given. + if (prefix == null || prefix.isEmpty()) { + attributes.addAll(allAttributes); + return attributes; } - if (prefix != null && prefix.isEmpty()) { - Iterator iterator = attributes.iterator(); - while (iterator.hasNext()) { - NodeAttribute attribute = iterator.next(); - if (!prefix.contains(attribute.getAttributePrefix())) { - iterator.remove(); - } + // Try search attributes by prefix and return valid ones. + Iterator iterator = allAttributes.iterator(); + while (iterator.hasNext()) { + NodeAttribute current = iterator.next(); + if (prefix.contains(current.getAttributePrefix())) { + attributes.add(current); } } return attributes;