3131import org .apache .hadoop .yarn .api .records .ApplicationAttemptId ;
3232import org .apache .hadoop .yarn .api .records .ApplicationId ;
3333import org .apache .hadoop .yarn .api .records .ContainerId ;
34+ import org .apache .hadoop .yarn .api .records .ContainerStatus ;
3435import org .apache .hadoop .yarn .api .records .NodeId ;
36+ import org .apache .hadoop .yarn .api .records .timelineservice .ContainerEntity ;
3537import org .apache .hadoop .yarn .api .records .timelineservice .TimelineEntity ;
3638import org .apache .hadoop .yarn .api .records .timelineservice .TimelineMetric ;
3739import org .apache .hadoop .yarn .client .api .impl .TimelineV2ClientImpl ;
3840import org .apache .hadoop .yarn .conf .YarnConfiguration ;
41+ import org .apache .hadoop .yarn .event .AsyncDispatcher ;
42+ import org .apache .hadoop .yarn .event .DrainDispatcher ;
3943import org .apache .hadoop .yarn .exceptions .YarnException ;
44+ import org .apache .hadoop .yarn .server .metrics .ContainerMetricsConstants ;
4045import org .apache .hadoop .yarn .server .nodemanager .Context ;
46+ import org .apache .hadoop .yarn .server .nodemanager .containermanager .application .ApplicationContainerFinishedEvent ;
4147import org .apache .hadoop .yarn .server .nodemanager .containermanager .container .Container ;
4248import org .apache .hadoop .yarn .util .ResourceCalculatorProcessTree ;
4349import org .junit .Assert ;
4450import org .junit .Test ;
51+ import org .junit .After ;
52+ import org .junit .Before ;
4553
4654public class TestNMTimelinePublisher {
4755 private static final String MEMORY_ID = "MEMORY" ;
4856 private static final String CPU_ID = "CPU" ;
4957
50- @ Test
51- public void testContainerResourceUsage () {
52- Context context = mock (Context .class );
53- @ SuppressWarnings ("unchecked" )
54- final DummyTimelineClient timelineClient = new DummyTimelineClient (null );
55- when (context .getNodeId ()).thenReturn (NodeId .newInstance ("localhost" , 0 ));
58+ private NMTimelinePublisher publisher ;
59+ private DummyTimelineClient timelineClient ;
60+ private Configuration conf ;
61+ private DrainDispatcher dispatcher ;
5662
57- Configuration conf = new Configuration ();
63+
64+ @ Before public void setup () throws Exception {
65+ conf = new Configuration ();
5866 conf .setBoolean (YarnConfiguration .TIMELINE_SERVICE_ENABLED , true );
5967 conf .setFloat (YarnConfiguration .TIMELINE_SERVICE_VERSION , 2.0f );
68+ conf .setLong (YarnConfiguration .ATS_APP_COLLECTOR_LINGER_PERIOD_IN_MS ,
69+ 3000L );
70+ timelineClient = new DummyTimelineClient (null );
71+ Context context = createMockContext ();
72+ dispatcher = new DrainDispatcher ();
6073
61- NMTimelinePublisher publisher = new NMTimelinePublisher (context ) {
74+ publisher = new NMTimelinePublisher (context ) {
6275 public void createTimelineClient (ApplicationId appId ) {
6376 if (!getAppToClientMap ().containsKey (appId )) {
6477 timelineClient .init (getConfig ());
6578 timelineClient .start ();
6679 getAppToClientMap ().put (appId , timelineClient );
6780 }
6881 }
82+
83+ @ Override protected AsyncDispatcher createDispatcher () {
84+ return dispatcher ;
85+ }
6986 };
7087 publisher .init (conf );
7188 publisher .start ();
89+ }
90+
91+ private Context createMockContext () {
92+ Context context = mock (Context .class );
93+ when (context .getNodeId ()).thenReturn (NodeId .newInstance ("localhost" , 0 ));
94+ return context ;
95+ }
96+
97+ @ After public void tearDown () throws Exception {
98+ if (publisher != null ) {
99+ publisher .stop ();
100+ }
101+ if (timelineClient != null ) {
102+ timelineClient .stop ();
103+ }
104+ }
105+
106+ @ Test public void testPublishContainerFinish () throws Exception {
107+ ApplicationId appId = ApplicationId .newInstance (0 , 2 );
108+ ApplicationAttemptId appAttemptId =
109+ ApplicationAttemptId .newInstance (appId , 1 );
110+ ContainerId cId = ContainerId .newContainerId (appAttemptId , 1 );
111+
112+ String diag = "test-diagnostics" ;
113+ int exitStatus = 0 ;
114+ ContainerStatus cStatus = mock (ContainerStatus .class );
115+ when (cStatus .getContainerId ()).thenReturn (cId );
116+ when (cStatus .getDiagnostics ()).thenReturn (diag );
117+ when (cStatus .getExitStatus ()).thenReturn (exitStatus );
118+ long timeStamp = System .currentTimeMillis ();
119+
120+ ApplicationContainerFinishedEvent finishedEvent =
121+ new ApplicationContainerFinishedEvent (cStatus , timeStamp );
122+
123+ publisher .createTimelineClient (appId );
124+ publisher .publishApplicationEvent (finishedEvent );
125+ publisher .stopTimelineClient (appId );
126+ dispatcher .await ();
127+
128+ ContainerEntity cEntity = new ContainerEntity ();
129+ cEntity .setId (cId .toString ());
130+ TimelineEntity [] lastPublishedEntities =
131+ timelineClient .getLastPublishedEntities ();
132+
133+ Assert .assertNotNull (lastPublishedEntities );
134+ Assert .assertEquals (1 , lastPublishedEntities .length );
135+ TimelineEntity entity = lastPublishedEntities [0 ];
136+ Assert .assertTrue (cEntity .equals (entity ));
137+ Assert .assertEquals (diag ,
138+ entity .getInfo ().get (ContainerMetricsConstants .DIAGNOSTICS_INFO ));
139+ Assert .assertEquals (exitStatus ,
140+ entity .getInfo ().get (ContainerMetricsConstants .EXIT_STATUS_INFO ));
141+ }
142+
143+ @ Test public void testContainerResourceUsage () {
72144 ApplicationId appId = ApplicationId .newInstance (0 , 1 );
73145 publisher .createTimelineClient (appId );
74146 Container aContainer = mock (Container .class );
75- when (aContainer .getContainerId ()).thenReturn (ContainerId .newContainerId (
76- ApplicationAttemptId .newInstance (appId , 1 ),
77- 0L ));
147+ when (aContainer .getContainerId ()).thenReturn (ContainerId
148+ .newContainerId (ApplicationAttemptId .newInstance (appId , 1 ), 0L ));
78149 publisher .reportContainerResourceUsage (aContainer , 1024L , 8F );
79150 verifyPublishedResourceUsageMetrics (timelineClient , 1024L , 8 );
80151 timelineClient .reset ();
@@ -91,7 +162,6 @@ public void createTimelineClient(ApplicationId appId) {
91162 (float ) ResourceCalculatorProcessTree .UNAVAILABLE );
92163 verifyPublishedResourceUsageMetrics (timelineClient , 1024L ,
93164 ResourceCalculatorProcessTree .UNAVAILABLE );
94- publisher .stop ();
95165 }
96166
97167 private void verifyPublishedResourceUsageMetrics (
@@ -151,8 +221,12 @@ public DummyTimelineClient(ApplicationId appId) {
151221
152222 private TimelineEntity [] lastPublishedEntities ;
153223
154- @ Override
155- public void putEntitiesAsync (TimelineEntity ... entities )
224+ @ Override public void putEntitiesAsync (TimelineEntity ... entities )
225+ throws IOException , YarnException {
226+ this .lastPublishedEntities = entities ;
227+ }
228+
229+ @ Override public void putEntities (TimelineEntity ... entities )
156230 throws IOException , YarnException {
157231 this .lastPublishedEntities = entities ;
158232 }
0 commit comments