org.apache.hadoop.yarn.server.resourcemanager.scheduler.event.NodeAddedSchedulerEvent Java Examples

@Test (expected = YarnException.class)
public void testMoveWouldViolateMaxResourcesConstraints() throws Exception {
  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(conf, resourceManager.getRMContext());

  QueueManager queueMgr = scheduler.getQueueManager();
  FSLeafQueue oldQueue = queueMgr.getLeafQueue("queue1", true);
  queueMgr.getLeafQueue("queue2", true);
  scheduler.getAllocationConfiguration().maxQueueResources.put("root.queue2",
      Resource.newInstance(1024, 1));

  ApplicationAttemptId appAttId =
      createSchedulingRequest(1024, 1, "queue1", "user1", 3);
  RMNode node = MockNodes.newNodeInfo(1, Resources.createResource(2048, 2));
  NodeAddedSchedulerEvent nodeEvent = new NodeAddedSchedulerEvent(node);
  NodeUpdateSchedulerEvent updateEvent = new NodeUpdateSchedulerEvent(node);
  scheduler.handle(nodeEvent);
  scheduler.handle(updateEvent);
  scheduler.handle(updateEvent);
  
  assertEquals(Resource.newInstance(2048, 2), oldQueue.getResourceUsage());
  scheduler.moveApplication(appAttId.getApplicationId(), "queue2");
}
 

@Test
public void testNoMoreCpuOnNode() throws IOException {
  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(conf, resourceManager.getRMContext());

  RMNode node1 = MockNodes.newNodeInfo(1, Resources.createResource(2048, 1, 1),
      1, "127.0.0.1");
  NodeAddedSchedulerEvent nodeEvent1 = new NodeAddedSchedulerEvent(node1);
  scheduler.handle(nodeEvent1);
  
  ApplicationAttemptId attId = createSchedulingRequest(1024, 1, 1, "default",
      "user1", 2);
  FSAppAttempt app = scheduler.getSchedulerApp(attId);
  scheduler.update();

  NodeUpdateSchedulerEvent updateEvent = new NodeUpdateSchedulerEvent(node1);
  scheduler.handle(updateEvent);
  assertEquals(1, app.getLiveContainers().size());
  scheduler.handle(updateEvent);
  assertEquals(1, app.getLiveContainers().size());
}
 

@Test
public void testNoMoreCpuOnNode() throws IOException {
  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(conf, resourceManager.getRMContext());

  RMNode node1 = MockNodes.newNodeInfo(1, Resources.createResource(2048, 1),
      1, "127.0.0.1");
  NodeAddedSchedulerEvent nodeEvent1 = new NodeAddedSchedulerEvent(node1);
  scheduler.handle(nodeEvent1);
  
  ApplicationAttemptId attId = createSchedulingRequest(1024, 1, "default",
      "user1", 2);
  FSAppAttempt app = scheduler.getSchedulerApp(attId);
  scheduler.update();

  NodeUpdateSchedulerEvent updateEvent = new NodeUpdateSchedulerEvent(node1);
  scheduler.handle(updateEvent);
  assertEquals(1, app.getLiveContainers().size());
  scheduler.handle(updateEvent);
  assertEquals(1, app.getLiveContainers().size());
}
 

private void registerNodeAndSubmitApp(
    int memory, int vcores, int appContainers, int appMemory) {
  RMNode node1 = MockNodes.newNodeInfo(
      1, Resources.createResource(memory, vcores), 1, "node1");
  NodeAddedSchedulerEvent nodeEvent1 = new NodeAddedSchedulerEvent(node1);
  scheduler.handle(nodeEvent1);

  assertEquals("Incorrect amount of resources in the cluster",
      memory, scheduler.rootMetrics.getAvailableMB());
  assertEquals("Incorrect amount of resources in the cluster",
      vcores, scheduler.rootMetrics.getAvailableVirtualCores());

  createSchedulingRequest(appMemory, "queueA", "user1", appContainers);
  scheduler.update();
  // Sufficient node check-ins to fully schedule containers
  for (int i = 0; i < 3; i++) {
    NodeUpdateSchedulerEvent nodeUpdate1 = new NodeUpdateSchedulerEvent(node1);
    scheduler.handle(nodeUpdate1);
  }
  assertEquals("app1's request is not met",
      memory - appContainers * appMemory,
      scheduler.rootMetrics.getAvailableMB());
}
 

@Test
public void testAggregateCapacityTracking() throws Exception {
  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(conf, resourceManager.getRMContext());

  // Add a node
  RMNode node1 =
      MockNodes
          .newNodeInfo(1, Resources.createResource(1024), 1, "127.0.0.1");
  NodeAddedSchedulerEvent nodeEvent1 = new NodeAddedSchedulerEvent(node1);
  scheduler.handle(nodeEvent1);
  assertEquals(1024, scheduler.getClusterResource().getMemory());

  // Add another node
  RMNode node2 =
      MockNodes.newNodeInfo(1, Resources.createResource(512), 2, "127.0.0.2");
  NodeAddedSchedulerEvent nodeEvent2 = new NodeAddedSchedulerEvent(node2);
  scheduler.handle(nodeEvent2);
  assertEquals(1536, scheduler.getClusterResource().getMemory());

  // Remove the first node
  NodeRemovedSchedulerEvent nodeEvent3 = new NodeRemovedSchedulerEvent(node1);
  scheduler.handle(nodeEvent3);
  assertEquals(512, scheduler.getClusterResource().getMemory());
}
 

