diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/nodelabels/CommonNodeLabelsManager.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/nodelabels/CommonNodeLabelsManager.java index f2ff0f6..4190777 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/nodelabels/CommonNodeLabelsManager.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/nodelabels/CommonNodeLabelsManager.java @@ -718,21 +718,19 @@ public void replaceLabelsOnNode(Map> replaceLabelsToNode) public Map> getNodeLabels() { try { readLock.lock(); - Map> nodeToLabels = - new HashMap>(); - for (Entry entry : nodeCollections.entrySet()) { - String hostName = entry.getKey(); - Host host = entry.getValue(); - for (NodeId nodeId : host.nms.keySet()) { - Set nodeLabels = getLabelsByNode(nodeId); - if (nodeLabels == null || nodeLabels.isEmpty()) { - continue; - } - nodeToLabels.put(nodeId, nodeLabels); + Map> nodeToLabels = new HashMap>(); + Set nodeIds = getNodeIds(); + for (NodeId nodeId : nodeIds) { + Set nodeLabels = getLabelsByNode(nodeId); + if (nodeLabels == null || nodeLabels.isEmpty()) { + continue; } - if (!host.labels.isEmpty()) { - nodeToLabels - .put(NodeId.newInstance(hostName, WILDCARD_PORT), host.labels); + nodeToLabels.put(nodeId, nodeLabels); + if (nodeId.getPort() == WILDCARD_PORT) { + Host host = nodeCollections.get(nodeId.getHost()); + if (!host.labels.isEmpty()) { + nodeToLabels.put(nodeId, host.labels); + } } } return Collections.unmodifiableMap(nodeToLabels); @@ -742,45 +740,139 @@ public void replaceLabelsOnNode(Map> replaceLabelsToNode) } /** + * Get list of node ids + * + * @return set of nodeid + */ + public Set getNodeIds() { + Set nodeIds = new HashSet(); + for (Entry entry : nodeCollections.entrySet()) { + nodeIds.addAll(entry.getValue().nms.keySet()); + if (!entry.getValue().labels.isEmpty()) { + nodeIds.add(NodeId.newInstance(entry.getKey(), WILDCARD_PORT)); + } + } + return nodeIds; + } + + public Map> getNodeLabelsInfo() { + try { + readLock.lock(); + Map> nodeToLabelsInfo = + new HashMap>(); + Set nodeIds = getNodeIds(); + for (NodeId nodeId : nodeIds) { + Set nodeLabelNames = getLabelsByNode(nodeId); + if (nodeId.getPort() == WILDCARD_PORT) { + Host host = nodeCollections.get(nodeId.getHost()); + if (!host.labels.isEmpty()) { + nodeLabelNames.addAll(host.labels); + } + } + if (nodeLabelNames == null || nodeLabelNames.isEmpty()) { + continue; + } + + Set nodeLabels = new HashSet(); + for (String label : nodeLabelNames) { + RMNodeLabel rmLabel = labelCollections.get(label); + if (rmLabel == null) { + continue; + } + nodeLabels.add(NodeLabel.newInstance(rmLabel.getLabelName(), + rmLabel.getIsExclusive())); + } + nodeToLabelsInfo.put(nodeId, nodeLabels); + } + return Collections.unmodifiableMap(nodeToLabelsInfo); + } finally { + readLock.unlock(); + } + } + + /** * Get mapping of labels to nodes for all the labels. * * @return labels to nodes map */ - public Map> getLabelsToNodes() { + public Map> getLabelsInfoToNodes() { try { readLock.lock(); - return getLabelsToNodes(labelCollections.keySet()); + return getLabelsInfoToNodes(labelCollections.keySet()); } finally { readLock.unlock(); } } /** - * Get mapping of labels to nodes for specified set of labels. + * Get mapping of labels info to nodes for specified set of labels. * * @param labels set of labels for which labels to nodes mapping will be * returned. * @return labels to nodes map */ - public Map> getLabelsToNodes(Set labels) { + public Map> getLabelsInfoToNodes(Set labels) { try { readLock.lock(); - Map> labelsToNodes = - new HashMap>(); + Map> labelsToNodes = + new HashMap>(); for (String label : labels) { - if(label.equals(NO_LABEL)) { + if (label.equals(NO_LABEL)) { continue; } RMNodeLabel nodeLabelInfo = labelCollections.get(label); - if(nodeLabelInfo != null) { + if (nodeLabelInfo != null) { Set nodeIds = nodeLabelInfo.getAssociatedNodeIds(); if (!nodeIds.isEmpty()) { - labelsToNodes.put(label, nodeIds); + labelsToNodes.put( + NodeLabel.newInstance(label, nodeLabelInfo.getIsExclusive()), + nodeIds); } } else { LOG.warn("getLabelsToNodes : Label [" + label + "] cannot be found"); } - } + } + return Collections.unmodifiableMap(labelsToNodes); + } finally { + readLock.unlock(); + } + } + + /** + * Get mapping of labels to nodes for all the labels. + * + * @return labels to nodes map + */ + @Deprecated + public Map> getLabelsToNodes() { + try { + readLock.lock(); + return getLabelsToNodes(labelCollections.keySet()); + } finally { + readLock.unlock(); + } + } + + /** + * Get mapping of labels to nodes for specified set of labels. + * + * @param labels set of labels for which labels to nodes mapping will be + * returned. + * @return labels to nodes map + */ + @Deprecated + public Map> getLabelsToNodes(Set labels) { + try { + readLock.lock(); + Map> labelsInfoToNodes = + getLabelsInfoToNodes(labels); + Map> labelsToNodes = + new HashMap>(); + for (Entry> entry : labelsInfoToNodes.entrySet()) { + NodeLabel label = entry.getKey(); + Set nodeIds = entry.getValue(); + labelsToNodes.put(label.getName(), nodeIds); + } return Collections.unmodifiableMap(labelsToNodes); } finally { readLock.unlock(); -- 1.9.4.msysgit.1