From 4458b2772f481259453ab5472a476bdebb8c835b Mon Sep 17 00:00:00 2001
From: bibinchundatt <bibinchundatt@apache.org>
Date: Fri, 2 Feb 2018 10:31:00 +0530
Subject: [PATCH] YARN-7840. Update PB for prefix support of node attributes.
 Contributed by Naganarasimha G R.

---
 .../yarn/api/records/NodeAttribute.java       | 22 ++++++++++-
 .../src/main/proto/yarn_protos.proto          |  7 ++--
 .../records/impl/pb/NodeAttributePBImpl.java  | 39 ++++++++++++++++---
 .../hadoop/yarn/api/TestPBImplRecords.java    |  7 ++--
 4 files changed, 61 insertions(+), 14 deletions(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/NodeAttribute.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/NodeAttribute.java
index 13081f3d6df..01c70b296eb 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/NodeAttribute.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/NodeAttribute.java
@@ -37,21 +37,41 @@ import org.apache.hadoop.yarn.util.Records;
  * Its not compulsory for all the attributes to have value, empty string is the
  * default value of the <code>NodeAttributeType.STRING</code>
  * </p>
- *
+ * <p>
+ * Node Attribute Prefix is used as namespace to segregate the attributes.
+ * </p>
  */
 @Public
 @Unstable
 public abstract class NodeAttribute {
 
+  public static final String DEFAULT_PREFIX = "";
+
   public static NodeAttribute newInstance(String attributeName,
       NodeAttributeType attributeType, String attributeValue) {
+    return newInstance(DEFAULT_PREFIX, attributeName, attributeType,
+        attributeValue);
+  }
+
+  public static NodeAttribute newInstance(String attributePrefix,
+      String attributeName, NodeAttributeType attributeType,
+      String attributeValue) {
     NodeAttribute nodeAttribute = Records.newRecord(NodeAttribute.class);
+    nodeAttribute.setAttributePrefix(attributePrefix);
     nodeAttribute.setAttributeName(attributeName);
     nodeAttribute.setAttributeType(attributeType);
     nodeAttribute.setAttributeValue(attributeValue);
     return nodeAttribute;
   }
 
+  @Public
+  @Unstable
+  public abstract String getAttributePrefix();
+
+  @Public
+  @Unstable
+  public abstract void setAttributePrefix(String attributePrefix);
+
   @Public
   @Unstable
   public abstract String getAttributeName();
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/proto/yarn_protos.proto b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/proto/yarn_protos.proto
index 41f5ccb8398..815e989cf81 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/proto/yarn_protos.proto
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/proto/yarn_protos.proto
@@ -377,9 +377,10 @@ enum NodeAttributeTypeProto {
 }
 
 message NodeAttributeProto {
-  optional string attributeName = 1;
-  optional NodeAttributeTypeProto attributeType = 2;
-  optional string attributeValue = 3;
+  optional string attributePrefix = 1;
+  required string attributeName = 2;
+  optional NodeAttributeTypeProto attributeType = 3 [default = STRING];
+  optional string attributeValue = 4 [default=""];
 }
 
 
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/NodeAttributePBImpl.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/NodeAttributePBImpl.java
index 11c9c48347f..7810939eea9 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/NodeAttributePBImpl.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/NodeAttributePBImpl.java
@@ -63,6 +63,10 @@ public class NodeAttributePBImpl extends NodeAttribute {
   @Override
   public void setAttributeName(String attributeName) {
     maybeInitBuilder();
+    if(attributeName == null) {
+      builder.clearAttributeName();
+      return;
+    }
     builder.setAttributeName(attributeName);
   }
 
@@ -78,6 +82,10 @@ public class NodeAttributePBImpl extends NodeAttribute {
   @Override
   public void setAttributeValue(String attributeValue) {
     maybeInitBuilder();
+    if(attributeValue == null) {
+      builder.clearAttributeValue();
+      return;
+    }
     builder.setAttributeValue(attributeValue);
   }
 
@@ -110,12 +118,6 @@ public class NodeAttributePBImpl extends NodeAttribute {
     return NodeAttributeType.valueOf(containerState.name());
   }
 
-  @Override
-  public String toString() {
-    return " name-" + getAttributeName() + ":value-" + getAttributeValue()
-        + ":type-" + getAttributeType();
-  }
-
   @Override
   public int hashCode() {
     return getProto().hashCode();
@@ -152,4 +154,29 @@ public class NodeAttributePBImpl extends NodeAttribute {
       return left.equals(right);
     }
   }
+
+  @Override
+  public String getAttributePrefix() {
+    NodeAttributeProtoOrBuilder p = viaProto ? proto : builder;
+    if (!p.hasAttributePrefix()) {
+      return null;
+    }
+    return p.getAttributePrefix();
+  }
+
+  @Override
+  public void setAttributePrefix(String attributePrefix) {
+    maybeInitBuilder();
+    if(attributePrefix == null) {
+      builder.clearAttributePrefix();
+      return;
+    }
+    builder.setAttributePrefix(attributePrefix);
+  }
+
+  @Override
+  public String toString() {
+    return "Prefix-" + getAttributePrefix() + " :Name-" + getAttributeName()
+        + ":Value-" + getAttributeValue() + ":Type-" + getAttributeType();
+  }
 }
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/api/TestPBImplRecords.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/api/TestPBImplRecords.java
index 80bbac4e1f6..1af4191cc60 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/api/TestPBImplRecords.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/api/TestPBImplRecords.java
@@ -1245,19 +1245,18 @@ public class TestPBImplRecords extends BasePBImplRecordsTest {
 
   @Test
   public void testNodeAttributePBImpl() throws Exception {
-    validatePBImplRecord(NodeAttributePBImpl.class,
-	NodeAttributeProto.class);
+    validatePBImplRecord(NodeAttributePBImpl.class, NodeAttributeProto.class);
   }
 
   @Test
   public void testNodeToAttributesPBImpl() throws Exception {
     validatePBImplRecord(NodeToAttributesPBImpl.class,
-	NodeToAttributesProto.class);
+        NodeToAttributesProto.class);
   }
 
   @Test
   public void testNodesToAttributesMappingRequestPBImpl() throws Exception {
     validatePBImplRecord(NodesToAttributesMappingRequestPBImpl.class,
-	NodesToAttributesMappingRequestProto.class);
+        NodesToAttributesMappingRequestProto.class);
   }
 }