private void registerNodeAndSubmitApp(
    int memory, int vcores, int gcores, int appContainers, int appMemory) {
  RMNode node1 = MockNodes.newNodeInfo(
      1, Resources.createResource(memory, vcores, gcores), 1, "node1");
  NodeAddedSchedulerEvent nodeEvent1 = new NodeAddedSchedulerEvent(node1);
  scheduler.handle(nodeEvent1);

  assertEquals("Incorrect amount of resources in the cluster",
      memory, scheduler.rootMetrics.getAvailableMB());
  assertEquals("Incorrect amount of resources in the cluster",
      vcores, scheduler.rootMetrics.getAvailableVirtualCores());

  createSchedulingRequest(appMemory, "queueA", "user1", appContainers);
  scheduler.update();
  // Sufficient node check-ins to fully schedule containers
  for (int i = 0; i < 3; i++) {
    NodeUpdateSchedulerEvent nodeUpdate1 = new NodeUpdateSchedulerEvent(node1);
    scheduler.handle(nodeUpdate1);
  }
  assertEquals("app1's request is not met",
      memory - appContainers * appMemory,
      scheduler.rootMetrics.getAvailableMB());
}
 

@Test
public void testAsyncScheduling() throws Exception {
  Configuration conf = new Configuration();
  conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class,
      ResourceScheduler.class);
  MockRM rm = new MockRM(conf);
  rm.start();
  CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();

  final int NODES = 100;
  
  // Register nodes
  for (int i=0; i < NODES; ++i) {
    String host = "192.168.1." + i;
    RMNode node =
        MockNodes.newNodeInfo(0, MockNodes.newResource(4 * GB), 1, host);
    cs.handle(new NodeAddedSchedulerEvent(node));
  }
  
  // Now directly exercise the scheduling loop
  for (int i=0; i < NODES; ++i) {
    CapacityScheduler.schedule(cs);
  }
}
 

@Override
public NodeState transition(RMNodeImpl rmNode, RMNodeEvent event) {
  RMNodeStatusEvent statusEvent = (RMNodeStatusEvent) event;

  // Switch the last heartbeatresponse.
  rmNode.latestNodeHeartBeatResponse = statusEvent.getLatestResponse();
  NodeHealthStatus remoteNodeHealthStatus = statusEvent.getNodeHealthStatus();
  rmNode.setHealthReport(remoteNodeHealthStatus.getHealthReport());
  rmNode.setLastHealthReportTime(
      remoteNodeHealthStatus.getLastHealthReportTime());
  if (remoteNodeHealthStatus.getIsNodeHealthy()) {
    rmNode.context.getDispatcher().getEventHandler().handle(
        new NodeAddedSchedulerEvent(rmNode));
    rmNode.context.getDispatcher().getEventHandler().handle(
            new NodesListManagerEvent(
                NodesListManagerEventType.NODE_USABLE, rmNode));
    // ??? how about updating metrics before notifying to ensure that
    // notifiers get update metadata because they will very likely query it
    // upon notification
    // Update metrics
    rmNode.updateMetricsForRejoinedNode(NodeState.UNHEALTHY);
    return NodeState.RUNNING;
  }

  return NodeState.UNHEALTHY;
}
 

@Test
public void testAsyncScheduling() throws Exception {
  Configuration conf = new Configuration();
  conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class,
      ResourceScheduler.class);
  MockRM rm = new MockRM(conf);
  rm.start();
  CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();

  final int NODES = 100;
  
  // Register nodes
  for (int i=0; i < NODES; ++i) {
    String host = "192.168.1." + i;
    RMNode node =
        MockNodes.newNodeInfo(0, MockNodes.newResource(4 * GB), 1, host);
    cs.handle(new NodeAddedSchedulerEvent(node));
  }
  
  // Now directly exercise the scheduling loop
  for (int i=0; i < NODES; ++i) {
    CapacityScheduler.schedule(cs);
  }
}
 

public static FairScheduler setupFairScheduler(
    ReservationSystemTestUtil testUtil,
    RMContext rmContext, Configuration conf, int numContainers) throws
    IOException {
  FairScheduler scheduler = new FairScheduler();
  scheduler.setRMContext(rmContext);

  when(rmContext.getScheduler()).thenReturn(scheduler);

  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(conf, rmContext);


  Resource resource = testUtil.calculateClusterResource(numContainers);
  RMNode node1 = MockNodes.newNodeInfo(1, resource, 1, "127.0.0.1");
  NodeAddedSchedulerEvent nodeEvent1 = new NodeAddedSchedulerEvent(node1);
  scheduler.handle(nodeEvent1);
  return scheduler;
}
 

