// This is to test container tokens are generated when the containers are
  // acquired by the AM, not when the containers are allocated
  @Test
  public void testContainerTokenGeneratedOnPullRequest() throws Exception {
    YarnConfiguration conf = new YarnConfiguration();
    YarnAPIStorageFactory.setConfiguration(conf);
    RMStorageFactory.setConfiguration(conf);
    conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);
    MockRM rm1 = new MockRM(conf);
    try {
      rm1.start();

      MockNM nm1 = rm1.registerNode("127.0.0.1:1234", 8000);
      RMApp app1 = rm1.submitApp(200);
      MockAM am1 = MockRM.launchAndRegisterAM(app1, rm1, nm1);
      // request a container.
      am1.allocate("127.0.0.1", 1024, 1, new ArrayList<ContainerId>());
      ContainerId containerId2 = ContainerId.newInstance(am1.getApplicationAttemptId(), 2);
      rm1.waitForState(nm1, containerId2, RMContainerState.ALLOCATED);

      RMContainer container = rm1.getResourceScheduler().getRMContainer(containerId2);
      // no container token is generated.
      Assert.assertEquals(containerId2, container.getContainerId());
      Assert.assertNull(container.getContainer().getContainerToken());

      // acquire the container.
      List<Container> containers =
          am1.allocate(new ArrayList<ResourceRequest>(), new ArrayList<ContainerId>())
              .getAllocatedContainers();
      Assert.assertEquals(containerId2, containers.get(0).getId());
      // container token is generated.
      Assert.assertNotNull(containers.get(0).getContainerToken());
    } finally {
      rm1.stop();
    }
  }
  @Test
  public void testNormalContainerAllocationWhenDNSUnavailable() throws Exception {
    YarnConfiguration conf = new YarnConfiguration();
    YarnAPIStorageFactory.setConfiguration(conf);
    RMStorageFactory.setConfiguration(conf);
    MockRM rm1 = new MockRM(conf);
    try {
      rm1.start();
      MockNM nm1 = rm1.registerNode("unknownhost:1234", 8000);
      RMApp app1 = rm1.submitApp(200);
      MockAM am1 = MockRM.launchAndRegisterAM(app1, rm1, nm1);

      // request a container.
      am1.allocate("127.0.0.1", 1024, 1, new ArrayList<ContainerId>());
      ContainerId containerId2 = ContainerId.newInstance(am1.getApplicationAttemptId(), 2);
      rm1.waitForState(nm1, containerId2, RMContainerState.ALLOCATED);

      // acquire the container.
      SecurityUtilTestHelper.setTokenServiceUseIp(true);
      List<Container> containers =
          am1.allocate(new ArrayList<ResourceRequest>(), new ArrayList<ContainerId>())
              .getAllocatedContainers();
      // not able to fetch the container;
      Assert.assertEquals(0, containers.size());

      SecurityUtilTestHelper.setTokenServiceUseIp(false);
      containers =
          am1.allocate(new ArrayList<ResourceRequest>(), new ArrayList<ContainerId>())
              .getAllocatedContainers();
      // should be able to fetch the container;
      Assert.assertEquals(1, containers.size());
    } finally {
      rm1.stop();
    }
  }
  @Test
  public void testPriorityWithPendingApplications() throws Exception {

    Configuration conf = new Configuration();
    conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);
    // Set Max Application Priority as 10
    conf.setInt(YarnConfiguration.MAX_CLUSTER_LEVEL_APPLICATION_PRIORITY, 10);
    MockRM rm = new MockRM(conf);
    rm.start();

    Priority appPriority1 = Priority.newInstance(5);
    MockNM nm1 = rm.registerNode("127.0.0.1:1234", 8 * GB);
    RMApp app1 = rm.submitApp(1 * GB, appPriority1);

    // kick the scheduler, 1 GB given to AM1, remaining 7GB on nm1
    MockAM am1 = MockRM.launchAM(app1, rm, nm1);
    am1.registerAppAttempt();

    // kick the scheduler, 7 containers will be allocated for App1
    List<Container> allocated1 = am1.allocateAndWaitForContainers("127.0.0.1", 7, 1 * GB, nm1);

    Assert.assertEquals(7, allocated1.size());
    Assert.assertEquals(1 * GB, allocated1.get(0).getResource().getMemory());

    // check node report, 8 GB used (1 AM and 7 containers) and 0 GB available
    SchedulerNodeReport report_nm1 = rm.getResourceScheduler().getNodeReport(nm1.getNodeId());
    Assert.assertEquals(8 * GB, report_nm1.getUsedResource().getMemory());
    Assert.assertEquals(0 * GB, report_nm1.getAvailableResource().getMemory());

    // Submit the second app App2 with priority 7
    Priority appPriority2 = Priority.newInstance(7);
    RMApp app2 = rm.submitApp(1 * GB, appPriority2);

    // Submit the third app App3 with priority 8
    Priority appPriority3 = Priority.newInstance(8);
    RMApp app3 = rm.submitApp(1 * GB, appPriority3);

    // Submit the second app App4 with priority 6
    Priority appPriority4 = Priority.newInstance(6);
    RMApp app4 = rm.submitApp(1 * GB, appPriority4);

    // Only one app can run as AM resource limit restricts it. Kill app1,
    // If app3 (highest priority among rest) gets active, it indicates that
    // priority is working with pendingApplications.
    rm.killApp(app1.getApplicationId());

    // kick the scheduler, app3 (high among pending) gets free space
    MockAM am3 = MockRM.launchAM(app3, rm, nm1);
    am3.registerAppAttempt();

    // check node report, 1 GB used and 7 GB available
    report_nm1 = rm.getResourceScheduler().getNodeReport(nm1.getNodeId());
    Assert.assertEquals(1 * GB, report_nm1.getUsedResource().getMemory());
    Assert.assertEquals(7 * GB, report_nm1.getAvailableResource().getMemory());

    rm.stop();
  }
  @Test
  public void testARRMResponseId() throws Exception {

    MockNM nm1 = rm.registerNode("h1:1234", 5000);

    RMApp app = rm.submitApp(2000);

    // Trigger the scheduling so the AM gets 'launched'
    nm1.nodeHeartbeat(true);

    RMAppAttempt attempt = app.getCurrentAppAttempt();
    MockAM am = rm.sendAMLaunched(attempt.getAppAttemptId());

    am.registerAppAttempt();

    AllocateRequest allocateRequest =
        BuilderUtils.newAllocateRequest(attempt.getAppAttemptId(), 0, 0F, null, null);

    AllocateResponse response = amService.allocate(allocateRequest);
    Assert.assertEquals(1, response.getResponseId());
    Assert.assertFalse(response.getReboot());
    allocateRequest =
        BuilderUtils.newAllocateRequest(
            attempt.getAppAttemptId(), response.getResponseId(), 0F, null, null);

    response = amService.allocate(allocateRequest);
    Assert.assertEquals(2, response.getResponseId());
    /* try resending */
    response = amService.allocate(allocateRequest);
    Assert.assertEquals(2, response.getResponseId());

    /** try sending old request again * */
    allocateRequest = BuilderUtils.newAllocateRequest(attempt.getAppAttemptId(), 0, 0F, null, null);
    response = amService.allocate(allocateRequest);
    Assert.assertTrue(response.getReboot());
  }
  @Test(timeout = 30000)
  public void testExcessReservationThanNodeManagerCapacity() throws Exception {
    YarnConfiguration conf = new YarnConfiguration();
    YarnAPIStorageFactory.setConfiguration(conf);
    RMStorageFactory.setConfiguration(conf);
    MockRM rm = new MockRM(conf);
    try {
      rm.start();

      // Register node1
      MockNM nm1 = rm.registerNode("127.0.0.1:1234", 2 * GB, 4);
      MockNM nm2 = rm.registerNode("127.0.0.1:2234", 3 * GB, 4);

      nm1.nodeHeartbeat(true);
      nm2.nodeHeartbeat(true);
      // HOP :: Sleep to allow previous events to be processed
      Thread.sleep(
          conf.getInt(
                  YarnConfiguration.HOPS_PENDING_EVENTS_RETRIEVAL_PERIOD,
                  YarnConfiguration.DEFAULT_HOPS_PENDING_EVENTS_RETRIEVAL_PERIOD)
              * 2);
      // wait..
      int waitCount = 20;
      int size = rm.getRMContext().getActiveRMNodes().size();
      while ((size = rm.getRMContext().getActiveRMNodes().size()) != 2 && waitCount-- > 0) {
        LOG.info("Waiting for node managers to register : " + size);
        Thread.sleep(100);
      }
      Assert.assertEquals(2, rm.getRMContext().getActiveRMNodes().size());
      // Submit an application
      RMApp app1 = rm.submitApp(128);

      // kick the scheduling
      nm1.nodeHeartbeat(true);
      RMAppAttempt attempt1 = app1.getCurrentAppAttempt();
      MockAM am1 = rm.sendAMLaunched(attempt1.getAppAttemptId(), nm1);
      am1.registerAppAttempt();

      LOG.info("sending container requests ");
      am1.addRequests(new String[] {"*"}, 3 * GB, 1, 1);
      AllocateResponse alloc1Response = am1.schedule(); // send the request

      // kick the scheduler
      nm1.nodeHeartbeat(true);
      int waitCounter = 20;
      LOG.info("heartbeating nm1");
      while (alloc1Response.getAllocatedContainers().size() < 1 && waitCounter-- > 0) {
        LOG.info("Waiting for containers to be created for app 1...");
        Thread.sleep(500);
        alloc1Response = am1.schedule();
      }
      LOG.info("received container : " + alloc1Response.getAllocatedContainers().size());

      // No container should be allocated.
      // Internally it should not been reserved.
      Assert.assertTrue(alloc1Response.getAllocatedContainers().size() == 0);

      LOG.info("heartbeating nm2");
      waitCounter = 20;
      nm2.nodeHeartbeat(true);
      while (alloc1Response.getAllocatedContainers().size() < 1 && waitCounter-- > 0) {
        LOG.info("Waiting for containers to be created for app 1...");
        Thread.sleep(500);
        alloc1Response = am1.schedule();
      }
      LOG.info("received container : " + alloc1Response.getAllocatedContainers().size());
      Assert.assertTrue(alloc1Response.getAllocatedContainers().size() == 1);
    } finally {
      rm.stop();
    }
  }
  @Test
  public void testApplicationPriorityAllocation() throws Exception {

    Configuration conf = new Configuration();
    conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);
    // Set Max Application Priority as 10
    conf.setInt(YarnConfiguration.MAX_CLUSTER_LEVEL_APPLICATION_PRIORITY, 10);
    MockRM rm = new MockRM(conf);
    rm.start();

    Priority appPriority1 = Priority.newInstance(5);
    MockNM nm1 = rm.registerNode("127.0.0.1:1234", 16 * GB);
    RMApp app1 = rm.submitApp(1 * GB, appPriority1);

    // kick the scheduler, 1 GB given to AM1, remaining 15GB on nm1
    MockAM am1 = MockRM.launchAM(app1, rm, nm1);
    am1.registerAppAttempt();

    // allocate 7 containers for App1
    List<Container> allocated1 = am1.allocateAndWaitForContainers("127.0.0.1", 7, 2 * GB, nm1);

    Assert.assertEquals(7, allocated1.size());
    Assert.assertEquals(2 * GB, allocated1.get(0).getResource().getMemory());

    // check node report, 15 GB used (1 AM and 7 containers) and 1 GB available
    SchedulerNodeReport report_nm1 = rm.getResourceScheduler().getNodeReport(nm1.getNodeId());
    Assert.assertEquals(15 * GB, report_nm1.getUsedResource().getMemory());
    Assert.assertEquals(1 * GB, report_nm1.getAvailableResource().getMemory());

    // Submit the second app App2 with priority 8 (Higher than App1)
    Priority appPriority2 = Priority.newInstance(8);
    RMApp app2 = rm.submitApp(1 * GB, appPriority2);

    // kick the scheduler, 1 GB which was free is given to AM of App2
    MockAM am2 = MockRM.launchAM(app2, rm, nm1);
    am2.registerAppAttempt();

    // check node report, 16 GB used and 0 GB available
    report_nm1 = rm.getResourceScheduler().getNodeReport(nm1.getNodeId());
    Assert.assertEquals(16 * GB, report_nm1.getUsedResource().getMemory());
    Assert.assertEquals(0 * GB, report_nm1.getAvailableResource().getMemory());

    // get scheduler
    CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();

    // get scheduler app
    FiCaSchedulerApp schedulerAppAttempt =
        cs.getSchedulerApplications().get(app1.getApplicationId()).getCurrentAppAttempt();

    // kill 2 containers of App1 to free up some space
    int counter = 0;
    for (Container c : allocated1) {
      if (++counter > 2) {
        break;
      }
      cs.killContainer(schedulerAppAttempt.getRMContainer(c.getId()));
    }

    // check node report, 12 GB used and 4 GB available
    report_nm1 = rm.getResourceScheduler().getNodeReport(nm1.getNodeId());
    Assert.assertEquals(12 * GB, report_nm1.getUsedResource().getMemory());
    Assert.assertEquals(4 * GB, report_nm1.getAvailableResource().getMemory());

    // send updated request for App1
    am1.allocate("127.0.0.1", 2 * GB, 10, new ArrayList<ContainerId>());

    // kick the scheduler, since App2 priority is more than App1, it will get
    // remaining cluster space.
    List<Container> allocated2 = am2.allocateAndWaitForContainers("127.0.0.1", 2, 2 * GB, nm1);

    // App2 has got 2 containers now.
    Assert.assertEquals(2, allocated2.size());

    // check node report, 16 GB used and 0 GB available
    report_nm1 = rm.getResourceScheduler().getNodeReport(nm1.getNodeId());
    Assert.assertEquals(16 * GB, report_nm1.getUsedResource().getMemory());
    Assert.assertEquals(0 * GB, report_nm1.getAvailableResource().getMemory());

    rm.stop();
  }
  @Test(timeout = 30000)
  public void testNodeUpdate() throws Exception {
    // set node -> label
    mgr.addToCluserNodeLabels(ImmutableSet.of("x", "y", "z"));

    // set mapping:
    // h1 -> x
    // h2 -> y
    mgr.addLabelsToNode(ImmutableMap.of(NodeId.newInstance("h1", 0), toSet("x")));
    mgr.addLabelsToNode(ImmutableMap.of(NodeId.newInstance("h2", 0), toSet("y")));

    // inject node label manager
    MockRM rm =
        new MockRM(getConfigurationWithQueueLabels(conf)) {
          @Override
          public RMNodeLabelsManager createNodeLabelManager() {
            return mgr;
          }
        };

    rm.getRMContext().setNodeLabelManager(mgr);
    rm.start();
    MockNM nm1 = rm.registerNode("h1:1234", 8000);
    MockNM nm2 = rm.registerNode("h2:1234", 8000);
    MockNM nm3 = rm.registerNode("h3:1234", 8000);

    ContainerId containerId;

    // launch an app to queue a1 (label = x), and check all container will
    // be allocated in h1
    RMApp app1 = rm.submitApp(GB, "app", "user", null, "a");
    MockAM am1 = MockRM.launchAndRegisterAM(app1, rm, nm3);

    // request a container.
    am1.allocate("*", GB, 1, new ArrayList<ContainerId>(), "x");
    containerId = ContainerId.newContainerId(am1.getApplicationAttemptId(), 2);
    Assert.assertTrue(rm.waitForState(nm1, containerId, RMContainerState.ALLOCATED, 10 * 1000));

    // check used resource:
    // queue-a used x=1G, ""=1G
    checkUsedResource(rm, "a", 1024, "x");
    checkUsedResource(rm, "a", 1024);

    // change h1's label to z, container should be killed
    mgr.replaceLabelsOnNode(ImmutableMap.of(NodeId.newInstance("h1", 0), toSet("z")));
    Assert.assertTrue(rm.waitForState(nm1, containerId, RMContainerState.KILLED, 10 * 1000));

    // check used resource:
    // queue-a used x=0G, ""=1G ("" not changed)
    checkUsedResource(rm, "a", 0, "x");
    checkUsedResource(rm, "a", 1024);

    // request a container with label = y
    am1.allocate("*", GB, 1, new ArrayList<ContainerId>(), "y");
    containerId = ContainerId.newContainerId(am1.getApplicationAttemptId(), 3);
    Assert.assertTrue(rm.waitForState(nm2, containerId, RMContainerState.ALLOCATED, 10 * 1000));

    // check used resource:
    // queue-a used y=1G, ""=1G
    checkUsedResource(rm, "a", 1024, "y");
    checkUsedResource(rm, "a", 1024);

    // change h2's label to no label, container should be killed
    mgr.replaceLabelsOnNode(
        ImmutableMap.of(NodeId.newInstance("h2", 0), CommonNodeLabelsManager.EMPTY_STRING_SET));
    Assert.assertTrue(rm.waitForState(nm1, containerId, RMContainerState.KILLED, 10 * 1000));

    // check used resource:
    // queue-a used x=0G, y=0G, ""=1G ("" not changed)
    checkUsedResource(rm, "a", 0, "x");
    checkUsedResource(rm, "a", 0, "y");
    checkUsedResource(rm, "a", 1024);

    containerId = ContainerId.newContainerId(am1.getApplicationAttemptId(), 1);

    // change h3's label to z, AM container should be killed
    mgr.replaceLabelsOnNode(ImmutableMap.of(NodeId.newInstance("h3", 0), toSet("z")));
    Assert.assertTrue(rm.waitForState(nm1, containerId, RMContainerState.KILLED, 10 * 1000));

    // check used resource:
    // queue-a used x=0G, y=0G, ""=1G ("" not changed)
    checkUsedResource(rm, "a", 0, "x");
    checkUsedResource(rm, "a", 0, "y");
    checkUsedResource(rm, "a", 0);

    rm.close();
  }
  @Test
  public void testAMRMUnusableNodes() throws Exception {

    MockNM nm1 = rm.registerNode("h1:1234", 5000);
    MockNM nm2 = rm.registerNode("h2:1234", 5000);
    MockNM nm3 = rm.registerNode("h3:1234", 5000);
    MockNM nm4 = rm.registerNode("h4:1234", 5000);

    RMApp app1 = rm.submitApp(2000);

    // Trigger the scheduling so the AM gets 'launched' on nm1
    nm1.nodeHeartbeat(true);

    RMAppAttempt attempt1 = app1.getCurrentAppAttempt();
    MockAM am1 = rm.sendAMLaunched(attempt1.getAppAttemptId());

    // register AM returns no unusable node
    am1.registerAppAttempt();

    // allocate request returns no updated node
    AllocateRequest allocateRequest1 =
        BuilderUtils.newAllocateRequest(attempt1.getAppAttemptId(), 0, 0F, null, null);
    AMResponse response1 = amService.allocate(allocateRequest1).getAMResponse();
    List<NodeReport> updatedNodes = response1.getUpdatedNodes();
    Assert.assertEquals(0, updatedNodes.size());

    syncNodeHeartbeat(nm4, false);

    // allocate request returns updated node
    allocateRequest1 =
        BuilderUtils.newAllocateRequest(
            attempt1.getAppAttemptId(), response1.getResponseId(), 0F, null, null);
    response1 = amService.allocate(allocateRequest1).getAMResponse();
    updatedNodes = response1.getUpdatedNodes();
    Assert.assertEquals(1, updatedNodes.size());
    NodeReport nr = updatedNodes.iterator().next();
    Assert.assertEquals(nm4.getNodeId(), nr.getNodeId());
    Assert.assertEquals(NodeState.UNHEALTHY, nr.getNodeState());

    // resending the allocate request returns the same result
    response1 = amService.allocate(allocateRequest1).getAMResponse();
    updatedNodes = response1.getUpdatedNodes();
    Assert.assertEquals(1, updatedNodes.size());
    nr = updatedNodes.iterator().next();
    Assert.assertEquals(nm4.getNodeId(), nr.getNodeId());
    Assert.assertEquals(NodeState.UNHEALTHY, nr.getNodeState());

    syncNodeLost(nm3);

    // subsequent allocate request returns delta
    allocateRequest1 =
        BuilderUtils.newAllocateRequest(
            attempt1.getAppAttemptId(), response1.getResponseId(), 0F, null, null);
    response1 = amService.allocate(allocateRequest1).getAMResponse();
    updatedNodes = response1.getUpdatedNodes();
    Assert.assertEquals(1, updatedNodes.size());
    nr = updatedNodes.iterator().next();
    Assert.assertEquals(nm3.getNodeId(), nr.getNodeId());
    Assert.assertEquals(NodeState.LOST, nr.getNodeState());

    // registering another AM gives it the complete failed list
    RMApp app2 = rm.submitApp(2000);
    // Trigger nm2 heartbeat so that AM gets launched on it
    nm2.nodeHeartbeat(true);
    RMAppAttempt attempt2 = app2.getCurrentAppAttempt();
    MockAM am2 = rm.sendAMLaunched(attempt2.getAppAttemptId());

    // register AM returns all unusable nodes
    am2.registerAppAttempt();

    // allocate request returns no updated node
    AllocateRequest allocateRequest2 =
        BuilderUtils.newAllocateRequest(attempt2.getAppAttemptId(), 0, 0F, null, null);
    AMResponse response2 = amService.allocate(allocateRequest2).getAMResponse();
    updatedNodes = response2.getUpdatedNodes();
    Assert.assertEquals(0, updatedNodes.size());

    syncNodeHeartbeat(nm4, true);

    // both AM's should get delta updated nodes
    allocateRequest1 =
        BuilderUtils.newAllocateRequest(
            attempt1.getAppAttemptId(), response1.getResponseId(), 0F, null, null);
    response1 = amService.allocate(allocateRequest1).getAMResponse();
    updatedNodes = response1.getUpdatedNodes();
    Assert.assertEquals(1, updatedNodes.size());
    nr = updatedNodes.iterator().next();
    Assert.assertEquals(nm4.getNodeId(), nr.getNodeId());
    Assert.assertEquals(NodeState.RUNNING, nr.getNodeState());

    allocateRequest2 =
        BuilderUtils.newAllocateRequest(
            attempt2.getAppAttemptId(), response2.getResponseId(), 0F, null, null);
    response2 = amService.allocate(allocateRequest2).getAMResponse();
    updatedNodes = response2.getUpdatedNodes();
    Assert.assertEquals(1, updatedNodes.size());
    nr = updatedNodes.iterator().next();
    Assert.assertEquals(nm4.getNodeId(), nr.getNodeId());
    Assert.assertEquals(NodeState.RUNNING, nr.getNodeState());

    // subsequent allocate calls should return no updated nodes
    allocateRequest2 =
        BuilderUtils.newAllocateRequest(
            attempt2.getAppAttemptId(), response2.getResponseId(), 0F, null, null);
    response2 = amService.allocate(allocateRequest2).getAMResponse();
    updatedNodes = response2.getUpdatedNodes();
    Assert.assertEquals(0, updatedNodes.size());

    // how to do the above for LOST node

  }