import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
@Test
public void testAddAndRemoveAppFromCapacityScheduler() throws Exception {
  CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
  setupQueueConfiguration(conf);
  conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class,
    ResourceScheduler.class);
  MockRM rm = new MockRM(conf);
  @SuppressWarnings("unchecked")
  AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode> cs =
      (AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>) rm
        .getResourceScheduler();

  SchedulerApplication<SchedulerApplicationAttempt> app =
      TestSchedulerUtils.verifyAppAddedAndRemovedFromScheduler(
        cs.getSchedulerApplications(), cs, "a1");
  Assert.assertEquals("a1", app.getQueue().getQueueName());
}
 

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
public static void waitForNumContainersToRecover(int num, MockRM rm,
    ApplicationAttemptId attemptId) throws Exception {
  AbstractYarnScheduler scheduler =
      (AbstractYarnScheduler) rm.getResourceScheduler();
  SchedulerApplicationAttempt attempt =
      scheduler.getApplicationAttempt(attemptId);
  while (attempt == null) {
    System.out.println("Wait for scheduler attempt " + attemptId
        + " to be created");
    Thread.sleep(200);
    attempt = scheduler.getApplicationAttempt(attemptId);
  }
  while (attempt.getLiveContainers().size() < num) {
    System.out.println("Wait for " + num
        + " containers to recover. currently: "
        + attempt.getLiveContainers().size());
    Thread.sleep(200);
  }
}
 

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
@SuppressWarnings("rawtypes")
private static void waitForSchedulerAppAttemptAdded(
    ApplicationAttemptId attemptId, MockRM rm) throws InterruptedException {
  int tick = 0;
  // Wait for at most 5 sec
  while (null == ((AbstractYarnScheduler) rm.getResourceScheduler())
      .getApplicationAttempt(attemptId) && tick < 50) {
    Thread.sleep(100);
    if (tick % 10 == 0) {
      System.out.println("waiting for SchedulerApplicationAttempt="
          + attemptId + " added.");
    }
    tick++;
  }
  Assert.assertNotNull("Timed out waiting for SchedulerApplicationAttempt=" +
    attemptId + " to be added.", ((AbstractYarnScheduler)
      rm.getResourceScheduler()).getApplicationAttempt(attemptId));
}
 

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
@SuppressWarnings("rawtypes")
public static void waitForNumContainersToRecover(int num, MockRM rm,
    ApplicationAttemptId attemptId) throws Exception {
  AbstractYarnScheduler scheduler =
      (AbstractYarnScheduler) rm.getResourceScheduler();
  SchedulerApplicationAttempt attempt =
      scheduler.getApplicationAttempt(attemptId);
  while (attempt == null) {
    System.out.println("Wait for scheduler attempt " + attemptId
        + " to be created");
    Thread.sleep(200);
    attempt = scheduler.getApplicationAttempt(attemptId);
  }
  while (attempt.getLiveContainers().size() < num) {
    System.out.println("Wait for " + num
        + " containers to recover. currently: "
        + attempt.getLiveContainers().size());
    Thread.sleep(200);
  }
}
 

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
/**
 * Wait until an app is removed from scheduler.
 * @param appId the id of an app
 * @param timeoutMsecs the length of timeout in milliseconds
 * @throws InterruptedException
 *         if interrupted while waiting for app removed
 */