@Test
  public void testRemoveNodeUpdatesRootQueueMetrics() throws IOException {
    scheduler.init(conf);
    scheduler.start();
    scheduler.reinitialize(conf, resourceManager.getRMContext());

    assertEquals(0, scheduler.getRootQueueMetrics().getAvailableMB());
	assertEquals(0, scheduler.getRootQueueMetrics().getAvailableVirtualCores());
    
    RMNode node1 = MockNodes.newNodeInfo(1, Resources.createResource(1024, 4), 1,
        "127.0.0.1");
    NodeAddedSchedulerEvent addEvent = new NodeAddedSchedulerEvent(node1);
    scheduler.handle(addEvent);
    
    assertEquals(1024, scheduler.getRootQueueMetrics().getAvailableMB());
    assertEquals(4, scheduler.getRootQueueMetrics().getAvailableVirtualCores());
    scheduler.update(); // update shouldn't change things
    assertEquals(1024, scheduler.getRootQueueMetrics().getAvailableMB());
    assertEquals(4, scheduler.getRootQueueMetrics().getAvailableVirtualCores());
    
    NodeRemovedSchedulerEvent removeEvent = new NodeRemovedSchedulerEvent(node1);
    scheduler.handle(removeEvent);
    
    assertEquals(0, scheduler.getRootQueueMetrics().getAvailableMB());
    assertEquals(0, scheduler.getRootQueueMetrics().getAvailableVirtualCores());
    scheduler.update(); // update shouldn't change things
    assertEquals(0, scheduler.getRootQueueMetrics().getAvailableMB());
    assertEquals(0, scheduler.getRootQueueMetrics().getAvailableVirtualCores());
}
 

@Test
public void testSimpleFairShareCalculation() throws IOException {
  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(conf, resourceManager.getRMContext());

  // Add one big node (only care about aggregate capacity)
  RMNode node1 =
      MockNodes.newNodeInfo(1, Resources.createResource(10 * 1024), 1,
          "127.0.0.1");
  NodeAddedSchedulerEvent nodeEvent1 = new NodeAddedSchedulerEvent(node1);
  scheduler.handle(nodeEvent1);

  // Have two queues which want entire cluster capacity
  createSchedulingRequest(10 * 1024, "queue1", "user1");
  createSchedulingRequest(10 * 1024, "queue2", "user1");
  createSchedulingRequest(10 * 1024, "root.default", "user1");

  scheduler.update();
  scheduler.getQueueManager().getRootQueue()
      .setSteadyFairShare(scheduler.getClusterResource());
  scheduler.getQueueManager().getRootQueue().recomputeSteadyShares();

  Collection<FSLeafQueue> queues = scheduler.getQueueManager().getLeafQueues();
  assertEquals(3, queues.size());
  
  // Divided three ways - between the two queues and the default queue
  for (FSLeafQueue p : queues) {
    assertEquals(3414, p.getFairShare().getMemory());
    assertEquals(3414, p.getMetrics().getFairShareMB());
    assertEquals(3414, p.getSteadyFairShare().getMemory());
    assertEquals(3414, p.getMetrics().getSteadyFairShareMB());
  }
}
 

@Test
public void testSteadyFairShareWithQueueCreatedRuntime() throws Exception {
  conf.setClass(CommonConfigurationKeys.HADOOP_SECURITY_GROUP_MAPPING,
      SimpleGroupsMapping.class, GroupMappingServiceProvider.class);
  conf.set(FairSchedulerConfiguration.USER_AS_DEFAULT_QUEUE, "true");
  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(conf, resourceManager.getRMContext());

  // Add one node
  RMNode node1 =
      MockNodes
          .newNodeInfo(1, Resources.createResource(6144), 1, "127.0.0.1");
  NodeAddedSchedulerEvent nodeEvent1 = new NodeAddedSchedulerEvent(node1);
  scheduler.handle(nodeEvent1);
  assertEquals(6144, scheduler.getClusterResource().getMemory());
  assertEquals(6144, scheduler.getQueueManager().getRootQueue()
      .getSteadyFairShare().getMemory());
  assertEquals(6144, scheduler.getQueueManager()
      .getLeafQueue("default", false).getSteadyFairShare().getMemory());

  // Submit one application
  ApplicationAttemptId appAttemptId1 = createAppAttemptId(1, 1);
  createApplicationWithAMResource(appAttemptId1, "default", "user1", null);
  assertEquals(3072, scheduler.getQueueManager()
      .getLeafQueue("default", false).getSteadyFairShare().getMemory());
  assertEquals(3072, scheduler.getQueueManager()
      .getLeafQueue("user1", false).getSteadyFairShare().getMemory());
}
 

