YARN-11656 RMStateStore event queue blocked

- fix tests

Author: K0K0V0K

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/recovery/RMStateStoreTestBase.java

@@ -33,6 +33,8 @@
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 
 import javax.crypto.SecretKey;
 
@@ -74,6 +76,7 @@
 import org.apache.hadoop.yarn.server.resourcemanager.reservation.ReservationSystemTestUtil;
 import org.apache.hadoop.yarn.server.resourcemanager.reservation.ReservationSystemUtil;
 import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMApp;
+import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMAppEvent;
 import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMAppState;
 import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.AggregateAppResourceUsage;
 import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.RMAppAttempt;
@@ -99,9 +102,8 @@ public class RMStateStoreTestBase {
 
   static class TestDispatcher implements Dispatcher, EventHandler<Event> {
 
-    ApplicationAttemptId attemptId;
-
-    boolean notified = false;
+    private final Set<ApplicationAttemptId> handledAttempt = ConcurrentHashMap.newKeySet();
+    private final Set<ApplicationId> handledApps = ConcurrentHashMap.newKeySet();
 
     @SuppressWarnings("rawtypes")
     @Override
@@ -113,11 +115,16 @@ public void register(Class<? extends Enum> eventType,
     public void handle(Event event) {
       if (event instanceof RMAppAttemptEvent) {
         RMAppAttemptEvent rmAppAttemptEvent = (RMAppAttemptEvent) event;
-        assertEquals(attemptId, rmAppAttemptEvent.getApplicationAttemptId());
-      }
-      notified = true;
-      synchronized (this) {
-        notifyAll();
+        synchronized (handledAttempt) {
+          handledAttempt.add(rmAppAttemptEvent.getApplicationAttemptId());
+          handledAttempt.notifyAll();
+        }
+      } else if (event instanceof RMAppEvent) {
+        RMAppEvent rmAppEvent = (RMAppEvent) event;
+        synchronized (handledApps) {
+          handledApps.add(rmAppEvent.getApplicationId());
+          handledApps.notifyAll();
+        }
       }
     }
 
@@ -127,6 +134,38 @@ public EventHandler<Event> getEventHandler() {
       return this;
     }
 
+    void waitNotify(ApplicationAttemptId attemptId) {
+      long startTime = System.currentTimeMillis();
+      while(!handledAttempt.contains(attemptId)) {
+        synchronized (handledAttempt) {
+          try {
+            handledAttempt.wait(100);
+          } catch (InterruptedException e) {
+            LOG.trace("Interrupted", e);
+          }
+        }
+        if(System.currentTimeMillis() - startTime > 1000*60) {
+          fail("Timed out attempt store notification");
+        }
+      }
+    }
+
+    void waitNotify(ApplicationId applicationId) {
+      long startTime = System.currentTimeMillis();
+      while(!handledApps.contains(applicationId)) {
+        synchronized (handledApps) {
+          try {
+            handledApps.wait(100);
+          } catch (InterruptedException e) {
+            LOG.trace("Interrupted", e);
+          }
+        }
+        if(System.currentTimeMillis() - startTime > 1000*60) {
+          fail("Timed out attempt store notification");
+        }
+      }
+    }
+
   }
 
   public static class StoreStateVerifier {
@@ -148,23 +187,6 @@ public long getEpochRange() {
     return epochRange;
   }
 
-  void waitNotify(TestDispatcher dispatcher) {
-    long startTime = System.currentTimeMillis();
-    while(!dispatcher.notified) {
-      synchronized (dispatcher) {
-        try {
-          dispatcher.wait(1000);
-        } catch (InterruptedException e) {
-          e.printStackTrace();
-        }
-      }
-      if(System.currentTimeMillis() - startTime > 1000*60) {
-        fail("Timed out attempt store notification");
-      }
-    }
-    dispatcher.notified = false;
-  }
-
   protected RMApp storeApp(RMStateStore store, ApplicationId appId,
       long submitTime, long startTime) throws Exception {
     ApplicationSubmissionContext context =
@@ -204,17 +226,15 @@ protected RMAppAttempt storeAttempt(RMStateStore store,
         .thenReturn(mockRmAppAttemptMetrics);
     when(mockRmAppAttemptMetrics.getAggregateAppResourceUsage())
         .thenReturn(new AggregateAppResourceUsage(new HashMap<>()));
-    dispatcher.attemptId = attemptId;
     store.storeNewApplicationAttempt(mockAttempt);
-    waitNotify(dispatcher);
+    dispatcher.waitNotify(attemptId);
     return mockAttempt;
   }
 
   protected void updateAttempt(RMStateStore store, TestDispatcher dispatcher,
       ApplicationAttemptStateData attemptState) {
-    dispatcher.attemptId = attemptState.getAttemptId();
     store.updateApplicationAttemptState(attemptState);
-    waitNotify(dispatcher);
+    dispatcher.waitNotify(attemptState.getAttemptId());
   }
 
   void testRMAppStateStore(RMStateStoreHelper stateStoreHelper)
@@ -642,6 +662,8 @@ private ArrayList<RMApp> createAndStoreApps(
   public void testDeleteStore(RMStateStoreHelper stateStoreHelper)
       throws Exception {
     RMStateStore store = stateStoreHelper.getRMStateStore();
+    TestDispatcher dispatcher = new TestDispatcher();
+    store.setRMDispatcher(dispatcher);
     ArrayList<RMApp> appList = createAndStoreApps(stateStoreHelper, store, 5);
     store.deleteStore();
     // verify apps deleted

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/recovery/TestZKRMStateStore.java

@@ -536,6 +536,7 @@ public void testFencedState() throws Exception {
     RMApp mockApp = mock(RMApp.class);
     ApplicationSubmissionContext context =
       new ApplicationSubmissionContextPBImpl();
+    context.setApplicationId(mock(ApplicationId.class));
     when(mockApp.getSubmitTime()).thenReturn(submitTime);
     when(mockApp.getStartTime()).thenReturn(startTime);
     when(mockApp.getApplicationSubmissionContext()).thenReturn(context);
@@ -875,7 +876,7 @@ private void finishAppWithAttempts(RMState state, RMStateStore store,
       appStateNew.attempts.putAll(appState.attempts);
     }
     store.updateApplicationState(appStateNew);
-    waitNotify(dispatcher);
+    dispatcher.waitNotify(appId);
     Container container = new ContainerPBImpl();
     container.setId(ContainerId.newContainerId(attemptId, 1));
     ApplicationAttemptStateData newAttemptState =
@@ -893,7 +894,7 @@ private void storeAppWithAttempts(RMStateStore store,
   private void storeApp(RMStateStore store, TestDispatcher dispatcher,
       ApplicationId appId, long submitTime, long startTime) throws Exception {
     storeApp(store, appId, submitTime, startTime);
-    waitNotify(dispatcher);
+    dispatcher.waitNotify(appId);
   }
 
   private void storeAppWithAttempts(RMStateStore store,
@@ -1000,7 +1001,7 @@ public void testAppNodeSplit() throws Exception {
     // Store app2 with app id application_1352994193343_120213.
     ApplicationId appId21 = ApplicationId.newInstance(1352994193343L, 120213);
     storeApp(store, appId21, submitTime, startTime);
-    waitNotify(dispatcher);
+    dispatcher.waitNotify(appId21);
 
     // Store another app which will be removed.
     ApplicationId appIdRemoved = ApplicationId.newInstance(1352994193343L, 2);
@@ -1170,7 +1171,7 @@ public void testAppNodeSplitChangeAcrossRestarts() throws Exception {
     ApplicationId appId71 = ApplicationId.newInstance(1442994195087L, 7);
     //storeApp(store, dispatcher, appId71, submitTime, startTime);
     storeApp(store, appId71, submitTime, startTime);
-    waitNotify(dispatcher);
+    dispatcher.waitNotify(appId71);
     ApplicationAttemptId attemptId71 =
         ApplicationAttemptId.newInstance(appId71, 1);
     storeAttempt(store, ApplicationAttemptId.newInstance(appId71, 1),

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/recovery/TestZKRMStateStorePerf.java

@@ -192,7 +192,7 @@ public int run(String[] args) {
             + e.getMessage());
         return -1;
       }
-      waitNotify(dispatcher);
+      dispatcher.waitNotify(appId);
       rmApps.add(app);
     }