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

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());
}
 

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 (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 (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, 1));

  ApplicationAttemptId appAttId =
      createSchedulingRequest(1024, 1, 1, "queue1", "user1", 3);
  RMNode node = MockNodes.newNodeInfo(1, Resources.createResource(2048, 2, 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, 2), oldQueue.getResourceUsage());
  scheduler.moveApplication(appAttId.getApplicationId(), "queue2");
}
 

@Override
public void afterSchedulerEventHandled(SchedulerEvent event) {

  try {
    switch (event.getType()) {
      case NODE_UPDATE:
        onNodeUpdated((NodeUpdateSchedulerEvent) event);
        break;

      case NODE_REMOVED:
        onNodeRemoved((NodeRemovedSchedulerEvent) event);
        break;

      default:
        break;
    }
  } catch (ClassCastException e) {
    LOGGER.error("incorrect event object", e);
  }
}
 

@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());
}
 

@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 (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();
}
 

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

  QueueManager queueMgr = scheduler.getQueueManager();
  FSLeafQueue oldQueue = queueMgr.getLeafQueue("queue1", true);
  FSLeafQueue targetQueue = queueMgr.getLeafQueue("queue2", true);

  ApplicationAttemptId appAttId =
      createSchedulingRequest(1024, 1, 0, "queue1", "user1", 3);
  ApplicationId appId = appAttId.getApplicationId();
  RMNode node = MockNodes.newNodeInfo(1, Resources.createResource(1024));
  NodeAddedSchedulerEvent nodeEvent = new NodeAddedSchedulerEvent(node);
  NodeUpdateSchedulerEvent updateEvent = new NodeUpdateSchedulerEvent(node);
  scheduler.handle(nodeEvent);
  scheduler.handle(updateEvent);
  
  assertEquals(Resource.newInstance(1024, 1, 0), oldQueue.getResourceUsage());
  scheduler.update();
  assertEquals(Resource.newInstance(3072, 3, 0), oldQueue.getDemand());
  
  scheduler.moveApplication(appId, "queue2");
  FSAppAttempt app = scheduler.getSchedulerApp(appAttId);
  assertSame(targetQueue, app.getQueue());
  assertFalse(oldQueue.isRunnableApp(app));
  assertTrue(targetQueue.isRunnableApp(app));
  assertEquals(Resource.newInstance(0, 0, 0), oldQueue.getResourceUsage());
  assertEquals(Resource.newInstance(1024, 1, 0), targetQueue.getResourceUsage());
  assertEquals(0, oldQueue.getNumRunnableApps());
  assertEquals(1, targetQueue.getNumRunnableApps());
  assertEquals(1, queueMgr.getRootQueue().getNumRunnableApps());
  
  scheduler.update();
  assertEquals(Resource.newInstance(0, 0, 0), oldQueue.getDemand());
  assertEquals(Resource.newInstance(3072, 3, 0), targetQueue.getDemand());
}
 

private NodeId getNodeIdForSchedulerEvent(SchedulerEvent event) {
  switch (event.getType()) {
    case NODE_ADDED:
      return ((NodeAddedSchedulerEvent) event).getAddedRMNode().getNodeID();
    case NODE_REMOVED:
      return ((NodeRemovedSchedulerEvent) event).getRemovedRMNode().getNodeID();
    case NODE_UPDATE:
      return ((NodeUpdateSchedulerEvent) event).getRMNode().getNodeID();
    case NODE_RESOURCE_UPDATE:
      return ((NodeResourceUpdateSchedulerEvent) event).getRMNode().getNodeID();
  }
  return null;
}
 