@Test(timeout = 3000)
public void testMaxAssignWithZeroMemoryContainers() throws Exception {
  conf.setBoolean(FairSchedulerConfiguration.ASSIGN_MULTIPLE, true);
  conf.setInt(YarnConfiguration.RM_SCHEDULER_MINIMUM_ALLOCATION_MB, 0);
  
  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(conf, resourceManager.getRMContext());

  RMNode node =
      MockNodes.newNodeInfo(1, Resources.createResource(16384, 16), 0,
          "127.0.0.1");
  NodeAddedSchedulerEvent nodeEvent = new NodeAddedSchedulerEvent(node);
  NodeUpdateSchedulerEvent updateEvent = new NodeUpdateSchedulerEvent(node);
  scheduler.handle(nodeEvent);

  ApplicationAttemptId attId =
      createSchedulingRequest(0, 1, "root.default", "user", 8);
  FSAppAttempt app = scheduler.getSchedulerApp(attId);

  // set maxAssign to 2: only 2 containers should be allocated
  scheduler.maxAssign = 2;
  scheduler.update();
  scheduler.handle(updateEvent);
  assertEquals("Incorrect number of containers allocated", 2, app
      .getLiveContainers().size());

  // set maxAssign to -1: all remaining containers should be allocated
  scheduler.maxAssign = -1;
  scheduler.update();
  scheduler.handle(updateEvent);
  assertEquals("Incorrect number of containers allocated", 8, app
      .getLiveContainers().size());
}
 

@Test(timeout = 3000)
public void testMaxAssignWithZeroMemoryContainers() throws Exception {
  conf.setBoolean(FairSchedulerConfiguration.ASSIGN_MULTIPLE, true);
  conf.setInt(YarnConfiguration.RM_SCHEDULER_MINIMUM_ALLOCATION_MB, 0);
  
  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(conf, resourceManager.getRMContext());

  RMNode node =
      MockNodes.newNodeInfo(1, Resources.createResource(16384, 16, 16), 0,
          "127.0.0.1");
  NodeAddedSchedulerEvent nodeEvent = new NodeAddedSchedulerEvent(node);
  NodeUpdateSchedulerEvent updateEvent = new NodeUpdateSchedulerEvent(node);
  scheduler.handle(nodeEvent);

  ApplicationAttemptId attId =
      createSchedulingRequest(0, 1, 1, "root.default", "user", 8);
  FSAppAttempt app = scheduler.getSchedulerApp(attId);

  // set maxAssign to 2: only 2 containers should be allocated
  scheduler.maxAssign = 2;
  scheduler.update();
  scheduler.handle(updateEvent);
  assertEquals("Incorrect number of containers allocated", 2, app
      .getLiveContainers().size());

  // set maxAssign to -1: all remaining containers should be allocated
  scheduler.maxAssign = -1;
  scheduler.update();
  scheduler.handle(updateEvent);
  assertEquals("Incorrect number of containers allocated", 8, app
      .getLiveContainers().size());
}
 

@Test
  public void testRemoveNodeUpdatesRootQueueMetrics() throws IOException {
    scheduler.init(conf);
    scheduler.start();
    scheduler.reinitialize(conf, resourceManager.getRMContext());

    assertEquals(0, scheduler.getRootQueueMetrics().getAvailableMB());
	assertEquals(0, scheduler.getRootQueueMetrics().getAvailableVirtualCores());
    
    RMNode node1 = MockNodes.newNodeInfo(1, Resources.createResource(1024, 4, 4), 1,
        "127.0.0.1");
    NodeAddedSchedulerEvent addEvent = new NodeAddedSchedulerEvent(node1);
    scheduler.handle(addEvent);
    
    assertEquals(1024, scheduler.getRootQueueMetrics().getAvailableMB());
    assertEquals(4, scheduler.getRootQueueMetrics().getAvailableVirtualCores());
    scheduler.update(); // update shouldn't change things
    assertEquals(1024, scheduler.getRootQueueMetrics().getAvailableMB());
    assertEquals(4, scheduler.getRootQueueMetrics().getAvailableVirtualCores());
    
    NodeRemovedSchedulerEvent removeEvent = new NodeRemovedSchedulerEvent(node1);
    scheduler.handle(removeEvent);
    
    assertEquals(0, scheduler.getRootQueueMetrics().getAvailableMB());
    assertEquals(0, scheduler.getRootQueueMetrics().getAvailableVirtualCores());
    scheduler.update(); // update shouldn't change things
    assertEquals(0, scheduler.getRootQueueMetrics().getAvailableMB());
    assertEquals(0, scheduler.getRootQueueMetrics().getAvailableVirtualCores());
}
 

