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 }