Java DrainDispatcher类代码示例

import org.apache.hadoop.yarn.event.DrainDispatcher; //导入依赖的package包/类
private ResourceLocalizationService createSpyService(
    DrainDispatcher dispatcher, LocalDirsHandlerService dirsHandler,
    NMStateStoreService stateStore) {
  ContainerExecutor exec = mock(ContainerExecutor.class);
  LocalizerTracker mockLocalizerTracker = mock(LocalizerTracker.class);
  DeletionService delService = mock(DeletionService.class);
  NMContext nmContext =
      new NMContext(new NMContainerTokenSecretManager(conf),
        new NMTokenSecretManagerInNM(), null,
        new ApplicationACLsManager(conf), stateStore);
  ResourceLocalizationService rawService =
    new ResourceLocalizationService(dispatcher, exec, delService,
                                    dirsHandler, nmContext);
  ResourceLocalizationService spyService = spy(rawService);
  doReturn(mockServer).when(spyService).createServer();
  doReturn(mockLocalizerTracker).when(spyService).createLocalizerTracker(
      isA(Configuration.class));
  doReturn(lfs).when(spyService)
      .getLocalFileContext(isA(Configuration.class));
  return spyService;
}
 

import org.apache.hadoop.yarn.event.DrainDispatcher; //导入依赖的package包/类
protected void waitForContainerCleanup(DrainDispatcher dispatcher, MockNM nm,
    NodeHeartbeatResponse resp) throws Exception {
  int waitCount = 0, cleanedConts = 0;
  List<ContainerId> contsToClean;
  do {
    dispatcher.await();
    contsToClean = resp.getContainersToCleanup();
    cleanedConts += contsToClean.size();
    if (cleanedConts >= 1) {
      break;
    }
    Thread.sleep(100);
    resp = nm.nodeHeartbeat(true);
  } while(waitCount++ < 200);

  if (contsToClean.isEmpty()) {
    LOG.error("Failed to get any containers to cleanup");
  } else {
    LOG.info("Got cleanup for " + contsToClean.get(0));
  }
  Assert.assertEquals(1, cleanedConts);
}
 

import org.apache.hadoop.yarn.event.DrainDispatcher; //导入依赖的package包/类
protected void startRMs() throws IOException {
  rm1 = new MockRM(confForRM1, null, false){
    @Override
    protected Dispatcher createDispatcher() {
      return new DrainDispatcher();
    }
  };
  rm2 = new MockRM(confForRM2, null, false){
    @Override
    protected Dispatcher createDispatcher() {
      return new DrainDispatcher();
    }
  };
  startRMs(rm1, confForRM1, rm2, confForRM2);

}
 

import org.apache.hadoop.yarn.event.DrainDispatcher; //导入依赖的package包/类
private ResourceLocalizationService createSpyService(
    DrainDispatcher dispatcher, LocalDirsHandlerService dirsHandler,
    NMStateStoreService stateStore) {
  ContainerExecutor exec = mock(ContainerExecutor.class);
  LocalizerTracker mockLocalizerTracker = mock(LocalizerTracker.class);
  DeletionService delService = mock(DeletionService.class);
  NMContext nmContext =
      new NMContext(new NMContainerTokenSecretManager(conf),
        new NMTokenSecretManagerInNM(), null,
        new ApplicationACLsManager(conf), stateStore,null);
  ResourceLocalizationService rawService =
    new ResourceLocalizationService(dispatcher, exec, delService,
                                    dirsHandler, nmContext);
  ResourceLocalizationService spyService = spy(rawService);
  doReturn(mockServer).when(spyService).createServer();
  doReturn(mockLocalizerTracker).when(spyService).createLocalizerTracker(
      isA(Configuration.class));
  doReturn(lfs).when(spyService)
      .getLocalFileContext(isA(Configuration.class));
  return spyService;
}
 

import org.apache.hadoop.yarn.event.DrainDispatcher; //导入依赖的package包/类
@Test
public void testVerifyAndCreateRemoteDirNonExistence()
    throws Exception {
  this.conf.set(YarnConfiguration.NM_LOG_DIRS, localLogDir.getAbsolutePath());
  File aNewFile = new File(String.valueOf("tmp"+System.currentTimeMillis()));
  this.conf.set(YarnConfiguration.NM_REMOTE_APP_LOG_DIR, 
      aNewFile.getAbsolutePath());
  
  DrainDispatcher dispatcher = createDispatcher();
  LogAggregationService logAggregationService = spy(
      new LogAggregationService(dispatcher, this.context, this.delSrvc,
                                super.dirsHandler));
  logAggregationService.init(this.conf);
  boolean existsBefore = aNewFile.exists();
  assertTrue("The new file already exists!", !existsBefore);

  logAggregationService.verifyAndCreateRemoteLogDir(this.conf);
  
  boolean existsAfter = aNewFile.exists();
  assertTrue("The new aggregate file is not successfully created", existsAfter);
  aNewFile.delete(); //housekeeping
}
 