@Test (timeout = 60000)
public void testSchedulingDelay() throws InterruptedException {
  // Add one node
  String host = "127.0.0.1";
  RMNode node1 = MockNodes.newNodeInfo(
      1, Resources.createResource(4096, 4), 1, host);
  NodeAddedSchedulerEvent nodeEvent1 = new NodeAddedSchedulerEvent(node1);
  scheduler.handle(nodeEvent1);
  NodeUpdateSchedulerEvent nodeUpdateEvent = new NodeUpdateSchedulerEvent(node1);
  scheduler.handle(nodeUpdateEvent);

  // Create one application and submit one each of node-local, rack-local
  // and ANY requests
  ApplicationAttemptId appAttemptId =
      createAppAttemptId(this.APP_ID++, this.ATTEMPT_ID++);
  createMockRMApp(appAttemptId);

  scheduler.addApplication(appAttemptId.getApplicationId(), "queue11", "user11", false);
  scheduler.addApplicationAttempt(appAttemptId, false, false);
  List<ResourceRequest> ask = new ArrayList<>();
  ask.add(createResourceRequest(1024, 1, ResourceRequest.ANY, 1, 1, true));
  scheduler.allocate(
      appAttemptId, ask, new ArrayList<ContainerId>(), null, null);
  FSAppAttempt app = scheduler.getSchedulerApp(appAttemptId);

  // Advance time and let continuous scheduling kick in
  mockClock.tick(1);
  while (1024 != app.getCurrentConsumption().getMemory()) {
    Thread.sleep(100);
  }
  assertEquals(1024, app.getCurrentConsumption().getMemory());
}
 

@Test
public void testNumClusterNodes() throws Exception {
  YarnConfiguration conf = new YarnConfiguration();
  CapacityScheduler cs = new CapacityScheduler();
  cs.setConf(conf);
  RMContext rmContext = TestUtils.getMockRMContext();
  cs.setRMContext(rmContext);
  CapacitySchedulerConfiguration csConf =
      new CapacitySchedulerConfiguration();
  setupQueueConfiguration(csConf);
  cs.init(csConf);
  cs.start();
  assertEquals(0, cs.getNumClusterNodes());

  RMNode n1 = MockNodes.newNodeInfo(0, MockNodes.newResource(4 * GB), 1);
  RMNode n2 = MockNodes.newNodeInfo(0, MockNodes.newResource(2 * GB), 2);
  cs.handle(new NodeAddedSchedulerEvent(n1));
  cs.handle(new NodeAddedSchedulerEvent(n2));
  assertEquals(2, cs.getNumClusterNodes());

  cs.handle(new NodeRemovedSchedulerEvent(n1));
  assertEquals(1, cs.getNumClusterNodes());
  cs.handle(new NodeAddedSchedulerEvent(n1));
  assertEquals(2, cs.getNumClusterNodes());
  cs.handle(new NodeRemovedSchedulerEvent(n2));
  cs.handle(new NodeRemovedSchedulerEvent(n1));
  assertEquals(0, cs.getNumClusterNodes());

  cs.stop();
}
 

@Override
public void transition(RMNodeImpl rmNode, RMNodeEvent event) {
  // Inform the scheduler
  RMNodeStartedEvent startEvent = (RMNodeStartedEvent) event;
  List<NMContainerStatus> containers = null;

  String host = rmNode.nodeId.getHost();
  if (rmNode.context.getInactiveRMNodes().containsKey(host)) {
    // Old node rejoining
    RMNode previouRMNode = rmNode.context.getInactiveRMNodes().get(host);
    rmNode.context.getInactiveRMNodes().remove(host);
    rmNode.updateMetricsForRejoinedNode(previouRMNode.getState());
  } else {
    // Increment activeNodes explicitly because this is a new node.
    ClusterMetrics.getMetrics().incrNumActiveNodes();
    containers = startEvent.getNMContainerStatuses();
    if (containers != null && !containers.isEmpty()) {
      for (NMContainerStatus container : containers) {
        if (container.getContainerState() == ContainerState.RUNNING) {
          rmNode.launchedContainers.add(container.getContainerId());
        }
      }
    }
  }
  
  if (null != startEvent.getRunningApplications()) {
    for (ApplicationId appId : startEvent.getRunningApplications()) {
      handleRunningAppOnNode(rmNode, rmNode.context, appId, rmNode.nodeId);
    }
  }

  rmNode.context.getDispatcher().getEventHandler()
    .handle(new NodeAddedSchedulerEvent(rmNode, containers));
  rmNode.context.getDispatcher().getEventHandler().handle(
    new NodesListManagerEvent(
        NodesListManagerEventType.NODE_USABLE, rmNode));
}
 

