YARN-6411. Clean up the overwrite of createDispatcher() in subclass of MockRM. Contributed by Yufei Gu

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/api/impl/TestAMRMClientOnRMRestart.java

@@ -49,8 +49,6 @@ import org.apache.hadoop.yarn.api.records.UpdateContainerRequest;
 import org.apache.hadoop.yarn.client.api.AMRMClient;
 import org.apache.hadoop.yarn.client.api.AMRMClient.ContainerRequest;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
-import org.apache.hadoop.yarn.event.Dispatcher;
-import org.apache.hadoop.yarn.event.DrainDispatcher;
 import org.apache.hadoop.yarn.event.EventHandler;
 import org.apache.hadoop.yarn.ipc.YarnRPC;
 import org.apache.hadoop.yarn.security.AMRMTokenIdentifier;
@@ -126,22 +124,20 @@ public class TestAMRMClientOnRMRestart {
     // Phase-1 Start 1st RM
     MyResourceManager rm1 = new MyResourceManager(conf, memStore);
     rm1.start();
-    DrainDispatcher dispatcher =
-        (DrainDispatcher) rm1.getRMContext().getDispatcher();
 
     // Submit the application
     RMApp app = rm1.submitApp(1024);
-    dispatcher.await();
+    rm1.drainEvents();
 
     MockNM nm1 = new MockNM("h1:1234", 15120, rm1.getResourceTrackerService());
     nm1.registerNode();
     nm1.nodeHeartbeat(true); // Node heartbeat
-    dispatcher.await();
+    rm1.drainEvents();
 
     ApplicationAttemptId appAttemptId =
         app.getCurrentAppAttempt().getAppAttemptId();
     rm1.sendAMLaunched(appAttemptId);
-    dispatcher.await();
+    rm1.drainEvents();
 
     org.apache.hadoop.security.token.Token<AMRMTokenIdentifier> token =
         rm1.getRMContext().getRMApps().get(appAttemptId.getApplicationId())
@@ -176,7 +172,7 @@ public class TestAMRMClientOnRMRestart {
     blacklistAdditions.remove("h2");// remove from local list
 
     AllocateResponse allocateResponse = amClient.allocate(0.1f);
-    dispatcher.await();
+    rm1.drainEvents();
     Assert.assertEquals("No of assignments must be 0", 0, allocateResponse
         .getAllocatedContainers().size());
 
@@ -189,10 +185,10 @@ public class TestAMRMClientOnRMRestart {
     // Step-2 : NM heart beat is sent.
     // On 2nd AM allocate request, RM allocates 3 containers to AM
     nm1.nodeHeartbeat(true); // Node heartbeat
-    dispatcher.await();
+    rm1.drainEvents();
 
     allocateResponse = amClient.allocate(0.2f);
-    dispatcher.await();
+    rm1.drainEvents();
     // 3 containers are allocated i.e for cRequest1, cRequest2 and cRequest3.
     Assert.assertEquals("No of assignments must be 0", 3, allocateResponse
         .getAllocatedContainers().size());
@@ -207,7 +203,7 @@ public class TestAMRMClientOnRMRestart {
     amClient.removeContainerRequest(cRequest3);
 
     allocateResponse = amClient.allocate(0.2f);
-    dispatcher.await();
+    rm1.drainEvents();
     Assert.assertEquals("No of assignments must be 0", 0, allocateResponse
         .getAllocatedContainers().size());
     assertAsksAndReleases(4, 0, rm1);
@@ -233,7 +229,7 @@ public class TestAMRMClientOnRMRestart {
     // request
     nm1.nodeHeartbeat(containerId.getApplicationAttemptId(),
         containerId.getContainerId(), ContainerState.RUNNING);
-    dispatcher.await();
+    rm1.drainEvents();
     amClient.requestContainerUpdate(
         container, UpdateContainerRequest.newInstance(
             container.getVersion(), container.getId(),
@@ -242,7 +238,7 @@ public class TestAMRMClientOnRMRestart {
     it.remove();
 
     allocateResponse = amClient.allocate(0.3f);
-    dispatcher.await();
+    rm1.drainEvents();
     Assert.assertEquals("No of assignments must be 0", 0, allocateResponse
         .getAllocatedContainers().size());
     assertAsksAndReleases(3, pendingRelease, rm1);
@@ -258,7 +254,6 @@ public class TestAMRMClientOnRMRestart {
     rm2.start();
     nm1.setResourceTrackerService(rm2.getResourceTrackerService());
     ((MyAMRMClientImpl) amClient).updateRMProxy(rm2);
-    dispatcher = (DrainDispatcher) rm2.getRMContext().getDispatcher();
 
     // NM should be rebooted on heartbeat, even first heartbeat for nm2
     NodeHeartbeatResponse hbResponse = nm1.nodeHeartbeat(true);
@@ -274,7 +269,7 @@ public class TestAMRMClientOnRMRestart {
         Collections.singletonList(
             containerId.getApplicationAttemptId().getApplicationId()));
     nm1.nodeHeartbeat(true);
-    dispatcher.await();
+    rm2.drainEvents();
 
     blacklistAdditions.add("h3");
     amClient.updateBlacklist(blacklistAdditions, null);
@@ -296,7 +291,7 @@ public class TestAMRMClientOnRMRestart {
     // containerRequest and blacklisted nodes.
     // Intern RM send resync command,AMRMClient resend allocate request
     allocateResponse = amClient.allocate(0.3f);
-    dispatcher.await();
+    rm2.drainEvents();
 
     completedContainer =
         allocateResponse.getCompletedContainersStatuses().size();
@@ -313,7 +308,7 @@ public class TestAMRMClientOnRMRestart {
 
     // Step-5 : Allocater after resync command
     allocateResponse = amClient.allocate(0.5f);
-    dispatcher.await();
+    rm2.drainEvents();
     Assert.assertEquals("No of assignments must be 0", 0, allocateResponse
         .getAllocatedContainers().size());
 
@@ -326,10 +321,10 @@ public class TestAMRMClientOnRMRestart {
     int count = 5;
     while (count-- > 0) {
       nm1.nodeHeartbeat(true);
-      dispatcher.await();
+      rm2.drainEvents();
 
       allocateResponse = amClient.allocate(0.5f);
-      dispatcher.await();
+      rm2.drainEvents();
       noAssignedContainer += allocateResponse.getAllocatedContainers().size();
       if (noAssignedContainer == 3) {
         break;
@@ -358,22 +353,20 @@ public class TestAMRMClientOnRMRestart {
     // Phase-1 Start 1st RM
     MyResourceManager rm1 = new MyResourceManager(conf, memStore);
     rm1.start();
-    DrainDispatcher dispatcher =
-        (DrainDispatcher) rm1.getRMContext().getDispatcher();
 
     // Submit the application
     RMApp app = rm1.submitApp(1024);
-    dispatcher.await();
+    rm1.drainEvents();
 
     MockNM nm1 = new MockNM("h1:1234", 15120, rm1.getResourceTrackerService());
     nm1.registerNode();
     nm1.nodeHeartbeat(true); // Node heartbeat
-    dispatcher.await();
+    rm1.drainEvents();
 
     ApplicationAttemptId appAttemptId =
         app.getCurrentAppAttempt().getAppAttemptId();
     rm1.sendAMLaunched(appAttemptId);
-    dispatcher.await();
+    rm1.drainEvents();
 
     org.apache.hadoop.security.token.Token<AMRMTokenIdentifier> token =
         rm1.getRMContext().getRMApps().get(appAttemptId.getApplicationId())
@@ -393,7 +386,6 @@ public class TestAMRMClientOnRMRestart {
     rm2.start();
     nm1.setResourceTrackerService(rm2.getResourceTrackerService());
     ((MyAMRMClientImpl) amClient).updateRMProxy(rm2);
-    dispatcher = (DrainDispatcher) rm2.getRMContext().getDispatcher();
 
     // NM should be rebooted on heartbeat, even first heartbeat for nm2
     NodeHeartbeatResponse hbResponse = nm1.nodeHeartbeat(true);
@@ -409,7 +401,7 @@ public class TestAMRMClientOnRMRestart {
             Priority.newInstance(0), 0);
     nm1.registerNode(Arrays.asList(containerReport), null);
     nm1.nodeHeartbeat(true);
-    dispatcher.await();
+    rm2.drainEvents();
 
     amClient.unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED,
         null, null);
@@ -421,7 +413,6 @@ public class TestAMRMClientOnRMRestart {
     amClient.stop();
     rm1.stop();
     rm2.stop();
-
   }
 
 
@@ -439,22 +430,20 @@ public class TestAMRMClientOnRMRestart {
     // start first RM
     MyResourceManager2 rm1 = new MyResourceManager2(conf, memStore);
     rm1.start();
-    DrainDispatcher dispatcher =
-        (DrainDispatcher) rm1.getRMContext().getDispatcher();
     Long startTime = System.currentTimeMillis();
     // Submit the application
     RMApp app = rm1.submitApp(1024);
-    dispatcher.await();
+    rm1.drainEvents();
 
     MockNM nm1 = new MockNM("h1:1234", 15120, rm1.getResourceTrackerService());
     nm1.registerNode();
     nm1.nodeHeartbeat(true); // Node heartbeat
-    dispatcher.await();
+    rm1.drainEvents();
 
     ApplicationAttemptId appAttemptId =
         app.getCurrentAppAttempt().getAppAttemptId();
     rm1.sendAMLaunched(appAttemptId);
-    dispatcher.await();
+    rm1.drainEvents();
 
     AMRMTokenSecretManager amrmTokenSecretManagerForRM1 =
         rm1.getRMContext().getAMRMTokenSecretManager();
@@ -513,7 +502,6 @@ public class TestAMRMClientOnRMRestart {
     rm2.start();
     nm1.setResourceTrackerService(rm2.getResourceTrackerService());
     ((MyAMRMClientImpl) amClient).updateRMProxy(rm2);
-    dispatcher = (DrainDispatcher) rm2.getRMContext().getDispatcher();
 
     AMRMTokenSecretManager amrmTokenSecretManagerForRM2 =
         rm2.getRMContext().getAMRMTokenSecretManager();
@@ -615,11 +603,6 @@ public class TestAMRMClientOnRMRestart {
       MyResourceManager.setClusterTimeStamp(fakeClusterTimeStamp);
     }
 
-    @Override
-    protected Dispatcher createDispatcher() {
-      return new DrainDispatcher();
-    }
-
     @Override
     protected EventHandler<SchedulerEvent> createSchedulerEventDispatcher() {
       // Dispatch inline for test sanity

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

@@ -30,14 +30,9 @@ import org.apache.hadoop.security.authorize.AccessControlList;
 import org.apache.hadoop.service.Service.STATE;
 import org.apache.hadoop.yarn.api.ApplicationClientProtocol;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
-import org.apache.hadoop.yarn.event.Dispatcher;
-import org.apache.hadoop.yarn.event.DrainDispatcher;
 import org.apache.hadoop.yarn.ipc.YarnRPC;
 import org.junit.Before;
 
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
-
 public abstract class ACLsTestBase {
 
   protected static final String COMMON_USER = "common_user";
@@ -80,11 +75,6 @@ public abstract class ACLsTestBase {
                 .getRMDelegationTokenSecretManager());
       }
 
-      @Override
-      protected Dispatcher createDispatcher() {
-        return new DrainDispatcher();
-      }
-
       @Override
       protected void doSecureLogin() throws IOException {
       }

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

@@ -26,22 +26,17 @@ import org.apache.hadoop.ha.HAServiceProtocol;
 import org.apache.hadoop.ha.HAServiceProtocol.HAServiceState;
 import org.apache.hadoop.ha.HAServiceProtocol.StateChangeRequestInfo;
 import org.apache.hadoop.yarn.api.records.ApplicationSubmissionContext;
-import org.apache.hadoop.yarn.api.records.ResourceRequest;
 import org.apache.hadoop.yarn.conf.HAUtil;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
-import org.apache.hadoop.yarn.event.Dispatcher;
-import org.apache.hadoop.yarn.event.DrainDispatcher;
 import org.apache.hadoop.yarn.exceptions.YarnException;
 import org.apache.hadoop.yarn.server.resourcemanager.recovery.ZKRMStateStore;
 import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMApp;
 import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMAppImpl;
 import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMAppState;
 import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.RMAppAttempt;
-import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.RMAppAttemptImpl;
 import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.RMAppAttemptState;
 import org.apache.hadoop.yarn.server.resourcemanager.scheduler.YarnScheduler;
 import org.apache.hadoop.yarn.server.security.ApplicationACLsManager;
-import org.apache.hadoop.yarn.server.utils.BuilderUtils;
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
@@ -108,20 +103,9 @@ public abstract class RMHATestBase extends ClientBaseWithFixes{
   }
 
   protected void startRMs() throws IOException {
-    rm1 = new MockRM(confForRM1, null, false, false){
-      @Override
-      protected Dispatcher createDispatcher() {
-        return new DrainDispatcher();
-      }
-    };
-    rm2 = new MockRM(confForRM2, null, false, false){
-      @Override
-      protected Dispatcher createDispatcher() {
-        return new DrainDispatcher();
-      }
-    };
+    rm1 = new MockRM(confForRM1, null, false, false);
+    rm2 = new MockRM(confForRM2, null, false, false);
     startRMs(rm1, confForRM1, rm2, confForRM2);
-
   }
 
   protected void startRMsWithCustomizedRMAppManager() throws IOException {

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

@@ -47,7 +47,6 @@ import org.apache.hadoop.yarn.api.records.ReservationRequest;
 import org.apache.hadoop.yarn.api.records.ReservationRequestInterpreter;
 import org.apache.hadoop.yarn.api.records.ReservationRequests;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
-import org.apache.hadoop.yarn.event.DrainDispatcher;
 import org.apache.hadoop.yarn.exceptions.YarnException;
 import org.apache.hadoop.yarn.server.resourcemanager.reservation.Plan;
 import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.CapacityScheduler;
@@ -463,9 +462,7 @@ public class ReservationACLsTestBase extends ACLsTestBase {
       int attempts = 10;
       Collection<Plan> plans;
       do {
-        DrainDispatcher dispatcher =
-                (DrainDispatcher) resourceManager.getRMContext().getDispatcher();
-        dispatcher.await();
+        resourceManager.drainEvents();
         LOG.info("Waiting for node capacity to be added to plan");
         plans = resourceManager.getRMContext().getReservationSystem()
                 .getAllPlans().values();

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

@@ -40,8 +40,6 @@ import org.apache.hadoop.yarn.api.records.Priority;
 import org.apache.hadoop.yarn.api.records.Resource;
 import org.apache.hadoop.yarn.api.records.ResourceRequest;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
-import org.apache.hadoop.yarn.event.Dispatcher;
-import org.apache.hadoop.yarn.event.DrainDispatcher;
 import org.apache.hadoop.yarn.server.api.protocolrecords.NMContainerStatus;
 import org.apache.hadoop.yarn.server.api.protocolrecords.NodeHeartbeatResponse;
 import org.apache.hadoop.yarn.server.resourcemanager.recovery.MemoryRMStateStore;
@@ -161,13 +159,7 @@ public class TestApplicationCleanup {
 
     Logger rootLogger = LogManager.getRootLogger();
     rootLogger.setLevel(Level.DEBUG);
-    final DrainDispatcher dispatcher = new DrainDispatcher();
-    MockRM rm = new MockRM() {
-      @Override
-      protected Dispatcher createDispatcher() {
-        return dispatcher;
-      }
-    };
+    MockRM rm = new MockRM();
     rm.start();
 
     MockNM nm1 = rm.registerNode("127.0.0.1:1234", 5000);
@@ -185,8 +177,8 @@ public class TestApplicationCleanup {
     int request = 2;
     am.allocate("127.0.0.1" , 1000, request, 
         new ArrayList<ContainerId>());
-    dispatcher.await();
-    
+    rm.drainEvents();
+
     //kick the scheduler
     nm1.nodeHeartbeat(true);
     List<Container> conts = am.allocate(new ArrayList<ResourceRequest>(),
@@ -199,7 +191,7 @@ public class TestApplicationCleanup {
       Thread.sleep(100);
       conts = am.allocate(new ArrayList<ResourceRequest>(),
           new ArrayList<ContainerId>()).getAllocatedContainers();
-      dispatcher.await();
+      rm.drainEvents();
       contReceived += conts.size();
       nm1.nodeHeartbeat(true);
     }
@@ -209,7 +201,7 @@ public class TestApplicationCleanup {
     ArrayList<ContainerId> release = new ArrayList<ContainerId>();
     release.add(conts.get(0).getId());
     am.allocate(new ArrayList<ResourceRequest>(), release);
-    dispatcher.await();
+    rm.drainEvents();
 
     // Send one more heartbeat with a fake running container. This is to
     // simulate the situation that can happen if the NM reports that container
@@ -224,7 +216,7 @@ public class TestApplicationCleanup {
     containerStatuses.put(app.getApplicationId(), containerStatusList);
 
     NodeHeartbeatResponse resp = nm1.nodeHeartbeat(containerStatuses, true);
-    waitForContainerCleanup(dispatcher, nm1, resp);
+    waitForContainerCleanup(rm, nm1, resp);
 
     // Now to test the case when RM already gave cleanup, and NM suddenly
     // realizes that the container is running.
@@ -240,17 +232,17 @@ public class TestApplicationCleanup {
     resp = nm1.nodeHeartbeat(containerStatuses, true);
     // The cleanup list won't be instantaneous as it is given out by scheduler
     // and not RMNodeImpl.
-    waitForContainerCleanup(dispatcher, nm1, resp);
+    waitForContainerCleanup(rm, nm1, resp);
 
     rm.stop();
   }
 
-  protected void waitForContainerCleanup(DrainDispatcher dispatcher, MockNM nm,
+  protected void waitForContainerCleanup(MockRM rm, MockNM nm,
       NodeHeartbeatResponse resp) throws Exception {
     int waitCount = 0, cleanedConts = 0;
     List<ContainerId> contsToClean;
     do {
-      dispatcher.await();
+      rm.drainEvents();
       contsToClean = resp.getContainersToCleanup();
       cleanedConts += contsToClean.size();
       if (cleanedConts >= 1) {
@@ -400,13 +392,7 @@ public class TestApplicationCleanup {
     memStore.init(conf);
 
     // start RM
-    final DrainDispatcher dispatcher = new DrainDispatcher();
-    MockRM rm1 = new MockRM(conf, memStore) {
-      @Override
-      protected Dispatcher createDispatcher() {
-        return dispatcher;
-      }
-    };
+    MockRM rm1 = new MockRM(conf, memStore);
     rm1.start();
     MockNM nm1 =
         new MockNM("127.0.0.1:1234", 15120, rm1.getResourceTrackerService());
@@ -419,13 +405,7 @@ public class TestApplicationCleanup {
     rm1.waitForState(app0.getApplicationId(), RMAppState.RUNNING);
 
     // start new RM
-    final DrainDispatcher dispatcher2 = new DrainDispatcher();
-    MockRM rm2 = new MockRM(conf, memStore) {
-      @Override
-      protected Dispatcher createDispatcher() {
-        return dispatcher2;
-      }
-    };
+    MockRM rm2 = new MockRM(conf, memStore);
     rm2.start();
 
     // nm1 register to rm2, and do a heartbeat
@@ -437,7 +417,7 @@ public class TestApplicationCleanup {
     NodeHeartbeatResponse response = nm1.nodeHeartbeat(am0
         .getApplicationAttemptId(), 2, ContainerState.RUNNING);
 
-    waitForContainerCleanup(dispatcher2, nm1, response);
+    waitForContainerCleanup(rm2, nm1, response);
 
     rm1.stop();
     rm2.stop();

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

@@ -59,8 +59,6 @@ import org.apache.hadoop.yarn.api.records.ResourceRequest;
 import org.apache.hadoop.yarn.api.records.SerializedException;
 import org.apache.hadoop.yarn.api.records.Token;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
-import org.apache.hadoop.yarn.event.Dispatcher;
-import org.apache.hadoop.yarn.event.DrainDispatcher;
 import org.apache.hadoop.yarn.exceptions.ApplicationAttemptNotFoundException;
 import org.apache.hadoop.yarn.exceptions.ApplicationMasterNotRegisteredException;
 import org.apache.hadoop.yarn.exceptions.NMNotYetReadyException;
@@ -260,7 +258,6 @@ public class TestApplicationMasterLauncher {
     Configuration conf = new Configuration();
     conf.setInt(YarnConfiguration.RM_AM_MAX_ATTEMPTS, 1);
     conf.setInt(YarnConfiguration.CLIENT_NM_CONNECT_RETRY_INTERVAL_MS, 1);
-    final DrainDispatcher dispatcher = new DrainDispatcher();
     MockRM rm = new MockRMWithCustomAMLauncher(conf, null) {
       @Override
       protected ApplicationMasterLauncher createAMLauncher() {
@@ -284,12 +281,8 @@ public class TestApplicationMasterLauncher {
           }
         };
       }
-
-      @Override
-      protected Dispatcher createDispatcher() {
-        return dispatcher;
-      }
     };
+
     rm.start();
     MockNM nm1 = rm.registerNode("127.0.0.1:1234", 5120);
 
@@ -297,7 +290,7 @@ public class TestApplicationMasterLauncher {
 
     // kick the scheduling
     nm1.nodeHeartbeat(true);
-    dispatcher.await();
+    rm.drainEvents();
 
     MockRM.waitForState(app.getCurrentAppAttempt(),
       RMAppAttemptState.LAUNCHED, 500);

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

@@ -42,8 +42,6 @@ import org.apache.hadoop.yarn.api.records.Priority;
 import org.apache.hadoop.yarn.api.records.ResourceRequest;
 import org.apache.hadoop.yarn.api.records.UpdateContainerRequest;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
-import org.apache.hadoop.yarn.event.Dispatcher;
-import org.apache.hadoop.yarn.event.DrainDispatcher;
 import org.apache.hadoop.yarn.exceptions.ApplicationMasterNotRegisteredException;
 import org.apache.hadoop.yarn.exceptions.InvalidContainerReleaseException;
 import org.apache.hadoop.yarn.proto.YarnServiceProtos.SchedulerResourceTypes;
@@ -327,10 +325,8 @@ public class TestApplicationMasterService {
 
   @Test(timeout=1200000)
   public void  testAllocateAfterUnregister() throws Exception {
-    MyResourceManager rm = new MyResourceManager(conf);
+    MockRM rm = new MockRM(conf);
     rm.start();
-    DrainDispatcher rmDispatcher = (DrainDispatcher) rm.getRMContext()
-            .getDispatcher();
     // Register node1
     MockNM nm1 = rm.registerNode("127.0.0.1:1234", 6 * GB);
 
@@ -351,7 +347,7 @@ public class TestApplicationMasterService {
     AllocateResponse alloc1Response = am1.schedule();
 
     nm1.nodeHeartbeat(true);
-    rmDispatcher.await();
+    rm.drainEvents();
     alloc1Response = am1.schedule();
     Assert.assertEquals(0, alloc1Response.getAllocatedContainers().size());
   }
@@ -474,17 +470,6 @@ public class TestApplicationMasterService {
     rm.stop();
   }
 
-  private static class MyResourceManager extends MockRM {
-
-    public MyResourceManager(YarnConfiguration conf) {
-      super(conf);
-    }
-    @Override
-    protected Dispatcher createDispatcher() {
-      return new DrainDispatcher();
-    }
-  }
-  
   private void sentRMContainerLaunched(MockRM rm, ContainerId containerId) {
     CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();
     RMContainer rmContainer = cs.getRMContainer(containerId);

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

@@ -33,8 +33,6 @@ import org.apache.hadoop.yarn.api.records.Priority;
 import org.apache.hadoop.yarn.api.records.Resource;
 import org.apache.hadoop.yarn.api.records.ResourceRequest;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
-import org.apache.hadoop.yarn.event.Dispatcher;
-import org.apache.hadoop.yarn.event.DrainDispatcher;
 import org.apache.hadoop.yarn.server.resourcemanager.applicationsmanager.TestAMRestart;
 import org.apache.hadoop.yarn.server.resourcemanager.recovery.MemoryRMStateStore;
 import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMApp;
@@ -65,8 +63,7 @@ public class TestNodeBlacklistingOnAMFailures {
     conf.setBoolean(YarnConfiguration.AM_SCHEDULING_NODE_BLACKLISTING_ENABLED,
         true);
 
-    DrainDispatcher dispatcher = new DrainDispatcher();
-    MockRM rm = startRM(conf, dispatcher);
+    MockRM rm = startRM(conf);
     CapacityScheduler scheduler = (CapacityScheduler) rm.getResourceScheduler();
 
     // Register 5 nodes, so that we can blacklist atleast one if AM container
@@ -122,7 +119,7 @@ public class TestNodeBlacklistingOnAMFailures {
     // Try the current node a few times
     for (int i = 0; i <= 2; i++) {
       currentNode.nodeHeartbeat(true);
-      dispatcher.await();
+      rm.drainEvents();
 
       Assert.assertEquals(
           "AppAttemptState should still be SCHEDULED if currentNode is "
@@ -132,7 +129,7 @@ public class TestNodeBlacklistingOnAMFailures {
 
     // Now try the other node
     otherNode.nodeHeartbeat(true);
-    dispatcher.await();
+    rm.drainEvents();
 
     // Now the AM container should be allocated
     MockRM.waitForState(attempt, RMAppAttemptState.ALLOCATED, 20000);
@@ -169,8 +166,7 @@ public class TestNodeBlacklistingOnAMFailures {
     conf.setBoolean(YarnConfiguration.AM_SCHEDULING_NODE_BLACKLISTING_ENABLED,
         true);
 
-    DrainDispatcher dispatcher = new DrainDispatcher();
-    MockRM rm = startRM(conf, dispatcher);
+    MockRM rm = startRM(conf);
     CapacityScheduler scheduler = (CapacityScheduler) rm.getResourceScheduler();
 
     // Register 5 nodes, so that we can blacklist atleast one if AM container
@@ -227,7 +223,7 @@ public class TestNodeBlacklistingOnAMFailures {
     System.out.println("New AppAttempt launched " + attempt.getAppAttemptId());
 
     nm2.nodeHeartbeat(true);
-    dispatcher.await();
+    rm.drainEvents();
 
     // Now the AM container should be allocated
     MockRM.waitForState(attempt, RMAppAttemptState.ALLOCATED, 20000);
@@ -257,8 +253,7 @@ public class TestNodeBlacklistingOnAMFailures {
     conf.setBoolean(YarnConfiguration.AM_SCHEDULING_NODE_BLACKLISTING_ENABLED,
         true);
 
-    DrainDispatcher dispatcher = new DrainDispatcher();
-    MockRM rm = startRM(conf, dispatcher);
+    MockRM rm = startRM(conf);
     CapacityScheduler scheduler = (CapacityScheduler) rm.getResourceScheduler();
 
     // Register 5 nodes, so that we can blacklist atleast one if AM container
@@ -319,7 +314,7 @@ public class TestNodeBlacklistingOnAMFailures {
     nm3.nodeHeartbeat(true);
     nm4.nodeHeartbeat(true);
     nm5.nodeHeartbeat(true);
-    dispatcher.await();
+    rm.drainEvents();
 
     // Now the AM container should be allocated
     MockRM.waitForState(attempt, RMAppAttemptState.ALLOCATED, 20000);
@@ -352,8 +347,7 @@ public class TestNodeBlacklistingOnAMFailures {
         1.5f);
     conf.setInt(YarnConfiguration.RM_AM_MAX_ATTEMPTS, 100);
 
-    DrainDispatcher dispatcher = new DrainDispatcher();
-    MockRM rm = startRM(conf, dispatcher);
+    MockRM rm = startRM(conf);
 
     MockNM node =
         new MockNM("127.0.0.1:1234", 8000, rm.getResourceTrackerService());
@@ -367,7 +361,7 @@ public class TestNodeBlacklistingOnAMFailures {
     // Now the AM container should be allocated
     RMAppAttempt attempt = MockRM.waitForAttemptScheduled(app, rm);
     node.nodeHeartbeat(true);
-    dispatcher.await();
+    rm.drainEvents();
     MockRM.waitForState(attempt, RMAppAttemptState.ALLOCATED, 20000);
     rm.sendAMLaunched(attempt.getAppAttemptId());
     rm.waitForState(attempt.getAppAttemptId(), RMAppAttemptState.LAUNCHED);
@@ -394,7 +388,7 @@ public class TestNodeBlacklistingOnAMFailures {
           .println("New AppAttempt launched " + attempt.getAppAttemptId());
 
       node.nodeHeartbeat(true);
-      dispatcher.await();
+      rm.drainEvents();
 
       MockRM.waitForState(attempt, RMAppAttemptState.ALLOCATED, 20000);
       rm.sendAMLaunched(attempt.getAppAttemptId());
@@ -418,20 +412,13 @@ public class TestNodeBlacklistingOnAMFailures {
     rm.waitForState(amAttemptID.getApplicationId(), RMAppState.ACCEPTED);
   }
 
-  private MockRM startRM(YarnConfiguration conf,
-      final DrainDispatcher dispatcher) {
-
+  private MockRM startRM(YarnConfiguration conf) {
     MemoryRMStateStore memStore = new MemoryRMStateStore();
     memStore.init(conf);
 
-    MockRM rm1 = new MockRM(conf, memStore) {
-      @Override
-      protected Dispatcher createDispatcher() {
-        return dispatcher;
-      }
-    };
+    MockRM rm = new MockRM(conf, memStore);
 
-    rm1.start();
-    return rm1;
+    rm.start();
+    return rm;
   }
 }

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

@@ -29,7 +29,6 @@ import org.apache.hadoop.yarn.api.records.QueueInfo;
 import org.apache.hadoop.yarn.api.records.ReservationDefinition;
 import org.apache.hadoop.yarn.api.records.ReservationId;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
-import org.apache.hadoop.yarn.event.DrainDispatcher;
 import org.apache.hadoop.yarn.proto.YarnProtos.ReservationAllocationStateProto;
 import org.apache.hadoop.yarn.server.resourcemanager.recovery.RMStateStore.RMState;
 import org.apache.hadoop.yarn.server.resourcemanager.reservation.Plan;
@@ -186,9 +185,7 @@ public class TestReservationSystemWithRMHA extends RMHATestBase {
       rm.registerNode("127.0.0.1:1", memory, vCores);
       int attempts = 10;
       do {
-        DrainDispatcher dispatcher =
-            (DrainDispatcher) rm1.getRMContext().getDispatcher();
-        dispatcher.await();
+        rm1.drainEvents();
         rm.getRMContext().getReservationSystem()
             .synchronizePlan(ReservationSystemTestUtil.reservationQ, false);
         if (rm.getRMContext().getReservationSystem()

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/applicationsmanager/TestAMRMRPCNodeUpdates.java

@@ -31,8 +31,6 @@ import org.apache.hadoop.yarn.api.protocolrecords.AllocateResponse;
 import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
 import org.apache.hadoop.yarn.api.records.NodeReport;
 import org.apache.hadoop.yarn.api.records.NodeState;
-import org.apache.hadoop.yarn.event.Dispatcher;
-import org.apache.hadoop.yarn.event.DrainDispatcher;
 import org.apache.hadoop.yarn.security.AMRMTokenIdentifier;
 import org.apache.hadoop.yarn.server.resourcemanager.ApplicationMasterService;
 import org.apache.hadoop.yarn.server.resourcemanager.MockAM;
@@ -47,12 +45,10 @@ import org.junit.Test;
 
 public class TestAMRMRPCNodeUpdates {
   private MockRM rm;
-  ApplicationMasterService amService = null;
-  DrainDispatcher dispatcher = null;
+  private ApplicationMasterService amService;
 
   @Before
   public void setUp() {
-    dispatcher = new DrainDispatcher();
     this.rm = new MockRM() {
       @Override
       public void init(Configuration conf) {
@@ -61,12 +57,8 @@ public class TestAMRMRPCNodeUpdates {
           "1.0");
         super.init(conf);
       }
-
-      @Override
-      protected Dispatcher createDispatcher() {
-        return dispatcher;
-      }
     };
+
     rm.start();
     amService = rm.getApplicationMasterService();
   }
@@ -80,14 +72,14 @@ public class TestAMRMRPCNodeUpdates {
   
   private void syncNodeHeartbeat(MockNM nm, boolean health) throws Exception {
     nm.nodeHeartbeat(health);
-    dispatcher.await();
+    rm.drainEvents();
   }
   
   private void syncNodeLost(MockNM nm) throws Exception {
     rm.sendNodeStarted(nm);
     rm.waitForState(nm.getNodeId(), NodeState.RUNNING);
     rm.sendNodeLost(nm);
-    dispatcher.await();
+    rm.drainEvents();
   }
 
   private AllocateResponse allocate(final ApplicationAttemptId attemptId,
@@ -113,7 +105,7 @@ public class TestAMRMRPCNodeUpdates {
     MockNM nm2 = rm.registerNode("127.0.0.2:1234", 10000);
     MockNM nm3 = rm.registerNode("127.0.0.3:1234", 10000);
     MockNM nm4 = rm.registerNode("127.0.0.4:1234", 10000);
-    dispatcher.await();
+    rm.drainEvents();
 
     RMApp app1 = rm.submitApp(2000);
 

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/resourcetracker/TestNMReconnect.java

@@ -28,7 +28,6 @@ import org.apache.hadoop.yarn.conf.ConfigurationProvider;
 import org.apache.hadoop.yarn.conf.ConfigurationProviderFactory;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
 import org.apache.hadoop.yarn.event.Dispatcher;
-import org.apache.hadoop.yarn.event.DrainDispatcher;
 import org.apache.hadoop.yarn.event.EventHandler;
 import org.apache.hadoop.yarn.event.InlineDispatcher;
 import org.apache.hadoop.yarn.factories.RecordFactory;
@@ -228,21 +227,16 @@ public class TestNMReconnect extends ParameterizedSchedulerTestBase {
     // The node(127.0.0.1:1234) reconnected with RM. When it registered with
     // RM, RM set its lastNodeHeartbeatResponse's id to 0 asynchronously. But
     // the node's heartbeat come before RM succeeded setting the id to 0.
-    final DrainDispatcher dispatcher = new DrainDispatcher();
-    MockRM rm = new MockRM(){
-      @Override
-      protected Dispatcher createDispatcher() {
-        return dispatcher;
-      }
-    };
+    MockRM rm = new MockRM();
     rm.start();
+
     MockNM nm1 =
         new MockNM("127.0.0.1:1234", 15120, rm.getResourceTrackerService());
     nm1.registerNode();
     int i = 0;
     while(i < 3) {
       nm1.nodeHeartbeat(true);
-      dispatcher.await();
+      rm.drainEvents();
       i++;
     }
 
@@ -251,7 +245,7 @@ public class TestNMReconnect extends ParameterizedSchedulerTestBase {
     nm2.registerNode();
     RMNode rmNode = rm.getRMContext().getRMNodes().get(nm2.getNodeId());
     nm2.nodeHeartbeat(true);
-    dispatcher.await();
+    rm.drainEvents();
     Assert.assertEquals("Node is Not in Running state.", NodeState.RUNNING,
         rmNode.getState());
     rm.stop();

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/rmcontainer/TestRMContainerImpl.java

@@ -46,7 +46,6 @@ import org.apache.hadoop.yarn.api.records.Priority;
 import org.apache.hadoop.yarn.api.records.Resource;
 import org.apache.hadoop.yarn.api.records.ResourceRequest;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
-import org.apache.hadoop.yarn.event.Dispatcher;
 import org.apache.hadoop.yarn.event.DrainDispatcher;
 import org.apache.hadoop.yarn.event.EventHandler;
 import org.apache.hadoop.yarn.server.resourcemanager.MockAM;

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestApplicationPriority.java

@@ -36,8 +36,6 @@ import org.apache.hadoop.yarn.api.records.Container;
 import org.apache.hadoop.yarn.api.records.ContainerId;
 import org.apache.hadoop.yarn.api.records.Priority;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
-import org.apache.hadoop.yarn.event.Dispatcher;
-import org.apache.hadoop.yarn.event.DrainDispatcher;
 import org.apache.hadoop.yarn.exceptions.YarnException;
 import org.apache.hadoop.yarn.server.resourcemanager.MockAM;
 import org.apache.hadoop.yarn.server.resourcemanager.MockNM;
@@ -612,24 +610,17 @@ public class TestApplicationPriority {
     conf.setInt(YarnConfiguration.RM_AM_MAX_ATTEMPTS,
         YarnConfiguration.DEFAULT_RM_AM_MAX_ATTEMPTS);
     conf.setInt(YarnConfiguration.MAX_CLUSTER_LEVEL_APPLICATION_PRIORITY, 10);
-    final DrainDispatcher dispatcher = new DrainDispatcher();
 
     MemoryRMStateStore memStore = new MemoryRMStateStore();
     memStore.init(conf);
 
-    MockRM rm1 = new MockRM(conf, memStore) {
-      @Override
-      protected Dispatcher createDispatcher() {
-        return dispatcher;
-      }
-    };
+    MockRM rm1 = new MockRM(conf, memStore);
     rm1.start();
 
     MockNM nm1 =
         new MockNM("127.0.0.1:1234", 16384, rm1.getResourceTrackerService());
     nm1.registerNode();
-
-    dispatcher.await();
+    rm1.drainEvents();
 
     ResourceScheduler scheduler = rm1.getRMContext().getScheduler();
     LeafQueue defaultQueue =
@@ -648,7 +639,7 @@ public class TestApplicationPriority {
     MockAM am2 = MockRM.launchAM(app2, rm1, nm1);
     am2.registerAppAttempt();
 
-    dispatcher.await();
+    rm1.drainEvents();
     Assert.assertEquals(2, defaultQueue.getNumActiveApplications());
     Assert.assertEquals(0, defaultQueue.getNumPendingApplications());
 
@@ -657,7 +648,7 @@ public class TestApplicationPriority {
     Priority appPriority3 = Priority.newInstance(7);
     RMApp app3 = rm1.submitApp(memory, appPriority3);
 
-    dispatcher.await();
+    rm1.drainEvents();
     Assert.assertEquals(2, defaultQueue.getNumActiveApplications());
     Assert.assertEquals(1, defaultQueue.getNumPendingApplications());
 
@@ -676,14 +667,8 @@ public class TestApplicationPriority {
     Assert.assertEquals(app3.getCurrentAppAttempt().getAppAttemptId(),
         fcApp3.getApplicationAttemptId());
 
-    final DrainDispatcher dispatcher1 = new DrainDispatcher();
     // create new RM to represent restart and recover state
-    MockRM rm2 = new MockRM(conf, memStore) {
-      @Override
-      protected Dispatcher createDispatcher() {
-        return dispatcher1;
-      }
-    };
+    MockRM rm2 = new MockRM(conf, memStore);
 
     // start new RM
     rm2.start();
@@ -693,7 +678,7 @@ public class TestApplicationPriority {
     // Verify RM Apps after this restart
     Assert.assertEquals(3, rm2.getRMContext().getRMApps().size());
 
-    dispatcher1.await();
+    rm2.drainEvents();
     scheduler = rm2.getRMContext().getScheduler();
     defaultQueue =
         (LeafQueue) ((CapacityScheduler) scheduler).getQueue("default");
@@ -714,7 +699,7 @@ public class TestApplicationPriority {
 
     // NM resync to new RM
     nm1.registerNode();
-    dispatcher1.await();
+    rm2.drainEvents();
 
     // wait for activating applications
     count = 50;

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/security/TestClientToAMTokens.java

@@ -50,8 +50,6 @@ import org.apache.hadoop.yarn.api.protocolrecords.StartContainersResponse;
 import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
 import org.apache.hadoop.yarn.api.records.ApplicationReport;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
-import org.apache.hadoop.yarn.event.Dispatcher;
-import org.apache.hadoop.yarn.event.DrainDispatcher;
 import org.apache.hadoop.yarn.exceptions.YarnRuntimeException;
 import org.apache.hadoop.yarn.security.client.ClientToAMTokenIdentifier;
 import org.apache.hadoop.yarn.security.client.ClientToAMTokenSecretManager;
@@ -199,7 +197,6 @@ public class TestClientToAMTokens extends ParameterizedSchedulerTestBase {
     StartContainersResponse mockResponse = mock(StartContainersResponse.class);
     when(containerManager.startContainers((StartContainersRequest) any()))
       .thenReturn(mockResponse);
-    final DrainDispatcher dispatcher = new DrainDispatcher();
 
     MockRM rm = new MockRMWithCustomAMLauncher(conf, containerManager) {
       protected ClientRMService createClientRMService() {
@@ -208,11 +205,6 @@ public class TestClientToAMTokens extends ParameterizedSchedulerTestBase {
           getRMContext().getRMDelegationTokenSecretManager());
       };
 
-      @Override
-      protected Dispatcher createDispatcher() {
-        return dispatcher;
-      }
-
       @Override
       protected void doSecureLogin() throws IOException {
       }
@@ -225,11 +217,10 @@ public class TestClientToAMTokens extends ParameterizedSchedulerTestBase {
     // Set up a node.
     MockNM nm1 = rm.registerNode("localhost:1234", 3072);
     nm1.nodeHeartbeat(true);
-    dispatcher.await();
+    rm.drainEvents();
     
-
     nm1.nodeHeartbeat(true);
-    dispatcher.await();
+    rm.drainEvents();
 
     ApplicationAttemptId appAttempt = app.getCurrentAppAttempt().getAppAttemptId();
     final MockAM mockAM =
@@ -436,7 +427,6 @@ public class TestClientToAMTokens extends ParameterizedSchedulerTestBase {
     StartContainersResponse mockResponse = mock(StartContainersResponse.class);
     when(containerManager.startContainers((StartContainersRequest) any()))
       .thenReturn(mockResponse);
-    final DrainDispatcher dispatcher = new DrainDispatcher();
 
     MockRM rm = new MockRMWithCustomAMLauncher(conf, containerManager) {
       protected ClientRMService createClientRMService() {
@@ -445,11 +435,6 @@ public class TestClientToAMTokens extends ParameterizedSchedulerTestBase {
           getRMContext().getRMDelegationTokenSecretManager());
       };
 
-      @Override
-      protected Dispatcher createDispatcher() {
-        return dispatcher;
-      }
-
       @Override
       protected void doSecureLogin() throws IOException {
       }
@@ -462,10 +447,10 @@ public class TestClientToAMTokens extends ParameterizedSchedulerTestBase {
     // Set up a node.
     MockNM nm1 = rm.registerNode("localhost:1234", 3072);
     nm1.nodeHeartbeat(true);
-    dispatcher.await();
+    rm.drainEvents();
 
     nm1.nodeHeartbeat(true);
-    dispatcher.await();
+    rm.drainEvents();
 
     ApplicationAttemptId appAttempt = app.getCurrentAppAttempt().getAppAttemptId();
     final MockAM mockAM =