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 19 package org.apache.hadoop.yarn.nodelabels; 20 21 import java.util.HashSet; 22 import java.util.Set; 23 import org.apache.commons.lang.StringUtils; 24 import org.apache.hadoop.yarn.api.records.NodeId; 25 import org.apache.hadoop.yarn.api.records.Resource; 26 import org.apache.hadoop.yarn.util.resource.Resources; 27 28 public class NodeLabel implements Comparable<NodeLabel> { 29 private Resource resource; 30 private int numActiveNMs; 31 private String labelName; 32 private Set<NodeId> nodeIds; 33 NodeLabel(String labelName)34 public NodeLabel(String labelName) { 35 this(labelName, Resource.newInstance(0, 0), 0); 36 } 37 NodeLabel(String labelName, Resource res, int activeNMs)38 protected NodeLabel(String labelName, Resource res, int activeNMs) { 39 this.labelName = labelName; 40 this.resource = res; 41 this.numActiveNMs = activeNMs; 42 this.nodeIds = new HashSet<NodeId>(); 43 } 44 addNodeId(NodeId node)45 public void addNodeId(NodeId node) { 46 nodeIds.add(node); 47 } 48 removeNodeId(NodeId node)49 public void removeNodeId(NodeId node) { 50 nodeIds.remove(node); 51 } 52 getAssociatedNodeIds()53 public Set<NodeId> getAssociatedNodeIds() { 54 return new HashSet<NodeId>(nodeIds); 55 } 56 addNode(Resource nodeRes)57 public void addNode(Resource nodeRes) { 58 Resources.addTo(resource, nodeRes); 59 numActiveNMs++; 60 } 61 removeNode(Resource nodeRes)62 public void removeNode(Resource nodeRes) { 63 Resources.subtractFrom(resource, nodeRes); 64 numActiveNMs--; 65 } 66 getResource()67 public Resource getResource() { 68 return this.resource; 69 } 70 getNumActiveNMs()71 public int getNumActiveNMs() { 72 return numActiveNMs; 73 } 74 getLabelName()75 public String getLabelName() { 76 return labelName; 77 } 78 getCopy()79 public NodeLabel getCopy() { 80 return new NodeLabel(labelName, resource, numActiveNMs); 81 } 82 83 @Override compareTo(NodeLabel o)84 public int compareTo(NodeLabel o) { 85 // We should always put empty label entry first after sorting 86 if (labelName.isEmpty() != o.getLabelName().isEmpty()) { 87 if (labelName.isEmpty()) { 88 return -1; 89 } 90 return 1; 91 } 92 93 return labelName.compareTo(o.getLabelName()); 94 } 95 96 @Override equals(Object obj)97 public boolean equals(Object obj) { 98 if (obj instanceof NodeLabel) { 99 NodeLabel other = (NodeLabel) obj; 100 return Resources.equals(resource, other.getResource()) 101 && StringUtils.equals(labelName, other.getLabelName()) 102 && (other.getNumActiveNMs() == numActiveNMs); 103 } 104 return false; 105 } 106 107 @Override hashCode()108 public int hashCode() { 109 final int prime = 502357; 110 return (int) ((((long) labelName.hashCode() << 8) 111 + (resource.hashCode() << 4) + numActiveNMs) % prime); 112 } 113 }