@Test (timeout = 50000)
public void testReconnectedNode() throws Exception {
  CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
  conf.setQueues("default", new String[] {"default"});
  conf.setCapacity("default", 100);
  FifoScheduler fs = new FifoScheduler();
  fs.init(conf);
  fs.start();
  // mock rmContext to avoid NPE.
  RMContext context = mock(RMContext.class);
  fs.reinitialize(conf, null);
  fs.setRMContext(context);

  RMNode n1 =
      MockNodes.newNodeInfo(0, MockNodes.newResource(4 * GB), 1, "127.0.0.2");
  RMNode n2 =
      MockNodes.newNodeInfo(0, MockNodes.newResource(2 * GB), 2, "127.0.0.3");

  fs.handle(new NodeAddedSchedulerEvent(n1));
  fs.handle(new NodeAddedSchedulerEvent(n2));
  fs.handle(new NodeUpdateSchedulerEvent(n1));
  Assert.assertEquals(6 * GB, fs.getRootQueueMetrics().getAvailableMB());

  // reconnect n1 with downgraded memory
  n1 =
      MockNodes.newNodeInfo(0, MockNodes.newResource(2 * GB), 1, "127.0.0.2");
  fs.handle(new NodeRemovedSchedulerEvent(n1));
  fs.handle(new NodeAddedSchedulerEvent(n1));
  fs.handle(new NodeUpdateSchedulerEvent(n1));

  Assert.assertEquals(4 * GB, fs.getRootQueueMetrics().getAvailableMB());
  fs.stop();
}
 

/**
 * Two apps on one queue, one app on another
 */
@Test
public void testBasicDRFWithQueues() throws Exception {
  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(conf, resourceManager.getRMContext());

  RMNode node = MockNodes.newNodeInfo(1, BuilderUtils.newResource(8192, 7, 7),
      1, "127.0.0.1");
  NodeAddedSchedulerEvent nodeEvent = new NodeAddedSchedulerEvent(node);
  scheduler.handle(nodeEvent);

  ApplicationAttemptId appAttId1 = createSchedulingRequest(3072, 1, 1, "queue1",
      "user1", 2);
  FSAppAttempt app1 = scheduler.getSchedulerApp(appAttId1);
  ApplicationAttemptId appAttId2 = createSchedulingRequest(2048, 2, 2, "queue1",
      "user1", 2);
  FSAppAttempt app2 = scheduler.getSchedulerApp(appAttId2);
  ApplicationAttemptId appAttId3 = createSchedulingRequest(1024, 2, 2, "queue2",
      "user1", 2);
  FSAppAttempt app3 = scheduler.getSchedulerApp(appAttId3);
  
  DominantResourceFairnessPolicy drfPolicy = new DominantResourceFairnessPolicy();
  drfPolicy.initialize(scheduler.getClusterResource());
  scheduler.getQueueManager().getQueue("root").setPolicy(drfPolicy);
  scheduler.getQueueManager().getQueue("queue1").setPolicy(drfPolicy);
  scheduler.update();

  NodeUpdateSchedulerEvent updateEvent = new NodeUpdateSchedulerEvent(node);
  scheduler.handle(updateEvent);
  Assert.assertEquals(1, app1.getLiveContainers().size());
  scheduler.handle(updateEvent);
  Assert.assertEquals(1, app3.getLiveContainers().size());
  scheduler.handle(updateEvent);
  Assert.assertEquals(2, app3.getLiveContainers().size());
  scheduler.handle(updateEvent);
  Assert.assertEquals(1, app2.getLiveContainers().size());
}
 

private org.apache.hadoop.yarn.server.resourcemanager.NodeManager
    registerNode(String hostName, int containerManagerPort, int httpPort,
        String rackName, Resource capability) throws IOException,
        YarnException {
  org.apache.hadoop.yarn.server.resourcemanager.NodeManager nm = 
      new org.apache.hadoop.yarn.server.resourcemanager.NodeManager(
          hostName, containerManagerPort, httpPort, rackName, capability,
          resourceManager);
  NodeAddedSchedulerEvent nodeAddEvent1 = 
      new NodeAddedSchedulerEvent(resourceManager.getRMContext()
          .getRMNodes().get(nm.getNodeId()));
  resourceManager.getResourceScheduler().handle(nodeAddEvent1);
  return nm;
}
 

@Test
public void testSimpleFairShareCalculation() throws IOException {
  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(conf, resourceManager.getRMContext());

  // Add one big node (only care about aggregate capacity)
  RMNode node1 =
      MockNodes.newNodeInfo(1, Resources.createResource(10 * 1024), 1,
          "127.0.0.1");
  NodeAddedSchedulerEvent nodeEvent1 = new NodeAddedSchedulerEvent(node1);
  scheduler.handle(nodeEvent1);

  // Have two queues which want entire cluster capacity
  createSchedulingRequest(10 * 1024, "queue1", "user1");
  createSchedulingRequest(10 * 1024, "queue2", "user1");
  createSchedulingRequest(10 * 1024, "root.default", "user1");

  scheduler.update();
  scheduler.getQueueManager().getRootQueue()
      .setSteadyFairShare(scheduler.getClusterResource());
  scheduler.getQueueManager().getRootQueue().recomputeSteadyShares();

  Collection<FSLeafQueue> queues = scheduler.getQueueManager().getLeafQueues();
  assertEquals(3, queues.size());
  
  // Divided three ways - between the two queues and the default queue
  for (FSLeafQueue p : queues) {
    assertEquals(3414, p.getFairShare().getMemory());
    assertEquals(3414, p.getMetrics().getFairShareMB());
    assertEquals(3414, p.getSteadyFairShare().getMemory());
    assertEquals(3414, p.getMetrics().getSteadyFairShareMB());
  }
}
 

