From 57e2498cd49686ce231fff2e0c8df142bbbf58d2 Mon Sep 17 00:00:00 2001
From: Zhijie Shen <zjshen@apache.org>
Date: Mon, 3 Aug 2015 16:55:44 -0700
Subject: [PATCH] YARN-3993. Changed to use the AM flag in ContainerContext
 determine AM container in TestPerNodeTimelineCollectorsAuxService.
 Contributed by Sunil G.

(cherry picked from commit 9e48f9ff2ce08f3dcdd8d60bacb697664b92196f)
---
 .../PerNodeTimelineCollectorsAuxService.java        | 13 +++----------
 .../TestPerNodeTimelineCollectorsAuxService.java    |  5 +++++
 2 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-timelineservice/src/main/java/org/apache/hadoop/yarn/server/timelineservice/collector/PerNodeTimelineCollectorsAuxService.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-timelineservice/src/main/java/org/apache/hadoop/yarn/server/timelineservice/collector/PerNodeTimelineCollectorsAuxService.java
index 3ede97aefed..befaa83ba67 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-timelineservice/src/main/java/org/apache/hadoop/yarn/server/timelineservice/collector/PerNodeTimelineCollectorsAuxService.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-timelineservice/src/main/java/org/apache/hadoop/yarn/server/timelineservice/collector/PerNodeTimelineCollectorsAuxService.java
@@ -38,6 +38,7 @@ import org.apache.hadoop.yarn.server.api.AuxiliaryService;
 import org.apache.hadoop.yarn.server.api.ContainerContext;
 import org.apache.hadoop.yarn.server.api.ContainerInitializationContext;
 import org.apache.hadoop.yarn.server.api.ContainerTerminationContext;
+import org.apache.hadoop.yarn.server.api.ContainerType;
 
 import com.google.common.annotations.VisibleForTesting;
 
@@ -119,7 +120,7 @@ public class PerNodeTimelineCollectorsAuxService extends AuxiliaryService {
   public void initializeContainer(ContainerInitializationContext context) {
     // intercept the event of the AM container being created and initialize the
     // app level collector service
-    if (isApplicationMaster(context)) {
+    if (context.getContainerType() == ContainerType.APPLICATION_MASTER) {
       ApplicationId appId = context.getContainerId().
           getApplicationAttemptId().getApplicationId();
       addApplication(appId);
@@ -135,21 +136,13 @@ public class PerNodeTimelineCollectorsAuxService extends AuxiliaryService {
   public void stopContainer(ContainerTerminationContext context) {
     // intercept the event of the AM container being stopped and remove the app
     // level collector service
-    if (isApplicationMaster(context)) {
+    if (context.getContainerType() == ContainerType.APPLICATION_MASTER) {
       ApplicationId appId = context.getContainerId().
           getApplicationAttemptId().getApplicationId();
       removeApplication(appId);
     }
   }
 
-  private boolean isApplicationMaster(ContainerContext context) {
-    // TODO this is based on a (shaky) assumption that the container id (the
-    // last field of the full container id) for an AM is always 1
-    // we want to make this much more reliable
-    ContainerId containerId = context.getContainerId();
-    return containerId.getContainerId() == 1L;
-  }
-
   @VisibleForTesting
   boolean hasApplication(ApplicationId appId) {
     return collectorManager.containsTimelineCollector(appId);
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-timelineservice/src/test/java/org/apache/hadoop/yarn/server/timelineservice/collector/TestPerNodeTimelineCollectorsAuxService.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-timelineservice/src/test/java/org/apache/hadoop/yarn/server/timelineservice/collector/TestPerNodeTimelineCollectorsAuxService.java
index b1a5b046b21..7cc612d9204 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-timelineservice/src/test/java/org/apache/hadoop/yarn/server/timelineservice/collector/TestPerNodeTimelineCollectorsAuxService.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-timelineservice/src/test/java/org/apache/hadoop/yarn/server/timelineservice/collector/TestPerNodeTimelineCollectorsAuxService.java
@@ -39,6 +39,7 @@ import org.apache.hadoop.yarn.exceptions.YarnException;
 import org.apache.hadoop.yarn.server.api.CollectorNodemanagerProtocol;
 import org.apache.hadoop.yarn.server.api.ContainerInitializationContext;
 import org.apache.hadoop.yarn.server.api.ContainerTerminationContext;
+import org.apache.hadoop.yarn.server.api.ContainerType;
 import org.apache.hadoop.yarn.server.api.protocolrecords.GetTimelineCollectorContextRequest;
 import org.apache.hadoop.yarn.server.api.protocolrecords.GetTimelineCollectorContextResponse;
 import org.junit.After;
@@ -94,6 +95,8 @@ public class TestPerNodeTimelineCollectorsAuxService {
     ContainerTerminationContext context =
         mock(ContainerTerminationContext.class);
     when(context.getContainerId()).thenReturn(containerId);
+    when(context.getContainerType()).thenReturn(
+        ContainerType.APPLICATION_MASTER);
     auxService.stopContainer(context);
     // auxService should not have that app
     assertFalse(auxService.hasApplication(appAttemptId.getApplicationId()));
@@ -138,6 +141,8 @@ public class TestPerNodeTimelineCollectorsAuxService {
     ContainerInitializationContext context =
         mock(ContainerInitializationContext.class);
     when(context.getContainerId()).thenReturn(containerId);
+    when(context.getContainerType()).thenReturn(
+        ContainerType.APPLICATION_MASTER);
     auxService.initializeContainer(context);
     return auxService;
   }