/**
 * 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());
}
 

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

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

  ApplicationAttemptId appAttId1 = createSchedulingRequest(2048, 1, 1, "queue1",
      "user1", 2);
  FSAppAttempt app1 = scheduler.getSchedulerApp(appAttId1);
  ApplicationAttemptId appAttId2 = createSchedulingRequest(1024, 2, 2, "queue1",
      "user1", 2);
  FSAppAttempt app2 = scheduler.getSchedulerApp(appAttId2);

  DominantResourceFairnessPolicy drfPolicy = new DominantResourceFairnessPolicy();
  drfPolicy.initialize(scheduler.getClusterResource());
  scheduler.getQueueManager().getQueue("queue1").setPolicy(drfPolicy);
  scheduler.update();

  // First both apps get a container
  // Then the first gets another container because its dominant share of
  // 2048/8192 is less than the other's of 2/5
  NodeUpdateSchedulerEvent updateEvent = new NodeUpdateSchedulerEvent(node);
  scheduler.handle(updateEvent);
  Assert.assertEquals(1, app1.getLiveContainers().size());
  Assert.assertEquals(0, app2.getLiveContainers().size());

  scheduler.handle(updateEvent);
  Assert.assertEquals(1, app1.getLiveContainers().size());
  Assert.assertEquals(1, app2.getLiveContainers().size());

  scheduler.handle(updateEvent);
  Assert.assertEquals(2, app1.getLiveContainers().size());
  Assert.assertEquals(1, app2.getLiveContainers().size());
}
 

@Override
public void afterSchedulerEventHandled(SchedulerEvent event) {
  switch (event.getType()) {
    case NODE_ADDED:
      if (!(event instanceof NodeAddedSchedulerEvent)) {
        LOGGER.error("{} not an instance of {}", event.getClass().getName(), NodeAddedSchedulerEvent.class.getName());
        return;
      }

      NodeAddedSchedulerEvent nodeAddedEvent = (NodeAddedSchedulerEvent) event;
      NodeId nodeId = nodeAddedEvent.getAddedRMNode().getNodeID();
      String host = nodeId.getHost();

      SchedulerNode node = yarnScheduler.getSchedulerNode(nodeId);
      nodeStore.add(node);
      LOGGER.info("afterSchedulerEventHandled: NM registration from node {}", host);
      break;

    case NODE_UPDATE:
      if (!(event instanceof NodeUpdateSchedulerEvent)) {
        LOGGER.error("{} not an instance of {}", event.getClass().getName(), NodeUpdateSchedulerEvent.class.getName());
        return;
      }

      RMNode rmNode = ((NodeUpdateSchedulerEvent) event).getRMNode();
      handleContainerAllocation(rmNode);

      break;

    default:
      break;
  }
}
 

@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(timeout = 3000)
public void testMaxAssign() throws Exception {
  conf.setBoolean(FairSchedulerConfiguration.ASSIGN_MULTIPLE, true);
  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(1024, "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 testNodeUpdateBeforeAppAttemptInit() throws Exception {
  FifoScheduler scheduler = new FifoScheduler();
  MockRM rm = new MockRM(conf);
  scheduler.setRMContext(rm.getRMContext());
  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(conf, rm.getRMContext());

  RMNode node = MockNodes.newNodeInfo(1,
          Resources.createResource(1024, 4, 4), 1, "127.0.0.1");
  scheduler.handle(new NodeAddedSchedulerEvent(node));

  ApplicationId appId = ApplicationId.newInstance(0, 1);
  scheduler.addApplication(appId, "queue1", "user1", false);

  NodeUpdateSchedulerEvent updateEvent = new NodeUpdateSchedulerEvent(node);
  try {
    scheduler.handle(updateEvent);
  } catch (NullPointerException e) {
      Assert.fail();
  }

  ApplicationAttemptId attId = ApplicationAttemptId.newInstance(appId, 1);
  scheduler.addApplicationAttempt(attId, false, false);

  rm.stop();
}
 

@Test (timeout = 5000)
public void testStatusChange(){
  //Start the node
  node.handle(new RMNodeStartedEvent(null, null, null));
  //Add info to the queue first
  node.setNextHeartBeat(false);

  ContainerId completedContainerId1 = BuilderUtils.newContainerId(
      BuilderUtils.newApplicationAttemptId(
          BuilderUtils.newApplicationId(0, 0), 0), 0);
  ContainerId completedContainerId2 = BuilderUtils.newContainerId(
      BuilderUtils.newApplicationAttemptId(
          BuilderUtils.newApplicationId(1, 1), 1), 1);
      
  RMNodeStatusEvent statusEvent1 = getMockRMNodeStatusEvent();
  RMNodeStatusEvent statusEvent2 = getMockRMNodeStatusEvent();

  ContainerStatus containerStatus1 = mock(ContainerStatus.class);
  ContainerStatus containerStatus2 = mock(ContainerStatus.class);

  doReturn(completedContainerId1).when(containerStatus1).getContainerId();
  doReturn(Collections.singletonList(containerStatus1))
      .when(statusEvent1).getContainers();
   
  doReturn(completedContainerId2).when(containerStatus2).getContainerId();
  doReturn(Collections.singletonList(containerStatus2))
      .when(statusEvent2).getContainers();

  verify(scheduler,times(1)).handle(any(NodeUpdateSchedulerEvent.class)); 
  node.handle(statusEvent1);
  node.handle(statusEvent2);
  verify(scheduler,times(1)).handle(any(NodeUpdateSchedulerEvent.class));
  Assert.assertEquals(2, node.getQueueSize());
  node.handle(new RMNodeEvent(node.getNodeID(), RMNodeEventType.EXPIRE));
  Assert.assertEquals(0, node.getQueueSize());
}
 

@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));     
}
 

@Test (timeout = 5000)
public void testStatusChange(){
  //Start the node
  node.handle(new RMNodeStartedEvent(null, null, null));
  //Add info to the queue first
  node.setNextHeartBeat(false);

  ContainerId completedContainerId1 = BuilderUtils.newContainerId(
      BuilderUtils.newApplicationAttemptId(
          BuilderUtils.newApplicationId(0, 0), 0), 0);
  ContainerId completedContainerId2 = BuilderUtils.newContainerId(
      BuilderUtils.newApplicationAttemptId(
          BuilderUtils.newApplicationId(1, 1), 1), 1);
      
  RMNodeStatusEvent statusEvent1 = getMockRMNodeStatusEvent();
  RMNodeStatusEvent statusEvent2 = getMockRMNodeStatusEvent();

  ContainerStatus containerStatus1 = mock(ContainerStatus.class);
  ContainerStatus containerStatus2 = mock(ContainerStatus.class);

  doReturn(completedContainerId1).when(containerStatus1).getContainerId();
  doReturn(Collections.singletonList(containerStatus1))
      .when(statusEvent1).getContainers();
   
  doReturn(completedContainerId2).when(containerStatus2).getContainerId();
  doReturn(Collections.singletonList(containerStatus2))
      .when(statusEvent2).getContainers();

  verify(scheduler,times(1)).handle(any(NodeUpdateSchedulerEvent.class)); 
  node.handle(statusEvent1);
  node.handle(statusEvent2);
  verify(scheduler,times(1)).handle(any(NodeUpdateSchedulerEvent.class));
  Assert.assertEquals(2, node.getQueueSize());
  node.handle(new RMNodeEvent(node.getNodeID(), RMNodeEventType.EXPIRE));
  Assert.assertEquals(0, node.getQueueSize());
}
 

@Test
public void testNodeUpdateBeforeAppAttemptInit() throws Exception {
  FifoScheduler scheduler = new FifoScheduler();
  MockRM rm = new MockRM(conf);
  scheduler.setRMContext(rm.getRMContext());
  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(conf, rm.getRMContext());

  RMNode node = MockNodes.newNodeInfo(1,
          Resources.createResource(1024, 4), 1, "127.0.0.1");
  scheduler.handle(new NodeAddedSchedulerEvent(node));

  ApplicationId appId = ApplicationId.newInstance(0, 1);
  scheduler.addApplication(appId, "queue1", "user1", false);

  NodeUpdateSchedulerEvent updateEvent = new NodeUpdateSchedulerEvent(node);
  try {
    scheduler.handle(updateEvent);
  } catch (NullPointerException e) {
      Assert.fail();
  }

  ApplicationAttemptId attId = ApplicationAttemptId.newInstance(appId, 1);
  scheduler.addApplicationAttempt(attId, false, false);

  rm.stop();
}
 

@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();
}
 

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

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

  // Queue 1 requests full capacity of node
  createSchedulingRequest(1024, "queue1", "user1", 1);
  scheduler.update();
  NodeUpdateSchedulerEvent updateEvent = new NodeUpdateSchedulerEvent(node1);
  scheduler.handle(updateEvent);

  // Now queue 2 requests likewise
  createSchedulingRequest(1024, "queue2", "user1", 1);
  scheduler.update();
  scheduler.handle(updateEvent);

  // Make sure reserved memory gets updated correctly
  assertEquals(1024, scheduler.rootMetrics.getReservedMB());
  
  // Now another node checks in with capacity
  RMNode node2 = MockNodes.newNodeInfo(1, Resources.createResource(1024));
  NodeAddedSchedulerEvent nodeEvent2 = new NodeAddedSchedulerEvent(node2);
  NodeUpdateSchedulerEvent updateEvent2 = new NodeUpdateSchedulerEvent(node2);
  scheduler.handle(nodeEvent2);
  scheduler.handle(updateEvent2);


  // The old reservation should still be there...
  assertEquals(1024, scheduler.rootMetrics.getReservedMB());

  // ... but it should disappear when we update the first node.
  scheduler.handle(updateEvent);
  assertEquals(0, scheduler.rootMetrics.getReservedMB());
}
 

@Test (timeout = 5000)
public void testMultipleContainersWaitingForReservation() throws IOException {
  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);

  // Request full capacity of node
  createSchedulingRequest(1024, "queue1", "user1", 1);
  scheduler.update();
  NodeUpdateSchedulerEvent updateEvent = new NodeUpdateSchedulerEvent(node1);
  scheduler.handle(updateEvent);

  ApplicationAttemptId attId1 = createSchedulingRequest(1024, "queue2", "user2", 1);
  ApplicationAttemptId attId2 = createSchedulingRequest(1024, "queue3", "user3", 1);
  
  scheduler.update();
  scheduler.handle(updateEvent);
  
  // One container should get reservation and the other should get nothing
  assertEquals(1024,
      scheduler.getSchedulerApp(attId1).getCurrentReservation().getMemory());
  assertEquals(0,
      scheduler.getSchedulerApp(attId2).getCurrentReservation().getMemory());
}
 

@Test(timeout = 3000)
public void testMaxAssign() throws Exception {
  conf.setBoolean(FairSchedulerConfiguration.ASSIGN_MULTIPLE, true);
  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(1024, "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), 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
public void testBasicDRFAssignment() throws Exception {
  scheduler.init(conf);
  scheduler.start();
  scheduler.reinitialize(conf, resourceManager.getRMContext());

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

  ApplicationAttemptId appAttId1 = createSchedulingRequest(2048, 1, "queue1",
      "user1", 2);
  FSAppAttempt app1 = scheduler.getSchedulerApp(appAttId1);
  ApplicationAttemptId appAttId2 = createSchedulingRequest(1024, 2, "queue1",
      "user1", 2);
  FSAppAttempt app2 = scheduler.getSchedulerApp(appAttId2);

  DominantResourceFairnessPolicy drfPolicy = new DominantResourceFairnessPolicy();
  drfPolicy.initialize(scheduler.getClusterResource());
  scheduler.getQueueManager().getQueue("queue1").setPolicy(drfPolicy);
  scheduler.update();

  // First both apps get a container
  // Then the first gets another container because its dominant share of
  // 2048/8192 is less than the other's of 2/5
  NodeUpdateSchedulerEvent updateEvent = new NodeUpdateSchedulerEvent(node);
  scheduler.handle(updateEvent);
  Assert.assertEquals(1, app1.getLiveContainers().size());
  Assert.assertEquals(0, app2.getLiveContainers().size());

  scheduler.handle(updateEvent);
  Assert.assertEquals(1, app1.getLiveContainers().size());
  Assert.assertEquals(1, app2.getLiveContainers().size());

  scheduler.handle(updateEvent);
  Assert.assertEquals(2, app1.getLiveContainers().size());
  Assert.assertEquals(1, app2.getLiveContainers().size());
}
 

/**
 * 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),
      1, "127.0.0.1");
  NodeAddedSchedulerEvent nodeEvent = new NodeAddedSchedulerEvent(node);
  scheduler.handle(nodeEvent);

  ApplicationAttemptId appAttId1 = createSchedulingRequest(3072, 1, "queue1",
      "user1", 2);
  FSAppAttempt app1 = scheduler.getSchedulerApp(appAttId1);
  ApplicationAttemptId appAttId2 = createSchedulingRequest(2048, 2, "queue1",
      "user1", 2);
  FSAppAttempt app2 = scheduler.getSchedulerApp(appAttId2);
  ApplicationAttemptId appAttId3 = createSchedulingRequest(1024, 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());
}
 

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

  QueueManager queueMgr = scheduler.getQueueManager();
  FSLeafQueue oldQueue = queueMgr.getLeafQueue("queue1", true);
  FSLeafQueue targetQueue = queueMgr.getLeafQueue("queue2", true);

  ApplicationAttemptId appAttId =
      createSchedulingRequest(1024, 1, "queue1", "user1", 3);
  ApplicationId appId = appAttId.getApplicationId();
  RMNode node = MockNodes.newNodeInfo(1, Resources.createResource(1024));
  NodeAddedSchedulerEvent nodeEvent = new NodeAddedSchedulerEvent(node);
  NodeUpdateSchedulerEvent updateEvent = new NodeUpdateSchedulerEvent(node);
  scheduler.handle(nodeEvent);
  scheduler.handle(updateEvent);
  
  assertEquals(Resource.newInstance(1024, 1), oldQueue.getResourceUsage());
  scheduler.update();
  assertEquals(Resource.newInstance(3072, 3), oldQueue.getDemand());
  
  scheduler.moveApplication(appId, "queue2");
  FSAppAttempt app = scheduler.getSchedulerApp(appAttId);
  assertSame(targetQueue, app.getQueue());
  assertFalse(oldQueue.isRunnableApp(app));
  assertTrue(targetQueue.isRunnableApp(app));
  assertEquals(Resource.newInstance(0, 0), oldQueue.getResourceUsage());
  assertEquals(Resource.newInstance(1024, 1), targetQueue.getResourceUsage());
  assertEquals(0, oldQueue.getNumRunnableApps());
  assertEquals(1, targetQueue.getNumRunnableApps());
  assertEquals(1, queueMgr.getRootQueue().getNumRunnableApps());
  
  scheduler.update();
  assertEquals(Resource.newInstance(0, 0), oldQueue.getDemand());
  assertEquals(Resource.newInstance(3072, 3), targetQueue.getDemand());
}
 

@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));     
}