public void waitForAppRemovedFromScheduler(ApplicationId appId,
    long timeoutMsecs) throws InterruptedException {
  int timeWaiting = 0;

  Map<ApplicationId, SchedulerApplication> apps  =
      ((AbstractYarnScheduler) getResourceScheduler())
          .getSchedulerApplications();
  while (apps.containsKey(appId)) {
    if (timeWaiting >= timeoutMsecs) {
      break;
    }
    LOG.info("wait for app removed, " + appId);
    Thread.sleep(WAIT_MS_PER_LOOP);
    timeWaiting += WAIT_MS_PER_LOOP;
  }
  Assert.assertTrue("app is not removed from scheduler (timeout).",
      !apps.containsKey(appId));
  LOG.info("app is removed from scheduler, " + appId);
}
 

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
public static void waitForNumContainersToRecover(int num, MockRM rm,
    ApplicationAttemptId attemptId) throws Exception {
  AbstractYarnScheduler scheduler =
      (AbstractYarnScheduler) rm.getResourceScheduler();
  SchedulerApplicationAttempt attempt =
      scheduler.getApplicationAttempt(attemptId);
  while (attempt == null) {
    LOG.info("Wait for scheduler attempt " + attemptId
        + " to be created");
    Thread.sleep(200);
    attempt = scheduler.getApplicationAttempt(attemptId);
  }
  while (attempt.getLiveContainers().size() < num) {
    LOG.info("Wait for " + num
        + " containers to recover. currently: "
        + attempt.getLiveContainers().size());
    Thread.sleep(200);
  }
}
 

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
@Inject
public NMHeartBeatHandler(InterceptorRegistry registry, AbstractYarnScheduler yarnScheduler, MyriadDriver myriadDriver,
                          YarnNodeCapacityManager yarnNodeCapacityMgr, OfferLifecycleManager offerLifecycleMgr,
                          NodeStore nodeStore, SchedulerState state, NodeManagerConfiguration conf) {

  if (registry != null) {
    registry.register(this);
  }

  this.yarnScheduler = yarnScheduler;
  this.myriadDriver = myriadDriver;
  this.yarnNodeCapacityMgr = yarnNodeCapacityMgr;
  this.offerLifecycleMgr = offerLifecycleMgr;
  this.nodeStore = nodeStore;
  this.state = state;
  this.conf = conf;
}
 

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
private RMContext generateRMContext(AbstractYarnScheduler<FiCaSchedulerApp, FiCaSchedulerNode> scheduler) throws Exception {
  Configuration conf = new Configuration();
  MockRMContext context = null;
  Dispatcher dispatcher = new MockDispatcher();

  RMApplicationHistoryWriter rmApplicationHistoryWriter = new RMApplicationHistoryWriter(); 
  AMLivelinessMonitor amLivelinessMonitor = new AMLivelinessMonitor(dispatcher);
  AMLivelinessMonitor amFinishingMonitor = new AMLivelinessMonitor(dispatcher);    
  RMDelegationTokenSecretManager delegationTokenSecretManager = new RMDelegationTokenSecretManager(1, 1, 1, 1, context);

  context = new MockRMContext();
  context.setStateStore(TestObjectFactory.getStateStore(conf, "tmp/myriad-operations-test"));
  context.setAmLivelinessMonitor(amLivelinessMonitor);
  context.setAmFinishingMonitor(amFinishingMonitor);
  context.setRMApplicationHistoryWriter(rmApplicationHistoryWriter);
  context.setRMDelegationTokenSecretManager(delegationTokenSecretManager);
  return context;
}
 

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
private static RMContext mockRMContext(List<RMAppState> states) {
  final ConcurrentMap<ApplicationId, RMApp> applicationsMaps = Maps
      .newConcurrentMap();
  int i = 0;
  for (RMAppState state : states) {
    MockRMApp app = new MockRMApp(i, i, state) {
      @Override
      public RMAppMetrics getRMAppMetrics() {
        return new RMAppMetrics(Resource.newInstance(0, 0, 0), 0, 0, 0, 0, 0);
      }
      @Override
      public YarnApplicationState createApplicationState() {
        return YarnApplicationState.ACCEPTED;
      }
    };
    RMAppAttempt attempt = mock(RMAppAttempt.class);
    app.setCurrentAppAttempt(attempt);
    applicationsMaps.put(app.getApplicationId(), app);
    i++;
  }

  RMContextImpl rmContext =  new RMContextImpl(null, null, null, null,
      null, null, null, null, null, null) {
    @Override
    public ConcurrentMap<ApplicationId, RMApp> getRMApps() {
      return applicationsMaps;
    }
    @Override
    public ResourceScheduler getScheduler() {
      return mock(AbstractYarnScheduler.class);
    }
  };
  return rmContext;
}
 

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
public static void initializeRMContext(int numContainers,
    AbstractYarnScheduler scheduler, RMContext mockRMContext) {

  when(mockRMContext.getScheduler()).thenReturn(scheduler);
  Resource r = calculateClusterResource(numContainers);
  doReturn(r).when(scheduler).getClusterResource();
}
 

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
@Test
public void testAddAndRemoveAppFromFiFoScheduler() throws Exception {
  Configuration conf = new Configuration();
  conf.setClass(YarnConfiguration.RM_SCHEDULER, FifoScheduler.class,
    ResourceScheduler.class);
  MockRM rm = new MockRM(conf);
  @SuppressWarnings("unchecked")
  AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode> fs =
      (AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>) rm
        .getResourceScheduler();
  TestSchedulerUtils.verifyAppAddedAndRemovedFromScheduler(
    fs.getSchedulerApplications(), fs, "queue");
}
 

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
@Test
public void testAddAndRemoveAppFromFairScheduler() throws Exception {
  AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode> scheduler =
      (AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>) resourceManager
        .getResourceScheduler();
  TestSchedulerUtils.verifyAppAddedAndRemovedFromScheduler(
    scheduler.getSchedulerApplications(), scheduler, "default");
}
 

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
@Test(timeout = 20000)
public void testContainersNotRecoveredForCompletedApps() throws Exception {
  MemoryRMStateStore memStore = new MemoryRMStateStore();
  memStore.init(conf);
  rm1 = new MockRM(conf, memStore);
  rm1.start();
  MockNM nm1 =
      new MockNM("127.0.0.1:1234", 8192, rm1.getResourceTrackerService());
  nm1.registerNode();
  RMApp app1 = rm1.submitApp(200);
  MockAM am1 = MockRM.launchAndRegisterAM(app1, rm1, nm1);
  MockRM.finishAMAndVerifyAppState(app1, rm1, nm1, am1);

  rm2 = new MockRM(conf, memStore);
  rm2.start();
  nm1.setResourceTrackerService(rm2.getResourceTrackerService());
  NMContainerStatus runningContainer =
      TestRMRestart.createNMContainerStatus(am1.getApplicationAttemptId(), 2,
        ContainerState.RUNNING);
  NMContainerStatus completedContainer =
      TestRMRestart.createNMContainerStatus(am1.getApplicationAttemptId(), 3,
        ContainerState.COMPLETE);
  nm1.registerNode(Arrays.asList(runningContainer, completedContainer), null);
  RMApp recoveredApp1 =
      rm2.getRMContext().getRMApps().get(app1.getApplicationId());
  assertEquals(RMAppState.FINISHED, recoveredApp1.getState());

  // Wait for RM to settle down on recovering containers;
  Thread.sleep(3000);

  AbstractYarnScheduler scheduler =
      (AbstractYarnScheduler) rm2.getResourceScheduler();

  // scheduler should not recover containers for finished apps.
  assertNull(scheduler.getRMContainer(runningContainer.getContainerId()));
  assertNull(scheduler.getRMContainer(completedContainer.getContainerId()));
}
 

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
@Test (timeout = 30000)
public void testAMContainerStatusWithRMRestart() throws Exception {  
  MemoryRMStateStore memStore = new MemoryRMStateStore();
  memStore.init(conf);
  rm1 = new MockRM(conf, memStore);
  rm1.start();
  MockNM nm1 =
      new MockNM("127.0.0.1:1234", 8192, rm1.getResourceTrackerService());
  nm1.registerNode();
  RMApp app1_1 = rm1.submitApp(1024);
  MockAM am1_1 = MockRM.launchAndRegisterAM(app1_1, rm1, nm1);
  
  RMAppAttempt attempt0 = app1_1.getCurrentAppAttempt();
  AbstractYarnScheduler scheduler =
      ((AbstractYarnScheduler) rm1.getResourceScheduler());
  
  Assert.assertTrue(scheduler.getRMContainer(
      attempt0.getMasterContainer().getId()).isAMContainer());

  // Re-start RM
  rm2 = new MockRM(conf, memStore);
  rm2.start();
  nm1.setResourceTrackerService(rm2.getResourceTrackerService());

  List<NMContainerStatus> am1_1Containers =
      createNMContainerStatusForApp(am1_1);
  nm1.registerNode(am1_1Containers, null);

  // Wait for RM to settle down on recovering containers;
  waitForNumContainersToRecover(2, rm2, am1_1.getApplicationAttemptId());

  scheduler = ((AbstractYarnScheduler) rm2.getResourceScheduler());
  Assert.assertTrue(scheduler.getRMContainer(
      attempt0.getMasterContainer().getId()).isAMContainer());
}
 

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
@Test (timeout = 20000)
public void testRecoverSchedulerAppAndAttemptSynchronously() throws Exception {
  // start RM
  MemoryRMStateStore memStore = new MemoryRMStateStore();
  memStore.init(conf);
  rm1 = new MockRM(conf, memStore);
  rm1.start();
  MockNM nm1 =
      new MockNM("127.0.0.1:1234", 15120, rm1.getResourceTrackerService());
  nm1.registerNode();

  // create app and launch the AM
  RMApp app0 = rm1.submitApp(200);
  MockAM am0 = MockRM.launchAndRegisterAM(app0, rm1, nm1);

  rm2 = new MockRM(conf, memStore);
  rm2.start();
  nm1.setResourceTrackerService(rm2.getResourceTrackerService());
  // scheduler app/attempt is immediately available after RM is re-started.
  Assert.assertNotNull(rm2.getResourceScheduler().getSchedulerAppInfo(
    am0.getApplicationAttemptId()));

  // getTransferredContainers should not throw NPE.
  ((AbstractYarnScheduler) rm2.getResourceScheduler())
    .getTransferredContainers(am0.getApplicationAttemptId());

  List<NMContainerStatus> containers = createNMContainerStatusForApp(am0);
  nm1.registerNode(containers, null);
  waitForNumContainersToRecover(2, rm2, am0.getApplicationAttemptId());
}
 

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
@SuppressWarnings("unchecked")
@Override
public void serviceInit(Configuration conf) throws Exception {
  ((AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>)
      scheduler).init(conf);
  super.serviceInit(conf);
}
 

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
@SuppressWarnings("unchecked")
@Override
public void serviceStart() throws Exception {
  ((AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>)
      scheduler).start();
  super.serviceStart();
}
 

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
@SuppressWarnings("unchecked")
@Override
public void serviceStop() throws Exception {
  ((AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>)
      scheduler).stop();
  super.serviceStop();
}
 

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
public AppAttemptInfo(ResourceManager rm, RMAppAttempt attempt, String user,
    String schemePrefix) {
  this.startTime = 0;
  this.containerId = "";
  this.nodeHttpAddress = "";
  this.nodeId = "";
  this.logsLink = "";
  this.blacklistedNodes = "";
  if (attempt != null) {
    this.id = attempt.getAppAttemptId().getAttemptId();
    this.startTime = attempt.getStartTime();
    Container masterContainer = attempt.getMasterContainer();
    if (masterContainer != null) {
      this.containerId = masterContainer.getId().toString();
      this.nodeHttpAddress = masterContainer.getNodeHttpAddress();
      this.nodeId = masterContainer.getNodeId().toString();
      this.logsLink = WebAppUtils.getRunningLogURL(schemePrefix
          + masterContainer.getNodeHttpAddress(),
          ConverterUtils.toString(masterContainer.getId()), user);
      if (rm.getResourceScheduler() instanceof AbstractYarnScheduler) {
        AbstractYarnScheduler ayScheduler =
            (AbstractYarnScheduler) rm.getResourceScheduler();
        SchedulerApplicationAttempt sattempt =
            ayScheduler.getApplicationAttempt(attempt.getAppAttemptId());
        if (sattempt != null) {
          blacklistedNodes =
              StringUtils.join(sattempt.getBlacklistedNodes(), ", ");
        }
      }
    }
  }
}
 

import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler; //导入依赖的package包/类
public static Set<String> getBlacklistedNodes(ResourceManager rm,
    ApplicationAttemptId appid) {
  if (rm.getResourceScheduler() instanceof AbstractYarnScheduler) {
    AbstractYarnScheduler ayScheduler =
        (AbstractYarnScheduler) rm.getResourceScheduler();
    SchedulerApplicationAttempt attempt =
        ayScheduler.getApplicationAttempt(appid);
    if (attempt != null) {
      return attempt.getBlacklistedNodes();
    }
  }
  return null;
}