import org.apache.hadoop.yarn.event.DrainDispatcher; //导入依赖的package包/类
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();
    }
  };
  startRMs(rm1, confForRM1, rm2, confForRM2);

}
 

import org.apache.hadoop.yarn.event.DrainDispatcher; //导入依赖的package包/类
@Before
public void setUp() {
  dispatcher = new DrainDispatcher();
  this.rm = new MockRM() {
    @Override
    protected EventHandler<SchedulerEvent> createSchedulerEventDispatcher() {
      return new SchedulerEventDispatcher(this.scheduler) {
        @Override
        public void handle(SchedulerEvent event) {
          scheduler.handle(event);
        }
      };
    }

    @Override
    protected Dispatcher createDispatcher() {
      return dispatcher;
    }
  };
  rm.start();
  amService = rm.getApplicationMasterService();
}
 

import org.apache.hadoop.yarn.event.DrainDispatcher; //导入依赖的package包/类
private
    List<TaskAttemptContainerAssignedEvent> getContainerOnHost(JobId jobId,
        int taskAttemptId, int memory, String[] hosts, MockNM mockNM,
        DrainDispatcher dispatcher, MyContainerAllocator allocator)
        throws Exception {
  ContainerRequestEvent reqEvent =
      createReq(jobId, taskAttemptId, memory, hosts);
  allocator.sendRequest(reqEvent);

  // Send the request to the RM
  List<TaskAttemptContainerAssignedEvent> assigned = allocator.schedule();
  dispatcher.await();
  Assert.assertEquals("No of assignments must be 0", 0, assigned.size());

  // Heartbeat from the required nodeManager
  mockNM.nodeHeartbeat(true);
  dispatcher.await();

  assigned = allocator.schedule();
  dispatcher.await();
  return assigned;
}
 

import org.apache.hadoop.yarn.event.DrainDispatcher; //导入依赖的package包/类
@Before
public void setUp() {
  dispatcher = new DrainDispatcher();
  this.rm = new MockRM() {
    @Override
    public void init(Configuration conf) {
      conf.set(
        CapacitySchedulerConfiguration.MAXIMUM_APPLICATION_MASTERS_RESOURCE_PERCENT,
        "1.0");
      super.init(conf);
    }
    @Override
    protected EventHandler<SchedulerEvent> createSchedulerEventDispatcher() {
      return new SchedulerEventDispatcher(this.scheduler) {
        @Override
        public void handle(SchedulerEvent event) {
          scheduler.handle(event);
        }
      };
    }

    @Override
    protected Dispatcher createDispatcher() {
      return dispatcher;
    }
  };
  rm.start();
  amService = rm.getApplicationMasterService();
}
 

import org.apache.hadoop.yarn.event.DrainDispatcher; //导入依赖的package包/类
@Test (timeout=10000)
public void testFailAbortDoesntHang() throws IOException {
  Configuration conf = new Configuration();
  conf.set(MRJobConfig.MR_AM_STAGING_DIR, stagingDir);
  conf.set(MRJobConfig.MR_AM_COMMITTER_CANCEL_TIMEOUT_MS, "1000");
  
  DrainDispatcher dispatcher = new DrainDispatcher();
  dispatcher.init(conf);
  dispatcher.start();
  OutputCommitter committer = Mockito.mock(OutputCommitter.class);
  CommitterEventHandler commitHandler =
      createCommitterEventHandler(dispatcher, committer);
  commitHandler.init(conf);
  commitHandler.start();
  //Job has only 1 mapper task. No reducers
  conf.setInt(MRJobConfig.NUM_REDUCES, 0);
  conf.setInt(MRJobConfig.MAP_MAX_ATTEMPTS, 1);
  JobImpl job = createRunningStubbedJob(conf, dispatcher, 1, null);

  //Fail / finish all the tasks. This should land the JobImpl directly in the
  //FAIL_ABORT state
  for(Task t: job.tasks.values()) {
    TaskImpl task = (TaskImpl) t;
    task.handle(new TaskEvent(task.getID(), TaskEventType.T_SCHEDULE));
    for(TaskAttempt ta: task.getAttempts().values()) {
      task.handle(new TaskTAttemptEvent(ta.getID(),
        TaskEventType.T_ATTEMPT_FAILED));
    }
  }

  dispatcher.await();
  //Verify abortJob is called once and the job failed
  Mockito.verify(committer, Mockito.timeout(2000).times(1))
    .abortJob((JobContext) Mockito.any(), (State) Mockito.any());
  assertJobState(job, JobStateInternal.FAILED);

  dispatcher.stop();
}
 

import org.apache.hadoop.yarn.event.DrainDispatcher; //导入依赖的package包/类
private void finishNextNTasks(DrainDispatcher rmDispatcher, MockNM node,
    MRApp mrApp, Iterator<Task> it, int nextN) throws Exception {
  Task task;
  for (int i=0; i<nextN; i++) {
    task = it.next();
    finishTask(rmDispatcher, node, mrApp, task);
  }
}
 

