YARN-3431. Sub resources of timeline entity needs to be passed to a separate endpoint. Contributed By Zhijie Shen.

(cherry picked from commit fa5cc75245a6dba549620a8b26c7b4a8aed9838e)

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/timelineservice/ApplicationAttemptEntity.java

@@ -20,16 +20,17 @@ package org.apache.hadoop.yarn.api.records.timelineservice;
 import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.classification.InterfaceStability;
 
-import javax.xml.bind.annotation.XmlAccessType;
-import javax.xml.bind.annotation.XmlAccessorType;
-import javax.xml.bind.annotation.XmlRootElement;
-
-@XmlRootElement(name = "appattempt")
-@XmlAccessorType(XmlAccessType.NONE)
 @InterfaceAudience.Public
 @InterfaceStability.Unstable
 public class ApplicationAttemptEntity extends HierarchicalTimelineEntity {
   public ApplicationAttemptEntity() {
     super(TimelineEntityType.YARN_APPLICATION_ATTEMPT.toString());
   }
+
+  public ApplicationAttemptEntity(TimelineEntity entity) {
+    super(entity);
+    if (!entity.getType().equals(TimelineEntityType.YARN_APPLICATION_ATTEMPT.toString())) {
+      throw new IllegalArgumentException("Incompatible entity type: " + getId());
+    }
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/timelineservice/ApplicationEntity.java

@@ -20,28 +20,28 @@ package org.apache.hadoop.yarn.api.records.timelineservice;
 import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.classification.InterfaceStability;
 
-import javax.xml.bind.annotation.XmlAccessType;
-import javax.xml.bind.annotation.XmlAccessorType;
-import javax.xml.bind.annotation.XmlElement;
-import javax.xml.bind.annotation.XmlRootElement;
-
-@XmlRootElement(name = "application")
-@XmlAccessorType(XmlAccessType.NONE)
 @InterfaceAudience.Public
 @InterfaceStability.Unstable
 public class ApplicationEntity extends HierarchicalTimelineEntity {
-  private String queue;
+  public static final String QUEUE_INFO_KEY =
+      TimelineEntity.SYSTEM_INFO_KEY_PREFIX + "QUEUE";
 
   public ApplicationEntity() {
     super(TimelineEntityType.YARN_APPLICATION.toString());
   }
 
-  @XmlElement(name = "queue")
+  public ApplicationEntity(TimelineEntity entity) {
+    super(entity);
+    if (!entity.getType().equals(TimelineEntityType.YARN_APPLICATION.toString())) {
+      throw new IllegalArgumentException("Incompatible entity type: " + getId());
+    }
+  }
+
   public String getQueue() {
-    return queue;
+    return getInfo().get(QUEUE_INFO_KEY).toString();
   }
 
   public void setQueue(String queue) {
-    this.queue = queue;
+    addInfo(QUEUE_INFO_KEY, queue);
   }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/timelineservice/ClusterEntity.java

@@ -20,12 +20,6 @@ package org.apache.hadoop.yarn.api.records.timelineservice;
 import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.classification.InterfaceStability;
 
-import javax.xml.bind.annotation.XmlAccessType;
-import javax.xml.bind.annotation.XmlAccessorType;
-import javax.xml.bind.annotation.XmlRootElement;
-
-@XmlRootElement(name = "cluster")
-@XmlAccessorType(XmlAccessType.NONE)
 @InterfaceAudience.Public
 @InterfaceStability.Unstable
 public class ClusterEntity extends HierarchicalTimelineEntity {
@@ -33,4 +27,10 @@ public class ClusterEntity extends HierarchicalTimelineEntity {
     super(TimelineEntityType.YARN_CLUSTER.toString());
   }
 
+  public ClusterEntity(TimelineEntity entity) {
+    super(entity);
+    if (!entity.getType().equals(TimelineEntityType.YARN_CLUSTER.toString())) {
+      throw new IllegalArgumentException("Incompatible entity type: " + getId());
+    }
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/timelineservice/ContainerEntity.java

@@ -20,16 +20,17 @@ package org.apache.hadoop.yarn.api.records.timelineservice;
 import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.classification.InterfaceStability;
 
-import javax.xml.bind.annotation.XmlAccessType;
-import javax.xml.bind.annotation.XmlAccessorType;
-import javax.xml.bind.annotation.XmlRootElement;
-
-@XmlRootElement(name = "container")
-@XmlAccessorType(XmlAccessType.NONE)
 @InterfaceAudience.Public
 @InterfaceStability.Unstable
 public class ContainerEntity extends HierarchicalTimelineEntity {
   public ContainerEntity() {
     super(TimelineEntityType.YARN_CONTAINER.toString());
   }
+
+  public ContainerEntity(TimelineEntity entity) {
+    super(entity);
+    if (!entity.getType().equals(TimelineEntityType.YARN_CONTAINER.toString())) {
+      throw new IllegalArgumentException("Incompatible entity type: " + getId());
+    }
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/timelineservice/FlowEntity.java

@@ -20,62 +20,84 @@ package org.apache.hadoop.yarn.api.records.timelineservice;
 import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.classification.InterfaceStability;
 
-import javax.xml.bind.annotation.XmlAccessType;
-import javax.xml.bind.annotation.XmlAccessorType;
 import javax.xml.bind.annotation.XmlElement;
-import javax.xml.bind.annotation.XmlRootElement;
 
-@XmlRootElement(name = "flow")
-@XmlAccessorType(XmlAccessType.NONE)
 @InterfaceAudience.Public
 @InterfaceStability.Unstable
 public class FlowEntity extends HierarchicalTimelineEntity {
-  private String user;
-  private String version;
-  private String run;
+  public static final String USER_INFO_KEY =
+      TimelineEntity.SYSTEM_INFO_KEY_PREFIX + "USER";
+  public static final String FLOW_NAME_INFO_KEY =
+      TimelineEntity.SYSTEM_INFO_KEY_PREFIX + "FLOW_NAME";
+  public static final String FLOW_VERSION_INFO_KEY =
+      TimelineEntity.SYSTEM_INFO_KEY_PREFIX +  "FLOW_VERSION";
+  public static final String FLOW_RUN_ID_INFO_KEY =
+      TimelineEntity.SYSTEM_INFO_KEY_PREFIX +  "FLOW_RUN_ID";
 
   public FlowEntity() {
     super(TimelineEntityType.YARN_FLOW.toString());
   }
 
-  @Override
-  public String getId() {
-    //Flow id schema: user@flow_name(or id)/version/run
-    StringBuilder sb = new StringBuilder();
-    sb.append(user);
-    sb.append('@');
-    sb.append(super.getId());
-    sb.append('/');
-    sb.append(version);
-    sb.append('/');
-    sb.append(run);
-    return sb.toString();
+  public FlowEntity(TimelineEntity entity) {
+    super(entity);
+    if (!entity.getType().equals(TimelineEntityType.YARN_FLOW.toString())) {
+      throw new IllegalArgumentException("Incompatible entity type: " + getId());
+    }
+  }
+
+  @XmlElement(name = "id")
+  @Override
+  public String getId() {
+    //Flow id schema: user@flow_name(or id)/version/run_id
+    String id = super.getId();
+    if (id == null) {
+      StringBuilder sb = new StringBuilder();
+      sb.append(getInfo().get(USER_INFO_KEY).toString());
+      sb.append('@');
+      sb.append(getInfo().get(FLOW_NAME_INFO_KEY).toString());
+      sb.append('/');
+      sb.append(getInfo().get(FLOW_VERSION_INFO_KEY).toString());
+      sb.append('/');
+      sb.append(getInfo().get(FLOW_RUN_ID_INFO_KEY).toString());
+      id = sb.toString();
+      setId(id);
+    }
+    return id;
   }
 
-  @XmlElement(name = "user")
   public String getUser() {
-    return user;
+    Object user = getInfo().get(USER_INFO_KEY);
+    return user == null ? null : user.toString();
   }
 
   public void setUser(String user) {
-    this.user = user;
+    addInfo(USER_INFO_KEY, user);
+  }
+
+  public String getName() {
+    Object name = getInfo().get(FLOW_NAME_INFO_KEY);
+    return name == null ? null : name.toString();
+  }
+
+  public void setName(String name) {
+    addInfo(FLOW_NAME_INFO_KEY, name);
   }
 
-  @XmlElement(name = "version")
   public String getVersion() {
-    return version;
+    Object version = getInfo().get(FLOW_VERSION_INFO_KEY);
+    return version == null ? null : version.toString();
   }
 
   public void setVersion(String version) {
-    this.version = version;
+    addInfo(FLOW_VERSION_INFO_KEY, version);
   }
 
-  @XmlElement(name = "run")
-  public String getRun() {
-    return run;
+  public long getRunId() {
+    Object runId = getInfo().get(FLOW_RUN_ID_INFO_KEY);
+    return runId == null ? 0L : (Long) runId;
   }
 
-  public void setRun(String run) {
-    this.run = run;
+  public void setRunId(long runId) {
+    addInfo(FLOW_RUN_ID_INFO_KEY, runId);
   }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/timelineservice/HierarchicalTimelineEntity.java

@@ -17,93 +17,98 @@
  */
 package org.apache.hadoop.yarn.api.records.timelineservice;
 
+import com.google.common.base.Joiner;
 import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.classification.InterfaceStability;
+import org.apache.hadoop.yarn.exceptions.YarnRuntimeException;
 
-import javax.xml.bind.annotation.XmlAccessType;
-import javax.xml.bind.annotation.XmlAccessorType;
-import javax.xml.bind.annotation.XmlElement;
-import java.util.HashMap;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashSet;
-import java.util.Map;
+import java.util.List;
 import java.util.Set;
 
-@XmlAccessorType(XmlAccessType.NONE)
 @InterfaceAudience.Public
 @InterfaceStability.Unstable
 public abstract class HierarchicalTimelineEntity extends TimelineEntity {
-  private Identifier parent;
-  private HashMap<String, Set<String>> children = new HashMap<>();
+  public static final String PARENT_INFO_KEY =
+      TimelineEntity.SYSTEM_INFO_KEY_PREFIX + "PARENT_ENTITY";
+  public static final String CHILDREN_INFO_KEY =
+      TimelineEntity.SYSTEM_INFO_KEY_PREFIX + "CHILDREN_ENTITY";
+
+  HierarchicalTimelineEntity(TimelineEntity entity) {
+    super(entity);
+  }
 
   HierarchicalTimelineEntity(String type) {
     super(type);
   }
 
-  @XmlElement(name = "parent")
   public Identifier getParent() {
-    return parent;
+    Object obj = getInfo().get(PARENT_INFO_KEY);
+    if (obj != null) {
+      if (obj instanceof Identifier) {
+        return (Identifier) obj;
+      } else {
+        throw new YarnRuntimeException(
+            "Parent info is invalid identifier object");
+      }
+    }
+    return null;
   }
 
   public void setParent(Identifier parent) {
     validateParent(parent.getType());
-    this.parent = parent;
+    addInfo(PARENT_INFO_KEY, parent);
   }
 
   public void setParent(String type, String id) {
-    validateParent(type);
-    parent = new Identifier();
-    parent.setType(type);
-    parent.setId(id);
+    setParent(new Identifier(type, id));
   }
 
-  // required by JAXB
-  @InterfaceAudience.Private
-  // comment out XmlElement here because it cause UnrecognizedPropertyException
-  // TODO we need a better fix
-  //@XmlElement(name = "children")
-  public HashMap<String, Set<String>> getChildrenJAXB() {
-    return children;
+  public Set<Identifier> getChildren() {
+    Object identifiers = getInfo().get(CHILDREN_INFO_KEY);
+    if (identifiers == null) {
+      return new HashSet<>();
     }
-
-  public Map<String, Set<String>> getChildren() {
-    return children;
-  }
-
-  public void setChildren(Map<String, Set<String>> children) {
-    validateChildren(children);
-    if (children != null && !(children instanceof HashMap)) {
-      this.children = new HashMap<String, Set<String>>(children);
+    TimelineEntityType thisType = TimelineEntityType.valueOf(getType());
+    if (identifiers instanceof Set<?>) {
+      for (Object identifier : (Set<?>) identifiers) {
+        if (!(identifier instanceof Identifier)) {
+          throw new YarnRuntimeException(
+              "Children info contains invalid identifier object");
         } else {
-      this.children = (HashMap) children;
+          validateChild((Identifier) identifier, thisType);
         }
       }
+    } else {
+      throw new YarnRuntimeException(
+          "Children info is invalid identifier set");
+    }
+    Set<Identifier> children = (Set<Identifier>) identifiers;
+    return children;
+  }
 
-  public void addChildren(Map<String, Set<String>> children) {
-    validateChildren(children);
-    for (Map.Entry<String, Set<String>> entry : children.entrySet()) {
-      Set<String> ids = this.children.get(entry.getKey());
-      if (ids == null) {
-        ids = new HashSet<>();
-        this.children.put(entry.getKey(), ids);
+  public void setChildren(Set<Identifier> children) {
+    addInfo(CHILDREN_INFO_KEY, children);
   }
-      ids.addAll(entry.getValue());
+
+  public void addChildren(Set<Identifier> children) {
+    TimelineEntityType thisType = TimelineEntityType.valueOf(getType());
+    for (Identifier child : children) {
+      validateChild(child, thisType);
     }
+    Set<Identifier> existingChildren = getChildren();
+    existingChildren.addAll(children);
+    setChildren(existingChildren);
+  }
+
+  public void addChild(Identifier child) {
+    addChildren(Collections.singleton(child));
   }
 
   public void addChild(String type, String id) {
-    TimelineEntityType thisType = TimelineEntityType.valueOf(getType());
-    TimelineEntityType childType = TimelineEntityType.valueOf(type);
-    if (thisType.isChild(childType)) {
-      Set<String> ids = children.get(type);
-      if (ids == null) {
-        ids = new HashSet<>();
-        children.put(type, ids);
-      }
-      ids.add(id);
-    } else {
-      throw new IllegalArgumentException(
-          type + " is not the acceptable child of " + this.getType());
-    }
+    addChild(new Identifier(type, id));
   }
 
   private void validateParent(String type) {
@@ -115,15 +120,12 @@ public abstract class HierarchicalTimelineEntity extends TimelineEntity {
     }
   }
 
-  private void validateChildren(Map<String, Set<String>> children) {
-    TimelineEntityType thisType = TimelineEntityType.valueOf(getType());
-    for (Map.Entry<String, Set<String>> entry : children.entrySet()) {
-      TimelineEntityType childType = TimelineEntityType.valueOf(entry.getKey());
+  private void validateChild(Identifier child, TimelineEntityType thisType) {
+    TimelineEntityType childType = TimelineEntityType.valueOf(child.getType());
     if (!thisType.isChild(childType)) {
       throw new IllegalArgumentException(
-            entry.getKey() + " is not the acceptable child of " +
+          child.getType() + " is not the acceptable child of " +
               this.getType());
     }
   }
 }
-}

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/timelineservice/QueueEntity.java

@@ -20,16 +20,17 @@ package org.apache.hadoop.yarn.api.records.timelineservice;
 import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.classification.InterfaceStability;
 
-import javax.xml.bind.annotation.XmlAccessType;
-import javax.xml.bind.annotation.XmlAccessorType;
-import javax.xml.bind.annotation.XmlRootElement;
-
-@XmlRootElement(name = "queue")
-@XmlAccessorType(XmlAccessType.NONE)
 @InterfaceAudience.Public
 @InterfaceStability.Unstable
-public class TimelineQueue extends HierarchicalTimelineEntity {
-  public TimelineQueue() {
+public class QueueEntity extends HierarchicalTimelineEntity {
+  public QueueEntity() {
     super(TimelineEntityType.YARN_QUEUE.toString());
   }
+
+  public QueueEntity(TimelineEntity entity) {
+    super(entity);
+    if (!entity.getType().equals(TimelineEntityType.YARN_QUEUE.toString())) {
+      throw new IllegalArgumentException("Incompatible entity type: " + getId());
+    }
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/timelineservice/TimelineEntity.java

@@ -34,6 +34,7 @@ import java.util.Set;
 @InterfaceAudience.Public
 @InterfaceStability.Unstable
 public class TimelineEntity {
+  protected final static String SYSTEM_INFO_KEY_PREFIX = "SYSTEM_INFO_";
 
   @XmlRootElement(name = "identifier")
   @XmlAccessorType(XmlAccessType.NONE)
@@ -41,6 +42,11 @@ public class TimelineEntity {
     private String type;
     private String id;
 
+    public Identifier(String type, String id) {
+      this.type = type;
+      this.id = id;
+    }
+
     public Identifier() {
 
     }
@@ -62,8 +68,16 @@ public class TimelineEntity {
     public void setId(String id) {
       this.id = id;
     }
+
+    @Override
+    public String toString() {
+      return "TimelineEntity[" +
+          "type='" + type + '\'' +
+          ", id='" + id + '\'' + "]";
+    }
   }
 
+  private TimelineEntity real;
   private Identifier identifier;
   private HashMap<String, Object> info = new HashMap<>();
   private HashMap<String, Object> configs = new HashMap<>();
@@ -78,6 +92,22 @@ public class TimelineEntity {
     identifier = new Identifier();
   }
 
+  /**
+   * <p>
+   * The constuctor is used to construct a proxy {@link TimelineEntity} or its
+   * subclass object from the real entity object that carries information.
+   * </p>
+   *
+   * <p>
+   * It is usually used in the case where we want to recover class polymorphism
+   * after deserializing the entity from its JSON form.
+   * </p>
+   * @param entity the real entity that carries information
+   */
+  public TimelineEntity(TimelineEntity entity) {
+    real = entity.getReal();
+  }
+
   protected TimelineEntity(String type) {
     this();
     identifier.type = type;
@@ -85,140 +115,252 @@ public class TimelineEntity {
 
   @XmlElement(name = "type")
   public String getType() {
+    if (real == null) {
       return identifier.type;
+    } else {
+      return real.getType();
+    }
   }
 
   public void setType(String type) {
+    if (real == null) {
       identifier.type = type;
+    } else {
+      real.setType(type);
+    }
   }
 
   @XmlElement(name = "id")
   public String getId() {
+    if (real == null) {
       return identifier.id;
+    } else {
+      return real.getId();
+    }
   }
 
   public void setId(String id) {
+    if (real == null) {
       identifier.id = id;
+    } else {
+      real.setId(id);
+    }
   }
 
   public Identifier getIdentifier() {
+    if (real == null) {
       return identifier;
+    } else {
+      return real.getIdentifier();
+    }
   }
 
   public void setIdentifier(Identifier identifier) {
+    if (real == null) {
       this.identifier = identifier;
+    } else {
+      real.setIdentifier(identifier);
+    }
   }
 
   // required by JAXB
   @InterfaceAudience.Private
   @XmlElement(name = "info")
   public HashMap<String, Object> getInfoJAXB() {
+    if (real == null) {
       return info;
+    } else {
+      return real.getInfoJAXB();
+    }
   }
 
   public Map<String, Object> getInfo() {
+    if (real == null) {
       return info;
+    } else {
+      return real.getInfo();
+    }
   }
 
   public void setInfo(Map<String, Object> info) {
+    if (real == null) {
       if (info != null && !(info instanceof HashMap)) {
         this.info = new HashMap<String, Object>(info);
       } else {
         this.info = (HashMap<String, Object>) info;
       }
+    } else {
+      real.setInfo(info);
+    }
   }
 
   public void addInfo(Map<String, Object> info) {
+    if (real == null) {
       this.info.putAll(info);
+    } else {
+      real.addInfo(info);
+    }
   }
 
   public void addInfo(String key, Object value) {
+    if (real == null) {
       info.put(key, value);
+    } else {
+      real.addInfo(key, value);
+    }
   }
 
   // required by JAXB
   @InterfaceAudience.Private
   @XmlElement(name = "configs")
   public HashMap<String, Object> getConfigsJAXB() {
+    if (real == null) {
       return configs;
+    } else {
+      return real.getConfigsJAXB();
+    }
   }
 
   public Map<String, Object> getConfigs() {
+    if (real == null) {
       return configs;
+    } else {
+      return real.getConfigs();
+    }
   }
 
   public void setConfigs(Map<String, Object> configs) {
+    if (real == null) {
       if (configs != null && !(configs instanceof HashMap)) {
         this.configs = new HashMap<String, Object>(configs);
       } else {
         this.configs = (HashMap<String, Object>) configs;
       }
+    } else {
+      real.setConfigs(configs);
+    }
   }
 
   public void addConfigs(Map<String, Object> configs) {
+    if (real == null) {
       this.configs.putAll(configs);
+    } else {
+      real.addConfigs(configs);
+    }
   }
 
   public void addConfig(String key, Object value) {
+    if (real == null) {
       configs.put(key, value);
+    } else {
+      real.addConfig(key, value);
+    }
   }
 
   @XmlElement(name = "metrics")
   public Set<TimelineMetric> getMetrics() {
+    if (real == null) {
       return metrics;
+    } else {
+      return real.getMetrics();
+    }
   }
 
   public void setMetrics(Set<TimelineMetric> metrics) {
+    if (real == null) {
       this.metrics = metrics;
+    } else {
+      real.setMetrics(metrics);
+    }
   }
 
   public void addMetrics(Set<TimelineMetric> metrics) {
+    if (real == null) {
       this.metrics.addAll(metrics);
+    } else {
+      real.addMetrics(metrics);
+    }
   }
 
   public void addMetric(TimelineMetric metric) {
+    if (real == null) {
       metrics.add(metric);
+    } else {
+      real.addMetric(metric);
+    }
   }
 
   @XmlElement(name = "events")
   public Set<TimelineEvent> getEvents() {
+    if (real == null) {
       return events;
+    } else {
+      return real.getEvents();
+    }
   }
 
   public void setEvents(Set<TimelineEvent> events) {
+    if (real == null) {
       this.events = events;
+    } else {
+      real.setEvents(events);
+    }
   }
 
   public void addEvents(Set<TimelineEvent> events) {
+    if (real == null) {
       this.events.addAll(events);
+    } else {
+      real.addEvents(events);
+    }
   }
 
   public void addEvent(TimelineEvent event) {
+    if (real == null) {
       events.add(event);
+    } else {
+      real.addEvent(event);
+    }
   }
 
   public Map<String, Set<String>> getIsRelatedToEntities() {
+    if (real == null) {
       return isRelatedToEntities;
+    } else {
+      return real.getIsRelatedToEntities();
+    }
   }
 
   // required by JAXB
   @InterfaceAudience.Private
   @XmlElement(name = "isrelatedto")
   public HashMap<String, Set<String>> getIsRelatedToEntitiesJAXB() {
+    if (real == null) {
       return isRelatedToEntities;
+    } else {
+      return real.getIsRelatedToEntitiesJAXB();
+    }
   }
 
   public void setIsRelatedToEntities(
       Map<String, Set<String>> isRelatedToEntities) {
-    if (isRelatedToEntities != null && !(isRelatedToEntities instanceof HashMap)) {
-      this.isRelatedToEntities = new HashMap<String, Set<String>>(isRelatedToEntities);
+    if (real == null) {
+      if (isRelatedToEntities != null &&
+          !(isRelatedToEntities instanceof HashMap)) {
+        this.isRelatedToEntities =
+            new HashMap<String, Set<String>>(isRelatedToEntities);
       } else {
-      this.isRelatedToEntities = (HashMap<String, Set<String>>) isRelatedToEntities;
+        this.isRelatedToEntities =
+            (HashMap<String, Set<String>>) isRelatedToEntities;
+      }
+    } else {
+      real.setIsRelatedToEntities(isRelatedToEntities);
     }
   }
 
   public void addIsRelatedToEntities(
       Map<String, Set<String>> isRelatedToEntities) {
+    if (real == null) {
       for (Map.Entry<String, Set<String>> entry : isRelatedToEntities
           .entrySet()) {
         Set<String> ids = this.isRelatedToEntities.get(entry.getKey());
@@ -228,30 +370,47 @@ public class TimelineEntity {
         }
         ids.addAll(entry.getValue());
       }
+    } else {
+      real.addIsRelatedToEntities(isRelatedToEntities);
+    }
   }
 
   public void addIsRelatedToEntity(String type, String id) {
+    if (real == null) {
       Set<String> ids = isRelatedToEntities.get(type);
       if (ids == null) {
         ids = new HashSet<>();
         isRelatedToEntities.put(type, ids);
       }
       ids.add(id);
+    } else {
+      real.addIsRelatedToEntity(type, id);
+    }
   }
 
   // required by JAXB
   @InterfaceAudience.Private
   @XmlElement(name = "relatesto")
   public HashMap<String, Set<String>> getRelatesToEntitiesJAXB() {
+    if (real == null) {
       return relatesToEntities;
+    } else {
+      return real.getRelatesToEntitiesJAXB();
+    }
   }
 
   public Map<String, Set<String>> getRelatesToEntities() {
+    if (real == null) {
       return relatesToEntities;
+    } else {
+      return real.getRelatesToEntities();
+    }
   }
 
   public void addRelatesToEntities(Map<String, Set<String>> relatesToEntities) {
-    for (Map.Entry<String, Set<String>> entry : relatesToEntities.entrySet()) {
+    if (real == null) {
+      for (Map.Entry<String, Set<String>> entry : relatesToEntities
+          .entrySet()) {
         Set<String> ids = this.relatesToEntities.get(entry.getKey());
         if (ids == null) {
           ids = new HashSet<>();
@@ -259,42 +418,75 @@ public class TimelineEntity {
         }
         ids.addAll(entry.getValue());
       }
+    } else {
+      real.addRelatesToEntities(relatesToEntities);
+    }
   }
 
   public void addRelatesToEntity(String type, String id) {
+    if (real == null) {
       Set<String> ids = relatesToEntities.get(type);
       if (ids == null) {
         ids = new HashSet<>();
         relatesToEntities.put(type, ids);
       }
       ids.add(id);
+    } else {
+      real.addRelatesToEntity(type, id);
+    }
   }
 
   public void setRelatesToEntities(Map<String, Set<String>> relatesToEntities) {
-    if (relatesToEntities != null && !(relatesToEntities instanceof HashMap)) {
-      this.relatesToEntities = new HashMap<String, Set<String>>(relatesToEntities);
+    if (real == null) {
+      if (relatesToEntities != null &&
+          !(relatesToEntities instanceof HashMap)) {
+        this.relatesToEntities =
+            new HashMap<String, Set<String>>(relatesToEntities);
       } else {
-      this.relatesToEntities = (HashMap<String, Set<String>>) relatesToEntities;
+        this.relatesToEntities =
+            (HashMap<String, Set<String>>) relatesToEntities;
+      }
+    } else {
+      real.setRelatesToEntities(relatesToEntities);
     }
   }
 
   @XmlElement(name = "createdtime")
   public long getCreatedTime() {
+    if (real == null) {
       return createdTime;
+    } else {
+      return real.getCreatedTime();
+    }
   }
 
   public void setCreatedTime(long createdTime) {
+    if (real == null) {
       this.createdTime = createdTime;
+    } else {
+      real.setCreatedTime(createdTime);
+    }
   }
 
   @XmlElement(name = "modifiedtime")
   public long getModifiedTime() {
+    if (real == null) {
       return modifiedTime;
+    } else {
+      return real.getModifiedTime();
+    }
   }
 
   public void setModifiedTime(long modifiedTime) {
+    if (real == null) {
       this.modifiedTime = modifiedTime;
+    } else {
+      real.setModifiedTime(modifiedTime);
+    }
   }
 
+  protected TimelineEntity getReal() {
+    return real == null ? this : real;
+  }
 
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/timelineservice/UserEntity.java

@@ -20,16 +20,17 @@ package org.apache.hadoop.yarn.api.records.timelineservice;
 import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.classification.InterfaceStability;
 
-import javax.xml.bind.annotation.XmlAccessType;
-import javax.xml.bind.annotation.XmlAccessorType;
-import javax.xml.bind.annotation.XmlRootElement;
-
-@XmlRootElement(name = "user")
-@XmlAccessorType(XmlAccessType.NONE)
 @InterfaceAudience.Public
 @InterfaceStability.Unstable
-public class TimelineUser extends TimelineEntity {
-  public TimelineUser() {
+public class UserEntity extends TimelineEntity {
+  public UserEntity() {
     super(TimelineEntityType.YARN_USER.toString());
   }
+
+  public UserEntity(TimelineEntity entity) {
+    super(entity);
+    if (!entity.getType().equals(TimelineEntityType.YARN_USER.toString())) {
+      throw new IllegalArgumentException("Incompatible entity type: " + getId());
+    }
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/api/records/timelineservice/TestTimelineServiceRecords.java

@@ -23,8 +23,12 @@ import org.apache.commons.logging.LogFactory;
 import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.hadoop.yarn.api.records.ContainerId;
+import org.apache.hadoop.yarn.exceptions.YarnRuntimeException;
 import org.apache.hadoop.yarn.util.timeline.TimelineUtils;
 import org.junit.Test;
+import org.junit.Assert;
+
+import java.util.Collections;
 
 
 public class TestTimelineServiceRecords {
@@ -87,10 +91,10 @@ public class TestTimelineServiceRecords {
 
   @Test
   public void testFirstClassCitizenEntities() throws Exception {
-    TimelineUser user = new TimelineUser();
+    UserEntity user = new UserEntity();
     user.setId("test user id");
 
-    TimelineQueue queue = new TimelineQueue();
+    QueueEntity queue = new QueueEntity();
     queue.setId("test queue id");
 
 
@@ -98,20 +102,26 @@ public class TestTimelineServiceRecords {
     cluster.setId("test cluster id");
 
     FlowEntity flow1 = new FlowEntity();
-    flow1.setId("test flow id");
+    //flow1.setId("test flow id 1");
     flow1.setUser(user.getId());
-    flow1.setVersion("test flow version");
-    flow1.setRun("test run 1");
+    flow1.setName("test flow name 1");
+    flow1.setVersion("test flow version 1");
+    flow1.setRunId(1L);
 
     FlowEntity flow2 = new FlowEntity();
-    flow2.setId("test flow run id2");
+    //flow2.setId("test flow run id 2");
     flow2.setUser(user.getId());
-    flow1.setVersion("test flow version2");
-    flow2.setRun("test run 2");
+    flow2.setName("test flow name 2");
+    flow2.setVersion("test flow version 2");
+    flow2.setRunId(2L);
 
-    ApplicationEntity app = new ApplicationEntity();
-    app.setId(ApplicationId.newInstance(0, 1).toString());
-    app.setQueue(queue.getId());
+    ApplicationEntity app1 = new ApplicationEntity();
+    app1.setId(ApplicationId.newInstance(0, 1).toString());
+    app1.setQueue(queue.getId());
+
+    ApplicationEntity app2 = new ApplicationEntity();
+    app2.setId(ApplicationId.newInstance(0, 2).toString());
+    app2.setQueue(queue.getId());
 
     ApplicationAttemptEntity appAttempt = new ApplicationAttemptEntity();
     appAttempt.setId(ApplicationAttemptId.newInstance(
@@ -127,12 +137,14 @@ public class TestTimelineServiceRecords {
         .setParent(TimelineEntityType.YARN_CLUSTER.toString(), cluster.getId());
     flow1.addChild(TimelineEntityType.YARN_FLOW.toString(), flow2.getId());
     flow2.setParent(TimelineEntityType.YARN_FLOW.toString(), flow1.getId());
-    flow2.addChild(TimelineEntityType.YARN_APPLICATION.toString(), app.getId());
-    app.setParent(TimelineEntityType.YARN_FLOW.toString(), flow2.getId());
-    app.addChild(TimelineEntityType.YARN_APPLICATION_ATTEMPT.toString(),
+    flow2.addChild(TimelineEntityType.YARN_APPLICATION.toString(), app1.getId());
+    flow2.addChild(TimelineEntityType.YARN_APPLICATION.toString(), app2.getId());
+    app1.setParent(TimelineEntityType.YARN_FLOW.toString(), flow2.getId());
+    app1.addChild(TimelineEntityType.YARN_APPLICATION_ATTEMPT.toString(),
         appAttempt.getId());
     appAttempt
-        .setParent(TimelineEntityType.YARN_APPLICATION.toString(), app.getId());
+        .setParent(TimelineEntityType.YARN_APPLICATION.toString(), app1.getId());
+    app2.setParent(TimelineEntityType.YARN_FLOW.toString(), flow2.getId());
     appAttempt.addChild(TimelineEntityType.YARN_CONTAINER.toString(),
         container.getId());
     container.setParent(TimelineEntityType.YARN_APPLICATION_ATTEMPT.toString(),
@@ -141,14 +153,57 @@ public class TestTimelineServiceRecords {
     LOG.info(TimelineUtils.dumpTimelineRecordtoJSON(cluster, true));
     LOG.info(TimelineUtils.dumpTimelineRecordtoJSON(flow1, true));
     LOG.info(TimelineUtils.dumpTimelineRecordtoJSON(flow2, true));
-    LOG.info(TimelineUtils.dumpTimelineRecordtoJSON(app, true));
+    LOG.info(TimelineUtils.dumpTimelineRecordtoJSON(app1, true));
+    LOG.info(TimelineUtils.dumpTimelineRecordtoJSON(app2, true));
     LOG.info(TimelineUtils.dumpTimelineRecordtoJSON(appAttempt, true));
     LOG.info(TimelineUtils.dumpTimelineRecordtoJSON(container, true));
+
+
+    // Check parent/children APIs
+    Assert.assertNotNull(app1.getParent());
+    Assert.assertEquals(flow2.getType(), app1.getParent().getType());
+    Assert.assertEquals(flow2.getId(), app1.getParent().getId());
+    app1.addInfo(ApplicationEntity.PARENT_INFO_KEY, "invalid parent object");
+    try {
+      app1.getParent();
+      Assert.fail();
+    } catch (Exception e) {
+      Assert.assertTrue(e instanceof YarnRuntimeException);
+      Assert.assertTrue(e.getMessage().contains(
+          "Parent info is invalid identifier object"));
+    }
+
+    Assert.assertNotNull(app1.getChildren());
+    Assert.assertEquals(1, app1.getChildren().size());
+    Assert.assertEquals(
+        appAttempt.getType(), app1.getChildren().iterator().next().getType());
+    Assert.assertEquals(
+        appAttempt.getId(), app1.getChildren().iterator().next().getId());
+    app1.addInfo(ApplicationEntity.CHILDREN_INFO_KEY,
+        Collections.singletonList("invalid children set"));
+    try {
+      app1.getChildren();
+      Assert.fail();
+    } catch (Exception e) {
+      Assert.assertTrue(e instanceof YarnRuntimeException);
+      Assert.assertTrue(e.getMessage().contains(
+          "Children info is invalid identifier set"));
+    }
+    app1.addInfo(ApplicationEntity.CHILDREN_INFO_KEY,
+        Collections.singleton("invalid child object"));
+    try {
+      app1.getChildren();
+      Assert.fail();
+    } catch (Exception e) {
+      Assert.assertTrue(e instanceof YarnRuntimeException);
+      Assert.assertTrue(e.getMessage().contains(
+          "Children info contains invalid identifier object"));
+    }
   }
 
   @Test
   public void testUser() throws Exception {
-    TimelineUser user = new TimelineUser();
+    UserEntity user = new UserEntity();
     user.setId("test user id");
     user.addInfo("test info key 1", "test info value 1");
     user.addInfo("test info key 2", "test info value 2");
@@ -157,7 +212,7 @@ public class TestTimelineServiceRecords {
 
   @Test
   public void testQueue() throws Exception {
-    TimelineQueue queue = new TimelineQueue();
+    QueueEntity queue = new QueueEntity();
     queue.setId("test queue id");
     queue.addInfo("test info key 1", "test info value 1");
     queue.addInfo("test info key 2", "test info value 2");

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-tests/src/test/java/org/apache/hadoop/yarn/server/timelineservice/TestTimelineServiceClientIntegration.java

@@ -24,9 +24,12 @@ import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import org.apache.hadoop.security.UserGroupInformation;
 import org.apache.hadoop.util.ExitUtil;
+import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
-import org.apache.hadoop.yarn.api.records.timelineservice.TimelineEntity;
+import org.apache.hadoop.yarn.api.records.ContainerId;
+import org.apache.hadoop.yarn.api.records.timelineservice.*;
 import org.apache.hadoop.yarn.client.api.TimelineClient;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
 import org.apache.hadoop.yarn.exceptions.YarnException;
@@ -85,6 +88,45 @@ public class TestTimelineServiceClientIntegration {
     }
   }
 
+  @Test
+  public void testPutExtendedEntities() throws Exception {
+    ApplicationId appId = ApplicationId.newInstance(0, 1);
+    TimelineClient client =
+        TimelineClient.createTimelineClient(appId);
+    try {
+      // set the timeline service address manually
+      client.setTimelineServiceAddress(
+          collectorManager.getRestServerBindAddress());
+      client.init(new YarnConfiguration());
+      client.start();
+      ClusterEntity cluster = new ClusterEntity();
+      cluster.setId(YarnConfiguration.DEFAULT_RM_CLUSTER_ID);
+      FlowEntity flow = new FlowEntity();
+      flow.setUser(UserGroupInformation.getCurrentUser().getShortUserName());
+      flow.setName("test_flow_name");
+      flow.setVersion("test_flow_version");
+      flow.setRunId(1L);
+      flow.setParent(cluster.getType(), cluster.getId());
+      ApplicationEntity app = new ApplicationEntity();
+      app.setId(appId.toString());
+      flow.addChild(app.getType(), app.getId());
+      ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(appId, 1);
+      ApplicationAttemptEntity appAttempt = new ApplicationAttemptEntity();
+      appAttempt.setId(attemptId.toString());
+      ContainerId containerId = ContainerId.newContainerId(attemptId, 1);
+      ContainerEntity container = new ContainerEntity();
+      container.setId(containerId.toString());
+      UserEntity user = new UserEntity();
+      user.setId(UserGroupInformation.getCurrentUser().getShortUserName());
+      QueueEntity queue = new QueueEntity();
+      queue.setId("default_queue");
+      client.putEntities(cluster, flow, app, appAttempt, container, user, queue);
+      client.putEntitiesAsync(cluster, flow, app, appAttempt, container, user, queue);
+    } finally {
+      client.stop();
+    }
+  }
+
   private static class MockNodeTimelineCollectorManager extends
       NodeTimelineCollectorManager {
     public MockNodeTimelineCollectorManager() {

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-timelineservice/src/main/java/org/apache/hadoop/yarn/server/timelineservice/collector/TimelineCollectorWebService.java

@@ -43,7 +43,16 @@ import org.apache.hadoop.classification.InterfaceAudience.Public;
 import org.apache.hadoop.classification.InterfaceStability.Unstable;
 import org.apache.hadoop.security.UserGroupInformation;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
+import org.apache.hadoop.yarn.api.records.timelineservice.ApplicationAttemptEntity;
+import org.apache.hadoop.yarn.api.records.timelineservice.ApplicationEntity;
+import org.apache.hadoop.yarn.api.records.timelineservice.ClusterEntity;
+import org.apache.hadoop.yarn.api.records.timelineservice.ContainerEntity;
+import org.apache.hadoop.yarn.api.records.timelineservice.FlowEntity;
+import org.apache.hadoop.yarn.api.records.timelineservice.QueueEntity;
 import org.apache.hadoop.yarn.api.records.timelineservice.TimelineEntities;
+import org.apache.hadoop.yarn.api.records.timelineservice.TimelineEntity;
+import org.apache.hadoop.yarn.api.records.timelineservice.TimelineEntityType;
+import org.apache.hadoop.yarn.api.records.timelineservice.UserEntity;
 import org.apache.hadoop.yarn.util.ConverterUtils;
 import org.apache.hadoop.yarn.webapp.ForbiddenException;
 import org.apache.hadoop.yarn.webapp.NotFoundException;
@@ -142,7 +151,8 @@ public class TimelineCollectorWebService {
         LOG.error("Application: "+ appId + " is not found");
         throw new NotFoundException(); // different exception?
       }
-      collector.putEntities(entities, callerUgi);
+
+      collector.putEntities(processTimelineEntities(entities), callerUgi);
       return Response.ok().build();
     } catch (Exception e) {
       LOG.error("Error putting entities", e);
@@ -151,7 +161,7 @@ public class TimelineCollectorWebService {
     }
   }
 
-  private ApplicationId parseApplicationId(String appId) {
+  private static ApplicationId parseApplicationId(String appId) {
     try {
       if (appId != null) {
         return ConverterUtils.toApplicationId(appId.trim());
@@ -159,15 +169,16 @@ public class TimelineCollectorWebService {
         return null;
       }
     } catch (Exception e) {
+      LOG.error("Invalid application ID: " + appId);
       return null;
     }
   }
 
-  private void init(HttpServletResponse response) {
+  private static void init(HttpServletResponse response) {
     response.setContentType(null);
   }
 
-  private UserGroupInformation getUser(HttpServletRequest req) {
+  private static UserGroupInformation getUser(HttpServletRequest req) {
     String remoteUser = req.getRemoteUser();
     UserGroupInformation callerUgi = null;
     if (remoteUser != null) {
@@ -175,4 +186,50 @@ public class TimelineCollectorWebService {
     }
     return callerUgi;
   }
+
+  // The process may not be necessary according to the way we write the backend,
+  // but let's keep it for now in case we need to use sub-classes APIs in the
+  // future (e.g., aggregation).
+  private static TimelineEntities processTimelineEntities(
+      TimelineEntities entities) {
+    TimelineEntities entitiesToReturn = new TimelineEntities();
+    for (TimelineEntity entity : entities.getEntities()) {
+      TimelineEntityType type = null;
+      try {
+        type = TimelineEntityType.valueOf(entity.getType());
+      } catch (IllegalArgumentException e) {
+        type = null;
+      }
+      if (type != null) {
+        switch (type) {
+          case YARN_CLUSTER:
+            entitiesToReturn.addEntity(new ClusterEntity(entity));
+            break;
+          case YARN_FLOW:
+            entitiesToReturn.addEntity(new FlowEntity(entity));
+            break;
+          case YARN_APPLICATION:
+            entitiesToReturn.addEntity(new ApplicationEntity(entity));
+            break;
+          case YARN_APPLICATION_ATTEMPT:
+            entitiesToReturn.addEntity(new ApplicationAttemptEntity(entity));
+            break;
+          case YARN_CONTAINER:
+            entitiesToReturn.addEntity(new ContainerEntity(entity));
+            break;
+          case YARN_QUEUE:
+            entitiesToReturn.addEntity(new QueueEntity(entity));
+            break;
+          case YARN_USER:
+            entitiesToReturn.addEntity(new UserEntity(entity));
+            break;
+          default:
+            break;
+        }
+      } else {
+        entitiesToReturn.addEntity(entity);
+      }
+    }
+    return entitiesToReturn;
+  }
 }