@Test
public void testReconnectedNode() throws Exception {
  CapacitySchedulerConfiguration csConf =
      new CapacitySchedulerConfiguration();
  setupQueueConfiguration(csConf);
  CapacityScheduler cs = new CapacityScheduler();
  cs.setConf(new YarnConfiguration());
  cs.setRMContext(resourceManager.getRMContext());
  cs.init(csConf);
  cs.start();
  cs.reinitialize(csConf, new RMContextImpl(null, null, null, null,
    null, null, new RMContainerTokenSecretManager(csConf),
    new NMTokenSecretManagerInRM(csConf),
    new ClientToAMTokenSecretManagerInRM(), null));

  RMNode n1 = MockNodes.newNodeInfo(0, MockNodes.newResource(4 * GB), 1);
  RMNode n2 = MockNodes.newNodeInfo(0, MockNodes.newResource(2 * GB), 2);

  cs.handle(new NodeAddedSchedulerEvent(n1));
  cs.handle(new NodeAddedSchedulerEvent(n2));

  Assert.assertEquals(6 * GB, cs.getClusterResource().getMemory());

  // reconnect n1 with downgraded memory
  n1 = MockNodes.newNodeInfo(0, MockNodes.newResource(2 * GB), 1);
  cs.handle(new NodeRemovedSchedulerEvent(n1));
  cs.handle(new NodeAddedSchedulerEvent(n1));

  Assert.assertEquals(4 * GB, cs.getClusterResource().getMemory());
  cs.stop();
}
 

@Test
public void testContinuousSchedulingWithNodeRemoved() throws Exception {
  // Disable continuous scheduling, will invoke continuous scheduling once manually
  scheduler.init(conf);
  scheduler.start();
  Assert.assertTrue("Continuous scheduling should be disabled.",
      !scheduler.isContinuousSchedulingEnabled());

  // Add two nodes
  RMNode node1 =
      MockNodes.newNodeInfo(1, Resources.createResource(8 * 1024, 8, 8), 1,
          "127.0.0.1");
  NodeAddedSchedulerEvent nodeEvent1 = new NodeAddedSchedulerEvent(node1);
  scheduler.handle(nodeEvent1);
  RMNode node2 =
      MockNodes.newNodeInfo(1, Resources.createResource(8 * 1024, 8, 8), 2,
          "127.0.0.2");
  NodeAddedSchedulerEvent nodeEvent2 = new NodeAddedSchedulerEvent(node2);
  scheduler.handle(nodeEvent2);
  Assert.assertEquals("We should have two alive nodes.",
      2, scheduler.getNumClusterNodes());

  // Remove one node
  NodeRemovedSchedulerEvent removeNode1 = new NodeRemovedSchedulerEvent(node1);
  scheduler.handle(removeNode1);
  Assert.assertEquals("We should only have one alive node.",
      1, scheduler.getNumClusterNodes());

  // Invoke the continuous scheduling once
  try {
    scheduler.continuousSchedulingAttempt();
  } catch (Exception e) {
    fail("Exception happened when doing continuous scheduling. " +
      e.toString());
  }
}
 

private org.apache.hadoop.yarn.server.resourcemanager.NodeManager
    registerNode(String hostName, int containerManagerPort, int httpPort,
        String rackName, Resource capability)
        throws IOException, YarnException {
  org.apache.hadoop.yarn.server.resourcemanager.NodeManager nm =
      new org.apache.hadoop.yarn.server.resourcemanager.NodeManager(
          hostName, containerManagerPort, httpPort, rackName, capability,
          resourceManager);
  NodeAddedSchedulerEvent nodeAddEvent1 = 
      new NodeAddedSchedulerEvent(resourceManager.getRMContext()
          .getRMNodes().get(nm.getNodeId()));
  resourceManager.getResourceScheduler().handle(nodeAddEvent1);
  return nm;
}
 

@Test
public void testContinuousSchedulingWithNodeRemoved() throws Exception {
  // Disable continuous scheduling, will invoke continuous scheduling once manually
  scheduler.init(conf);
  scheduler.start();
  Assert.assertTrue("Continuous scheduling should be disabled.",
      !scheduler.isContinuousSchedulingEnabled());

  // Add two nodes
  RMNode node1 =
      MockNodes.newNodeInfo(1, Resources.createResource(8 * 1024, 8), 1,
          "127.0.0.1");
  NodeAddedSchedulerEvent nodeEvent1 = new NodeAddedSchedulerEvent(node1);
  scheduler.handle(nodeEvent1);
  RMNode node2 =
      MockNodes.newNodeInfo(1, Resources.createResource(8 * 1024, 8), 2,
          "127.0.0.2");
  NodeAddedSchedulerEvent nodeEvent2 = new NodeAddedSchedulerEvent(node2);
  scheduler.handle(nodeEvent2);
  Assert.assertEquals("We should have two alive nodes.",
      2, scheduler.getNumClusterNodes());

  // Remove one node
  NodeRemovedSchedulerEvent removeNode1 = new NodeRemovedSchedulerEvent(node1);
  scheduler.handle(removeNode1);
  Assert.assertEquals("We should only have one alive node.",
      1, scheduler.getNumClusterNodes());

  // Invoke the continuous scheduling once
  try {
    scheduler.continuousSchedulingAttempt();
  } catch (Exception e) {
    fail("Exception happened when doing continuous scheduling. " +
      e.toString());
  }
}
 