import org.apache.hadoop.yarn.event.DrainDispatcher; //导入依赖的package包/类
private void finishTask(DrainDispatcher rmDispatcher, MockNM node,
    MRApp mrApp, Task task) throws Exception {
  TaskAttempt attempt = task.getAttempts().values().iterator().next();
  List<ContainerStatus> contStatus = new ArrayList<ContainerStatus>(1);
  contStatus.add(ContainerStatus.newInstance(attempt.getAssignedContainerID(),
      ContainerState.COMPLETE, "", 0));
  Map<ApplicationId,List<ContainerStatus>> statusUpdate =
      new HashMap<ApplicationId,List<ContainerStatus>>(1);
  statusUpdate.put(mrApp.getAppID(), contStatus);
  node.nodeHeartbeat(statusUpdate, true);
  rmDispatcher.await();
  mrApp.getContext().getEventHandler().handle(
        new TaskAttemptEvent(attempt.getID(), TaskAttemptEventType.TA_DONE));
  mrApp.waitForState(task, TaskState.SUCCEEDED);
}
 

import org.apache.hadoop.yarn.event.DrainDispatcher; //导入依赖的package包/类
private
    List<TaskAttemptContainerAssignedEvent> getContainerOnHost(JobId jobId,
        int taskAttemptId, int memory, String[] hosts, MockNM mockNM,
        DrainDispatcher dispatcher, MyContainerAllocator allocator,
        int expectedAdditions1, int expectedRemovals1,
        int expectedAdditions2, int expectedRemovals2, MyResourceManager rm)
        throws Exception {
  ContainerRequestEvent reqEvent =
      createReq(jobId, taskAttemptId, memory, hosts);
  allocator.sendRequest(reqEvent);

  // Send the request to the RM
  List<TaskAttemptContainerAssignedEvent> assigned = allocator.schedule();
  dispatcher.await();
  assertBlacklistAdditionsAndRemovals(
      expectedAdditions1, expectedRemovals1, rm);
  Assert.assertEquals("No of assignments must be 0", 0, assigned.size());

  // Heartbeat from the required nodeManager
  mockNM.nodeHeartbeat(true);
  dispatcher.await();

  assigned = allocator.schedule();
  dispatcher.await();
  assertBlacklistAdditionsAndRemovals(
      expectedAdditions2, expectedRemovals2, rm);
  return assigned;
}
 

import org.apache.hadoop.yarn.event.DrainDispatcher; //导入依赖的package包/类
private void addNodeCapacityToPlan(MockRM rm, int memory, int vCores) {
  try {
    rm.registerNode("127.0.0.1:1", memory, vCores);
    int attempts = 10;
    do {
      DrainDispatcher dispatcher =
          (DrainDispatcher) rm1.getRMContext().getDispatcher();
      dispatcher.await();
      rm.getRMContext().getReservationSystem()
          .synchronizePlan(ReservationSystemTestUtil.reservationQ, false);
      if (rm.getRMContext().getReservationSystem()
          .getPlan(ReservationSystemTestUtil.reservationQ).getTotalCapacity()
          .getMemory() > 0) {
        break;
      }
      LOG.info("Waiting for node capacity to be added to plan");
      Thread.sleep(100);
    } while (attempts-- > 0);
    if (attempts <= 0) {
      Assert.fail("Exhausted attempts in checking if node capacity was "
          + "added to the plan");
    }

  } catch (Exception e) {
    Assert.fail(e.getMessage());
  }
}
 

import org.apache.hadoop.yarn.event.DrainDispatcher; //导入依赖的package包/类
@Test(timeout=1200000)
public void  testAllocateAfterUnregister() throws Exception {
  MyResourceManager rm = new MyResourceManager(conf);
  rm.start();
  DrainDispatcher rmDispatcher = (DrainDispatcher) rm.getRMContext()
          .getDispatcher();
  // Register node1
  MockNM nm1 = rm.registerNode("127.0.0.1:1234", 6 * GB);

  // Submit an application
  RMApp app1 = rm.submitApp(2048);

  nm1.nodeHeartbeat(true);
  RMAppAttempt attempt1 = app1.getCurrentAppAttempt();
  MockAM am1 = rm.sendAMLaunched(attempt1.getAppAttemptId());
  am1.registerAppAttempt();
  // unregister app attempt
  FinishApplicationMasterRequest req =
      FinishApplicationMasterRequest.newInstance(
         FinalApplicationStatus.KILLED, "", "");
  am1.unregisterAppAttempt(req, false);
  // request container after unregister
  am1.addRequests(new String[] { "127.0.0.1" }, GB, 1, 1);
  AllocateResponse alloc1Response = am1.schedule();

  nm1.nodeHeartbeat(true);
  rmDispatcher.await();
  alloc1Response = am1.schedule();
  Assert.assertEquals(0, alloc1Response.getAllocatedContainers().size());
}