private void createClusterWithQueuesAndOneNode(int mem, int vCores, int gCores,
    String policy) throws IOException {
  PrintWriter out = new PrintWriter(new FileWriter(ALLOC_FILE));
  out.println("<?xml version=\"1.0\"?>");
  out.println("<allocations>");
  out.println("<queue name=\"root\" >");
  out.println("   <queue name=\"parentA\" >");
  out.println("       <weight>8</weight>");
  out.println("       <queue name=\"childA1\" />");
  out.println("       <queue name=\"childA2\" />");
  out.println("       <queue name=\"childA3\" />");
  out.println("       <queue name=\"childA4\" />");
  out.println("   </queue>");
  out.println("   <queue name=\"parentB\" >");
  out.println("       <weight>1</weight>");
  out.println("       <queue name=\"childB1\" />");
  out.println("       <queue name=\"childB2\" />");
  out.println("   </queue>");
  out.println("</queue>");
  out.println("<defaultQueueSchedulingPolicy>" + policy
      + "</defaultQueueSchedulingPolicy>");
  out.println("</allocations>");
  out.close();

  resourceManager = new MockRM(conf);
  resourceManager.start();
  scheduler = (FairScheduler) resourceManager.getResourceScheduler();

  RMNode node1 = MockNodes.newNodeInfo(1,
      Resources.createResource(mem, vCores, gCores), 1, "127.0.0.1");
  NodeAddedSchedulerEvent nodeEvent1 = new NodeAddedSchedulerEvent(node1);
  scheduler.handle(nodeEvent1);
}
 

@Test (timeout = 60000)
public void testSchedulingDelay() throws InterruptedException {
  // Add one node
  String host = "127.0.0.1";
  RMNode node1 = MockNodes.newNodeInfo(
      1, Resources.createResource(4096, 4, 4), 1, host);
  NodeAddedSchedulerEvent nodeEvent1 = new NodeAddedSchedulerEvent(node1);
  scheduler.handle(nodeEvent1);
  NodeUpdateSchedulerEvent nodeUpdateEvent = new NodeUpdateSchedulerEvent(node1);
  scheduler.handle(nodeUpdateEvent);

  // Create one application and submit one each of node-local, rack-local
  // and ANY requests
  ApplicationAttemptId appAttemptId =
      createAppAttemptId(this.APP_ID++, this.ATTEMPT_ID++);
  createMockRMApp(appAttemptId);

  scheduler.addApplication(appAttemptId.getApplicationId(), "queue11", "user11", false);
  scheduler.addApplicationAttempt(appAttemptId, false, false);
  List<ResourceRequest> ask = new ArrayList<>();
  ask.add(createResourceRequest(1024, 1, 1, ResourceRequest.ANY, 1, 1, true));
  scheduler.allocate(
      appAttemptId, ask, new ArrayList<ContainerId>(), null, null);
  FSAppAttempt app = scheduler.getSchedulerApp(appAttemptId);

  // Advance time and let continuous scheduling kick in
  mockClock.tick(1);
  while (1024 != app.getCurrentConsumption().getMemory()) {
    Thread.sleep(100);
  }
  assertEquals(1024, app.getCurrentConsumption().getMemory());
}
 

@Test (timeout = 5000)
public void testExpiredContainer() {
  // Start the node
  node.handle(new RMNodeStartedEvent(null, null, null));
  verify(scheduler).handle(any(NodeAddedSchedulerEvent.class));
  
  // Expire a container
  ContainerId completedContainerId = BuilderUtils.newContainerId(
      BuilderUtils.newApplicationAttemptId(
          BuilderUtils.newApplicationId(0, 0), 0), 0);
  node.handle(new RMNodeCleanContainerEvent(null, completedContainerId));
  Assert.assertEquals(1, node.getContainersToCleanUp().size());
  
  // Now verify that scheduler isn't notified of an expired container
  // by checking number of 'completedContainers' it got in the previous event
  RMNodeStatusEvent statusEvent = getMockRMNodeStatusEvent();
  ContainerStatus containerStatus = mock(ContainerStatus.class);
  doReturn(completedContainerId).when(containerStatus).getContainerId();
  doReturn(Collections.singletonList(containerStatus)).
      when(statusEvent).getContainers();
  node.handle(statusEvent);
  /* Expect the scheduler call handle function 2 times
   * 1. RMNode status from new to Running, handle the add_node event
   * 2. handle the node update event
   */
  verify(scheduler,times(2)).handle(any(NodeUpdateSchedulerEvent.class));     
}