org.apache.hadoop.yarn.server.resourcemanager.rmcontainer.RMContainerImpl Java Examples

@SuppressWarnings("unchecked")
private Container allocateApplicationAttempt() {
  scheduleApplicationAttempt();
  
  // Mock the allocation of AM container 
  Container container = mock(Container.class);
  Resource resource = BuilderUtils.newResource(2048, 1);
  when(container.getId()).thenReturn(
      BuilderUtils.newContainerId(applicationAttempt.getAppAttemptId(), 1));
  when(container.getResource()).thenReturn(resource);
  Allocation allocation = mock(Allocation.class);
  when(allocation.getContainers()).
      thenReturn(Collections.singletonList(container));
  when(
      scheduler.allocate(
          any(ApplicationAttemptId.class), 
          any(List.class), 
          any(List.class), 
          any(List.class), 
          any(List.class))).
  thenReturn(allocation);
  RMContainer rmContainer = mock(RMContainerImpl.class);
  when(scheduler.getRMContainer(container.getId())).
      thenReturn(rmContainer);
  
  applicationAttempt.handle(
      new RMAppAttemptContainerAllocatedEvent(
          applicationAttempt.getAppAttemptId()));
  
  assertEquals(RMAppAttemptState.ALLOCATED_SAVING, 
      applicationAttempt.getAppAttemptState());
  applicationAttempt.handle(
      new RMAppAttemptEvent(applicationAttempt.getAppAttemptId(),
          RMAppAttemptEventType.ATTEMPT_NEW_SAVED));
  
  testAppAttemptAllocatedState(container);
  
  return container;
}
 

private RMContainer recoverAndCreateContainer(NMContainerStatus status,
    RMNode node) {
  Container container =
      Container.newInstance(status.getContainerId(), node.getNodeID(),
        node.getHttpAddress(), status.getAllocatedResource(),
        status.getPriority(), null);
  ApplicationAttemptId attemptId =
      container.getId().getApplicationAttemptId();
  RMContainer rmContainer =
      new RMContainerImpl(container, attemptId, node.getNodeID(),
        applications.get(attemptId.getApplicationId()).getUser(), rmContext,
        status.getCreationTime());
  return rmContainer;
}
 

@SuppressWarnings("unchecked")
private Container allocateApplicationAttempt() {
  scheduleApplicationAttempt();
  
  // Mock the allocation of AM container 
  Container container = mock(Container.class);
  Resource resource = BuilderUtils.newResource(2048, 1, 1);
  when(container.getId()).thenReturn(
      BuilderUtils.newContainerId(applicationAttempt.getAppAttemptId(), 1));
  when(container.getResource()).thenReturn(resource);
  Allocation allocation = mock(Allocation.class);
  when(allocation.getContainers()).
      thenReturn(Collections.singletonList(container));
  when(
      scheduler.allocate(
          any(ApplicationAttemptId.class), 
          any(List.class), 
          any(List.class), 
          any(List.class), 
          any(List.class))).
  thenReturn(allocation);
  RMContainer rmContainer = mock(RMContainerImpl.class);
  when(scheduler.getRMContainer(container.getId())).
      thenReturn(rmContainer);
  
  applicationAttempt.handle(
      new RMAppAttemptContainerAllocatedEvent(
          applicationAttempt.getAppAttemptId()));
  
  assertEquals(RMAppAttemptState.ALLOCATED_SAVING, 
      applicationAttempt.getAppAttemptState());
  applicationAttempt.handle(
      new RMAppAttemptEvent(applicationAttempt.getAppAttemptId(),
          RMAppAttemptEventType.ATTEMPT_NEW_SAVED));
  
  testAppAttemptAllocatedState(container);
  
  return container;
}
 

private RMContainer recoverAndCreateContainer(NMContainerStatus status,
    RMNode node) {
  Container container =
      Container.newInstance(status.getContainerId(), node.getNodeID(),
        node.getHttpAddress(), status.getAllocatedResource(),
        status.getPriority(), null);
  ApplicationAttemptId attemptId =
      container.getId().getApplicationAttemptId();
  RMContainer rmContainer =
      new RMContainerImpl(container, attemptId, node.getNodeID(),
        applications.get(attemptId.getApplicationId()).getUser(), rmContext,
        status.getCreationTime());
  return rmContainer;
}
 

@Test
@AgentAttachmentRule.Enforce
public void ContainerResourceMonitoringModule_should_attach_to_recordCpuUsage() throws ClassNotFoundException, NoSuchMethodException, IllegalAccessException, InstantiationException, InvocationTargetException, NoSuchFieldException {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));

    final Header[] header = new Header[1];
    final Object[] event = new Object[1];
    RMContainerTracer.initEventHandler((t, h, o) -> {
        header[0] = h;
        event[0] = o;
    });

    ClassFileTransformer classFileTransformer = new RMContainerTracer.RMContainerImplTracer().installOnByteBuddyAgent();

    try {
        Class<?> clazz = classLoader.loadClass(RMContainerImpl.class.getName());
        Constructor<?> constructor = clazz.getDeclaredConstructor(Container.class,
            ApplicationAttemptId.class, NodeId.class, String.class,
            RMContext.class);
        constructor.setAccessible(true);

        RMContext rmContext = mock(RMContext.class);
        Configuration yarnConf = new Configuration();
        when(rmContext.getYarnConfiguration())
            .thenReturn(yarnConf);
        Dispatcher dispatcher = mock(Dispatcher.class);
        EventHandler eventHandler = mock(EventHandler.class);
        when(dispatcher.getEventHandler())
            .thenReturn(eventHandler);
        when(rmContext.getDispatcher())
            .thenReturn(dispatcher);

        RMApplicationHistoryWriter rmApplicationHistoryWriter = mock(RMApplicationHistoryWriter.class);
        when(rmContext.getRMApplicationHistoryWriter())
            .thenReturn(rmApplicationHistoryWriter);

        SystemMetricsPublisher systemMetricsPublisher = mock(SystemMetricsPublisher.class);
        when(rmContext.getSystemMetricsPublisher())
            .thenReturn(systemMetricsPublisher);

        Object inFormat = constructor.newInstance(mock(Container.class), mock(ApplicationAttemptId.class), mock(NodeId.class), "user", rmContext);

        Method m = clazz.getDeclaredMethod("handle", RMContainerEvent.class);

        ContainerId cid = mock(ContainerId.class);
        ApplicationAttemptId applicationAttemptId = mock(ApplicationAttemptId.class);
        when(cid.toString())
            .thenReturn("cid");
        when(cid.getApplicationAttemptId())
            .thenReturn(applicationAttemptId);
        when(applicationAttemptId.toString())
            .thenReturn("appattempt_id");
        ApplicationId applicationId = mock(ApplicationId.class);
        when(applicationAttemptId.getApplicationId())
            .thenReturn(applicationId);
        when(applicationId.toString())
            .thenReturn("app_id");
        m.invoke(inFormat, new RMContainerEvent(cid, RMContainerEventType.START));

        assertNotNull(header[0]);
        assertNotNull(event[0]);

    } finally {
        ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer);
    }
}
 

private RMAppImpl getRMApp(RMContext rmContext, YarnScheduler yarnScheduler,
    ApplicationId applicationId3, YarnConfiguration config, String queueName,
    final long memorySeconds, final long vcoreSeconds) {
  ApplicationSubmissionContext asContext = mock(ApplicationSubmissionContext.class);
  when(asContext.getMaxAppAttempts()).thenReturn(1);

  RMAppImpl app =
      spy(new RMAppImpl(applicationId3, rmContext, config, null, null,
          queueName, asContext, yarnScheduler, null,
          System.currentTimeMillis(), "YARN", null,
          BuilderUtils.newResourceRequest(
              RMAppAttemptImpl.AM_CONTAINER_PRIORITY, ResourceRequest.ANY,
              Resource.newInstance(1024, 1), 1)){
                @Override
                public ApplicationReport createAndGetApplicationReport(
                    String clientUserName, boolean allowAccess) {
                  ApplicationReport report = super.createAndGetApplicationReport(
                      clientUserName, allowAccess);
                  ApplicationResourceUsageReport usageReport = 
                      report.getApplicationResourceUsageReport();
                  usageReport.setMemorySeconds(memorySeconds);
                  usageReport.setVcoreSeconds(vcoreSeconds);
                  report.setApplicationResourceUsageReport(usageReport);
                  return report;
                }
            });

  ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(
      ApplicationId.newInstance(123456, 1), 1);
  RMAppAttemptImpl rmAppAttemptImpl = spy(new RMAppAttemptImpl(attemptId,
      rmContext, yarnScheduler, null, asContext, config, false, null));
  Container container = Container.newInstance(
      ContainerId.newContainerId(attemptId, 1), null, "", null, null, null);
  RMContainerImpl containerimpl = spy(new RMContainerImpl(container,
      attemptId, null, "", rmContext));
  Map<ApplicationAttemptId, RMAppAttempt> attempts = 
    new HashMap<ApplicationAttemptId, RMAppAttempt>();
  attempts.put(attemptId, rmAppAttemptImpl);
  when(app.getCurrentAppAttempt()).thenReturn(rmAppAttemptImpl);
  when(app.getAppAttempts()).thenReturn(attempts);
  when(rmAppAttemptImpl.getMasterContainer()).thenReturn(container);
  ResourceScheduler rs = mock(ResourceScheduler.class);
  when(rmContext.getScheduler()).thenReturn(rs);
  when(rmContext.getScheduler().getRMContainer(any(ContainerId.class)))
      .thenReturn(containerimpl);
  SchedulerAppReport sAppReport = mock(SchedulerAppReport.class);
  when(
      rmContext.getScheduler().getSchedulerAppInfo(
          any(ApplicationAttemptId.class))).thenReturn(sAppReport);
  List<RMContainer> rmContainers = new ArrayList<RMContainer>();
  rmContainers.add(containerimpl);
  when(
      rmContext.getScheduler().getSchedulerAppInfo(attemptId)
          .getLiveContainers()).thenReturn(rmContainers);
  ContainerStatus cs = mock(ContainerStatus.class);
  when(containerimpl.getFinishedStatus()).thenReturn(cs);
  when(containerimpl.getDiagnosticsInfo()).thenReturn("N/A");
  when(containerimpl.getContainerExitStatus()).thenReturn(0);
  when(containerimpl.getContainerState()).thenReturn(ContainerState.COMPLETE);
  return app;
}
 

@Test
public void testFindNodeToUnreserve() throws Exception {

  CapacitySchedulerConfiguration csConf = new CapacitySchedulerConfiguration();
  setup(csConf);
  final String user_0 = "user_0";
  final ApplicationAttemptId appAttemptId_0 = TestUtils
      .getMockApplicationAttemptId(0, 0);
  LeafQueue a = stubLeafQueue((LeafQueue) queues.get(A));
  FiCaSchedulerApp app_0 = new FiCaSchedulerApp(appAttemptId_0, user_0, a,
      mock(ActiveUsersManager.class), spyRMContext);

  String host_1 = "host_1";
  FiCaSchedulerNode node_1 = TestUtils.getMockNode(host_1, DEFAULT_RACK, 0,
      8 * GB);

  // Setup resource-requests
  Priority priorityMap = TestUtils.createMockPriority(5);
  Resource capability = Resources.createResource(2 * GB, 0);

  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  SystemMetricsPublisher publisher = mock(SystemMetricsPublisher.class);
  RMContext rmContext = mock(RMContext.class);
  ContainerAllocationExpirer expirer =
    mock(ContainerAllocationExpirer.class);
  DrainDispatcher drainDispatcher = new DrainDispatcher();
  when(rmContext.getContainerAllocationExpirer()).thenReturn(expirer);
  when(rmContext.getDispatcher()).thenReturn(drainDispatcher);
  when(rmContext.getRMApplicationHistoryWriter()).thenReturn(writer);
  when(rmContext.getSystemMetricsPublisher()).thenReturn(publisher);
  ApplicationAttemptId appAttemptId = BuilderUtils.newApplicationAttemptId(
      app_0.getApplicationId(), 1);
  ContainerId containerId = BuilderUtils.newContainerId(appAttemptId, 1);
  Container container = TestUtils.getMockContainer(containerId,
      node_1.getNodeID(), Resources.createResource(2*GB), priorityMap);
  RMContainer rmContainer = new RMContainerImpl(container, appAttemptId,
      node_1.getNodeID(), "user", rmContext);

  // nothing reserved
  boolean res = a.findNodeToUnreserve(csContext.getClusterResource(),
      node_1, app_0, priorityMap, capability);
  assertFalse(res);

  // reserved but scheduler doesn't know about that node.
  app_0.reserve(node_1, priorityMap, rmContainer, container);
  node_1.reserveResource(app_0, priorityMap, rmContainer);
  res = a.findNodeToUnreserve(csContext.getClusterResource(), node_1, app_0,
      priorityMap, capability);
  assertFalse(res);
}
 

@Test
public void testGetAppToUnreserve() throws Exception {

  CapacitySchedulerConfiguration csConf = new CapacitySchedulerConfiguration();
  setup(csConf);
  final String user_0 = "user_0";
  final ApplicationAttemptId appAttemptId_0 = TestUtils
      .getMockApplicationAttemptId(0, 0);
  LeafQueue a = stubLeafQueue((LeafQueue) queues.get(A));
  FiCaSchedulerApp app_0 = new FiCaSchedulerApp(appAttemptId_0, user_0, a,
      mock(ActiveUsersManager.class), spyRMContext);

  String host_0 = "host_0";
  FiCaSchedulerNode node_0 = TestUtils.getMockNode(host_0, DEFAULT_RACK, 0,
      8 * GB);
  String host_1 = "host_1";
  FiCaSchedulerNode node_1 = TestUtils.getMockNode(host_1, DEFAULT_RACK, 0,
      8 * GB);
  
  Resource clusterResource = Resources.createResource(2 * 8 * GB);

  // Setup resource-requests
  Priority priorityMap = TestUtils.createMockPriority(5);
  Resource capability = Resources.createResource(2*GB, 0);

  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  SystemMetricsPublisher publisher = mock(SystemMetricsPublisher.class);
  RMContext rmContext = mock(RMContext.class);
  ContainerAllocationExpirer expirer =
    mock(ContainerAllocationExpirer.class);
  DrainDispatcher drainDispatcher = new DrainDispatcher();
  when(rmContext.getContainerAllocationExpirer()).thenReturn(expirer);
  when(rmContext.getDispatcher()).thenReturn(drainDispatcher);
  when(rmContext.getRMApplicationHistoryWriter()).thenReturn(writer);
  when(rmContext.getSystemMetricsPublisher()).thenReturn(publisher);
  ApplicationAttemptId appAttemptId = BuilderUtils.newApplicationAttemptId(
      app_0.getApplicationId(), 1);
  ContainerId containerId = BuilderUtils.newContainerId(appAttemptId, 1);
  Container container = TestUtils.getMockContainer(containerId,
      node_1.getNodeID(), Resources.createResource(2*GB), priorityMap);
  RMContainer rmContainer = new RMContainerImpl(container, appAttemptId,
      node_1.getNodeID(), "user", rmContext);

  Container container_1 = TestUtils.getMockContainer(containerId,
      node_0.getNodeID(), Resources.createResource(1*GB), priorityMap);
  RMContainer rmContainer_1 = new RMContainerImpl(container_1, appAttemptId,
      node_0.getNodeID(), "user", rmContext);

  // no reserved containers
  NodeId unreserveId =
      app_0.getNodeIdToUnreserve(priorityMap, capability,
          cs.getResourceCalculator(), clusterResource);
  assertEquals(null, unreserveId);

  // no reserved containers - reserve then unreserve
  app_0.reserve(node_0, priorityMap, rmContainer_1, container_1);
  app_0.unreserve(node_0, priorityMap);
  unreserveId = app_0.getNodeIdToUnreserve(priorityMap, capability,
      cs.getResourceCalculator(), clusterResource);
  assertEquals(null, unreserveId);

  // no container large enough is reserved
  app_0.reserve(node_0, priorityMap, rmContainer_1, container_1);
  unreserveId = app_0.getNodeIdToUnreserve(priorityMap, capability,
      cs.getResourceCalculator(), clusterResource);
  assertEquals(null, unreserveId);

  // reserve one that is now large enough
  app_0.reserve(node_1, priorityMap, rmContainer, container);
  unreserveId = app_0.getNodeIdToUnreserve(priorityMap, capability,
      cs.getResourceCalculator(), clusterResource);
  assertEquals(node_1.getNodeID(), unreserveId);
}
 

synchronized public RMContainer allocate(NodeType type, FiCaSchedulerNode node,
    Priority priority, ResourceRequest request, 
    Container container) {

  if (isStopped) {
    return null;
  }
  
  // Required sanity check - AM can call 'allocate' to update resource 
  // request without locking the scheduler, hence we need to check
  if (getTotalRequiredResources(priority) <= 0) {
    return null;
  }
  
  // Create RMContainer
  RMContainer rmContainer = new RMContainerImpl(container, this
      .getApplicationAttemptId(), node.getNodeID(),
      appSchedulingInfo.getUser(), this.rmContext);

  // Add it to allContainers list.
  newlyAllocatedContainers.add(rmContainer);
  liveContainers.put(container.getId(), rmContainer);    

  // Update consumption and track allocations
  List<ResourceRequest> resourceRequestList = appSchedulingInfo.allocate(
      type, node, priority, request, container);
  Resources.addTo(currentConsumption, container.getResource());
  
  // Update resource requests related to "request" and store in RMContainer 
  ((RMContainerImpl)rmContainer).setResourceRequests(resourceRequestList);

  // Inform the container
  rmContainer.handle(
      new RMContainerEvent(container.getId(), RMContainerEventType.START));

  if (LOG.isDebugEnabled()) {
    LOG.debug("allocate: applicationAttemptId=" 
        + container.getId().getApplicationAttemptId() 
        + " container=" + container.getId() + " host="
        + container.getNodeId().getHost() + " type=" + type);
  }
  RMAuditLogger.logSuccess(getUser(), 
      AuditConstants.ALLOC_CONTAINER, "SchedulerApp", 
      getApplicationId(), container.getId());
  
  return rmContainer;
}
 

public synchronized void recoverContainersOnNode(
    List<NMContainerStatus> containerReports, RMNode nm) {
  if (!rmContext.isWorkPreservingRecoveryEnabled()
      || containerReports == null
      || (containerReports != null && containerReports.isEmpty())) {
    return;
  }

  for (NMContainerStatus container : containerReports) {
    ApplicationId appId =
        container.getContainerId().getApplicationAttemptId().getApplicationId();
    RMApp rmApp = rmContext.getRMApps().get(appId);
    if (rmApp == null) {
      LOG.error("Skip recovering container " + container
          + " for unknown application.");
      killOrphanContainerOnNode(nm, container);
      continue;
    }

    // Unmanaged AM recovery is addressed in YARN-1815
    if (rmApp.getApplicationSubmissionContext().getUnmanagedAM()) {
      LOG.info("Skip recovering container " + container + " for unmanaged AM."
          + rmApp.getApplicationId());
      killOrphanContainerOnNode(nm, container);
      continue;
    }

    SchedulerApplication<T> schedulerApp = applications.get(appId);
    if (schedulerApp == null) {
      LOG.info("Skip recovering container  " + container
          + " for unknown SchedulerApplication. Application current state is "
          + rmApp.getState());
      killOrphanContainerOnNode(nm, container);
      continue;
    }

    LOG.info("Recovering container " + container);
    SchedulerApplicationAttempt schedulerAttempt =
        schedulerApp.getCurrentAppAttempt();

    if (!rmApp.getApplicationSubmissionContext()
      .getKeepContainersAcrossApplicationAttempts()) {
      // Do not recover containers for stopped attempt or previous attempt.
      if (schedulerAttempt.isStopped()
          || !schedulerAttempt.getApplicationAttemptId().equals(
            container.getContainerId().getApplicationAttemptId())) {
        LOG.info("Skip recovering container " + container
            + " for already stopped attempt.");
        killOrphanContainerOnNode(nm, container);
        continue;
      }
    }

    // create container
    RMContainer rmContainer = recoverAndCreateContainer(container, nm);

    // recover RMContainer
    rmContainer.handle(new RMContainerRecoverEvent(container.getContainerId(),
      container));

    // recover scheduler node
    nodes.get(nm.getNodeID()).recoverContainer(rmContainer);

    // recover queue: update headroom etc.
    Queue queue = schedulerAttempt.getQueue();
    queue.recoverContainer(clusterResource, schedulerAttempt, rmContainer);

    // recover scheduler attempt
    schedulerAttempt.recoverContainer(rmContainer);
          
    // set master container for the current running AMContainer for this
    // attempt.
    RMAppAttempt appAttempt = rmApp.getCurrentAppAttempt();
    if (appAttempt != null) {
      Container masterContainer = appAttempt.getMasterContainer();

      // Mark current running AMContainer's RMContainer based on the master
      // container ID stored in AppAttempt.
      if (masterContainer != null
          && masterContainer.getId().equals(rmContainer.getContainerId())) {
        ((RMContainerImpl)rmContainer).setAMContainer(true);
      }
    }

    synchronized (schedulerAttempt) {
      Set<ContainerId> releases = schedulerAttempt.getPendingRelease();
      if (releases.contains(container.getContainerId())) {
        // release the container
        rmContainer.handle(new RMContainerFinishedEvent(container
          .getContainerId(), SchedulerUtils.createAbnormalContainerStatus(
          container.getContainerId(), SchedulerUtils.RELEASED_CONTAINER),
          RMContainerEventType.RELEASED));
        releases.remove(container.getContainerId());
        LOG.info(container.getContainerId() + " is released by application.");
      }
    }
  }
}
 

synchronized public RMContainer allocate(NodeType type, FSSchedulerNode node,
    Priority priority, ResourceRequest request,
    Container container) {
  // Update allowed locality level
  NodeType allowed = allowedLocalityLevel.get(priority);
  if (allowed != null) {
    if (allowed.equals(NodeType.OFF_SWITCH) &&
        (type.equals(NodeType.NODE_LOCAL) ||
            type.equals(NodeType.RACK_LOCAL))) {
      this.resetAllowedLocalityLevel(priority, type);
    }
    else if (allowed.equals(NodeType.RACK_LOCAL) &&
        type.equals(NodeType.NODE_LOCAL)) {
      this.resetAllowedLocalityLevel(priority, type);
    }
  }

  // Required sanity check - AM can call 'allocate' to update resource 
  // request without locking the scheduler, hence we need to check
  if (getTotalRequiredResources(priority) <= 0) {
    return null;
  }
  
  // Create RMContainer
  RMContainer rmContainer = new RMContainerImpl(container, 
      getApplicationAttemptId(), node.getNodeID(),
      appSchedulingInfo.getUser(), rmContext);

  // Add it to allContainers list.
  newlyAllocatedContainers.add(rmContainer);
  liveContainers.put(container.getId(), rmContainer);    

  // Update consumption and track allocations
  List<ResourceRequest> resourceRequestList = appSchedulingInfo.allocate(
      type, node, priority, request, container);
  Resources.addTo(currentConsumption, container.getResource());

  // Update resource requests related to "request" and store in RMContainer
  ((RMContainerImpl) rmContainer).setResourceRequests(resourceRequestList);

  // Inform the container
  rmContainer.handle(
      new RMContainerEvent(container.getId(), RMContainerEventType.START));

  if (LOG.isDebugEnabled()) {
    LOG.debug("allocate: applicationAttemptId=" 
        + container.getId().getApplicationAttemptId() 
        + " container=" + container.getId() + " host="
        + container.getNodeId().getHost() + " type=" + type);
  }
  RMAuditLogger.logSuccess(getUser(), 
      AuditConstants.ALLOC_CONTAINER, "SchedulerApp", 
      getApplicationId(), container.getId());
  
  return rmContainer;
}
 

@Override
public RMAppAttemptState transition(RMAppAttemptImpl appAttempt,
    RMAppAttemptEvent event) {
  // Acquire the AM container from the scheduler.
  Allocation amContainerAllocation =
      appAttempt.scheduler.allocate(appAttempt.applicationAttemptId,
        EMPTY_CONTAINER_REQUEST_LIST, EMPTY_CONTAINER_RELEASE_LIST, null,
        null);
  // There must be at least one container allocated, because a
  // CONTAINER_ALLOCATED is emitted after an RMContainer is constructed,
  // and is put in SchedulerApplication#newlyAllocatedContainers.

  // Note that YarnScheduler#allocate is not guaranteed to be able to
  // fetch it since container may not be fetchable for some reason like
  // DNS unavailable causing container token not generated. As such, we
  // return to the previous state and keep retry until am container is
  // fetched.
  if (amContainerAllocation.getContainers().size() == 0) {
    appAttempt.retryFetchingAMContainer(appAttempt);
    return RMAppAttemptState.SCHEDULED;
  }

  // Set the masterContainer
  appAttempt.setMasterContainer(amContainerAllocation.getContainers()
      .get(0));
  RMContainerImpl rmMasterContainer = (RMContainerImpl)appAttempt.scheduler
      .getRMContainer(appAttempt.getMasterContainer().getId());
  rmMasterContainer.setAMContainer(true);
  // The node set in NMTokenSecrentManager is used for marking whether the
  // NMToken has been issued for this node to the AM.
  // When AM container was allocated to RM itself, the node which allocates
  // this AM container was marked as the NMToken already sent. Thus,
  // clear this node set so that the following allocate requests from AM are
  // able to retrieve the corresponding NMToken.
  appAttempt.rmContext.getNMTokenSecretManager()
    .clearNodeSetForAttempt(appAttempt.applicationAttemptId);
  appAttempt.getSubmissionContext().setResource(
    appAttempt.getMasterContainer().getResource());
  appAttempt.storeAttempt();
  return RMAppAttemptState.ALLOCATED_SAVING;
}
 

private RMAppImpl getRMApp(RMContext rmContext, YarnScheduler yarnScheduler,
    ApplicationId applicationId3, YarnConfiguration config, String queueName,
    final long memorySeconds, final long vcoreSeconds, final long gcoreSeconds) {
  ApplicationSubmissionContext asContext = mock(ApplicationSubmissionContext.class);
  when(asContext.getMaxAppAttempts()).thenReturn(1);

  RMAppImpl app =
      spy(new RMAppImpl(applicationId3, rmContext, config, null, null,
          queueName, asContext, yarnScheduler, null,
          System.currentTimeMillis(), "YARN", null,
          BuilderUtils.newResourceRequest(
              RMAppAttemptImpl.AM_CONTAINER_PRIORITY, ResourceRequest.ANY,
              Resource.newInstance(1024, 1, 1), 1)){
                @Override
                public ApplicationReport createAndGetApplicationReport(
                    String clientUserName, boolean allowAccess) {
                  ApplicationReport report = super.createAndGetApplicationReport(
                      clientUserName, allowAccess);
                  ApplicationResourceUsageReport usageReport = 
                      report.getApplicationResourceUsageReport();
                  usageReport.setMemorySeconds(memorySeconds);
                  usageReport.setVcoreSeconds(vcoreSeconds);
                  usageReport.setGcoreSeconds(gcoreSeconds);
                  report.setApplicationResourceUsageReport(usageReport);
                  return report;
                }
            });

  ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(
      ApplicationId.newInstance(123456, 1), 1);
  RMAppAttemptImpl rmAppAttemptImpl = spy(new RMAppAttemptImpl(attemptId,
      rmContext, yarnScheduler, null, asContext, config, false, null));
  Container container = Container.newInstance(
      ContainerId.newContainerId(attemptId, 1), null, "", null, null, null);
  RMContainerImpl containerimpl = spy(new RMContainerImpl(container,
      attemptId, null, "", rmContext));
  Map<ApplicationAttemptId, RMAppAttempt> attempts = 
    new HashMap<ApplicationAttemptId, RMAppAttempt>();
  attempts.put(attemptId, rmAppAttemptImpl);
  when(app.getCurrentAppAttempt()).thenReturn(rmAppAttemptImpl);
  when(app.getAppAttempts()).thenReturn(attempts);
  when(rmAppAttemptImpl.getMasterContainer()).thenReturn(container);
  ResourceScheduler rs = mock(ResourceScheduler.class);
  when(rmContext.getScheduler()).thenReturn(rs);
  when(rmContext.getScheduler().getRMContainer(any(ContainerId.class)))
      .thenReturn(containerimpl);
  SchedulerAppReport sAppReport = mock(SchedulerAppReport.class);
  when(
      rmContext.getScheduler().getSchedulerAppInfo(
          any(ApplicationAttemptId.class))).thenReturn(sAppReport);
  List<RMContainer> rmContainers = new ArrayList<RMContainer>();
  rmContainers.add(containerimpl);
  when(
      rmContext.getScheduler().getSchedulerAppInfo(attemptId)
          .getLiveContainers()).thenReturn(rmContainers);
  ContainerStatus cs = mock(ContainerStatus.class);
  when(containerimpl.getFinishedStatus()).thenReturn(cs);
  when(containerimpl.getDiagnosticsInfo()).thenReturn("N/A");
  when(containerimpl.getContainerExitStatus()).thenReturn(0);
  when(containerimpl.getContainerState()).thenReturn(ContainerState.COMPLETE);
  return app;
}
 

@Test
public void testFindNodeToUnreserve() throws Exception {

  CapacitySchedulerConfiguration csConf = new CapacitySchedulerConfiguration();
  setup(csConf);
  final String user_0 = "user_0";
  final ApplicationAttemptId appAttemptId_0 = TestUtils
      .getMockApplicationAttemptId(0, 0);
  LeafQueue a = stubLeafQueue((LeafQueue) queues.get(A));
  FiCaSchedulerApp app_0 = new FiCaSchedulerApp(appAttemptId_0, user_0, a,
      mock(ActiveUsersManager.class), spyRMContext);

  String host_1 = "host_1";
  FiCaSchedulerNode node_1 = TestUtils.getMockNode(host_1, DEFAULT_RACK, 0,
      8 * GB);

  // Setup resource-requests
  Priority priorityMap = TestUtils.createMockPriority(5);
  Resource capability = Resources.createResource(2 * GB, 0);

  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  SystemMetricsPublisher publisher = mock(SystemMetricsPublisher.class);
  RMContext rmContext = mock(RMContext.class);
  ContainerAllocationExpirer expirer =
    mock(ContainerAllocationExpirer.class);
  DrainDispatcher drainDispatcher = new DrainDispatcher();
  when(rmContext.getContainerAllocationExpirer()).thenReturn(expirer);
  when(rmContext.getDispatcher()).thenReturn(drainDispatcher);
  when(rmContext.getRMApplicationHistoryWriter()).thenReturn(writer);
  when(rmContext.getSystemMetricsPublisher()).thenReturn(publisher);
  ApplicationAttemptId appAttemptId = BuilderUtils.newApplicationAttemptId(
      app_0.getApplicationId(), 1);
  ContainerId containerId = BuilderUtils.newContainerId(appAttemptId, 1);
  Container container = TestUtils.getMockContainer(containerId,
      node_1.getNodeID(), Resources.createResource(2*GB), priorityMap);
  RMContainer rmContainer = new RMContainerImpl(container, appAttemptId,
      node_1.getNodeID(), "user", rmContext);

  // nothing reserved
  boolean res = a.findNodeToUnreserve(csContext.getClusterResource(),
      node_1, app_0, priorityMap, capability);
  assertFalse(res);

  // reserved but scheduler doesn't know about that node.
  app_0.reserve(node_1, priorityMap, rmContainer, container);
  node_1.reserveResource(app_0, priorityMap, rmContainer);
  res = a.findNodeToUnreserve(csContext.getClusterResource(), node_1, app_0,
      priorityMap, capability);
  assertFalse(res);
}
 

@Test
public void testGetAppToUnreserve() throws Exception {

  CapacitySchedulerConfiguration csConf = new CapacitySchedulerConfiguration();
  setup(csConf);
  final String user_0 = "user_0";
  final ApplicationAttemptId appAttemptId_0 = TestUtils
      .getMockApplicationAttemptId(0, 0);
  LeafQueue a = stubLeafQueue((LeafQueue) queues.get(A));
  FiCaSchedulerApp app_0 = new FiCaSchedulerApp(appAttemptId_0, user_0, a,
      mock(ActiveUsersManager.class), spyRMContext);

  String host_0 = "host_0";
  FiCaSchedulerNode node_0 = TestUtils.getMockNode(host_0, DEFAULT_RACK, 0,
      8 * GB);
  String host_1 = "host_1";
  FiCaSchedulerNode node_1 = TestUtils.getMockNode(host_1, DEFAULT_RACK, 0,
      8 * GB);
  
  Resource clusterResource = Resources.createResource(2 * 8 * GB);

  // Setup resource-requests
  Priority priorityMap = TestUtils.createMockPriority(5);
  Resource capability = Resources.createResource(2*GB, 0);

  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  SystemMetricsPublisher publisher = mock(SystemMetricsPublisher.class);
  RMContext rmContext = mock(RMContext.class);
  ContainerAllocationExpirer expirer =
    mock(ContainerAllocationExpirer.class);
  DrainDispatcher drainDispatcher = new DrainDispatcher();
  when(rmContext.getContainerAllocationExpirer()).thenReturn(expirer);
  when(rmContext.getDispatcher()).thenReturn(drainDispatcher);
  when(rmContext.getRMApplicationHistoryWriter()).thenReturn(writer);
  when(rmContext.getSystemMetricsPublisher()).thenReturn(publisher);
  ApplicationAttemptId appAttemptId = BuilderUtils.newApplicationAttemptId(
      app_0.getApplicationId(), 1);
  ContainerId containerId = BuilderUtils.newContainerId(appAttemptId, 1);
  Container container = TestUtils.getMockContainer(containerId,
      node_1.getNodeID(), Resources.createResource(2*GB), priorityMap);
  RMContainer rmContainer = new RMContainerImpl(container, appAttemptId,
      node_1.getNodeID(), "user", rmContext);

  Container container_1 = TestUtils.getMockContainer(containerId,
      node_0.getNodeID(), Resources.createResource(1*GB), priorityMap);
  RMContainer rmContainer_1 = new RMContainerImpl(container_1, appAttemptId,
      node_0.getNodeID(), "user", rmContext);

  // no reserved containers
  NodeId unreserveId =
      app_0.getNodeIdToUnreserve(priorityMap, capability,
          cs.getResourceCalculator(), clusterResource);
  assertEquals(null, unreserveId);

  // no reserved containers - reserve then unreserve
  app_0.reserve(node_0, priorityMap, rmContainer_1, container_1);
  app_0.unreserve(node_0, priorityMap);
  unreserveId = app_0.getNodeIdToUnreserve(priorityMap, capability,
      cs.getResourceCalculator(), clusterResource);
  assertEquals(null, unreserveId);

  // no container large enough is reserved
  app_0.reserve(node_0, priorityMap, rmContainer_1, container_1);
  unreserveId = app_0.getNodeIdToUnreserve(priorityMap, capability,
      cs.getResourceCalculator(), clusterResource);
  assertEquals(null, unreserveId);

  // reserve one that is now large enough
  app_0.reserve(node_1, priorityMap, rmContainer, container);
  unreserveId = app_0.getNodeIdToUnreserve(priorityMap, capability,
      cs.getResourceCalculator(), clusterResource);
  assertEquals(node_1.getNodeID(), unreserveId);
}
 

synchronized public RMContainer allocate(NodeType type, FiCaSchedulerNode node,
    Priority priority, ResourceRequest request, 
    Container container) {

  if (isStopped) {
    return null;
  }
  
  // Required sanity check - AM can call 'allocate' to update resource 
  // request without locking the scheduler, hence we need to check
  if (getTotalRequiredResources(priority) <= 0) {
    return null;
  }
  
  // Create RMContainer
  RMContainer rmContainer = new RMContainerImpl(container, this
      .getApplicationAttemptId(), node.getNodeID(),
      appSchedulingInfo.getUser(), this.rmContext);

  // Add it to allContainers list.
  newlyAllocatedContainers.add(rmContainer);
  liveContainers.put(container.getId(), rmContainer);    

  // Update consumption and track allocations
  List<ResourceRequest> resourceRequestList = appSchedulingInfo.allocate(
      type, node, priority, request, container);

  attemptResourceUsage.incUsed(node.getPartition(), container.getResource());

  // Update resource requests related to "request" and store in RMContainer
  ((RMContainerImpl)rmContainer).setResourceRequests(resourceRequestList);

  // Inform the container
  rmContainer.handle(
      new RMContainerEvent(container.getId(), RMContainerEventType.START));

  if (LOG.isDebugEnabled()) {
    LOG.debug("allocate: applicationAttemptId=" 
        + container.getId().getApplicationAttemptId() 
        + " container=" + container.getId() + " host="
        + container.getNodeId().getHost() + " type=" + type);
  }
  RMAuditLogger.logSuccess(getUser(),
      AuditConstants.ALLOC_CONTAINER, "SchedulerApp",
      getApplicationId(), container.getId());
  
  return rmContainer;
}
 

public synchronized void recoverContainersOnNode(
    List<NMContainerStatus> containerReports, RMNode nm) {
  if (!rmContext.isWorkPreservingRecoveryEnabled()
      || containerReports == null
      || (containerReports != null && containerReports.isEmpty())) {
    return;
  }

  for (NMContainerStatus container : containerReports) {
    ApplicationId appId =
        container.getContainerId().getApplicationAttemptId().getApplicationId();
    RMApp rmApp = rmContext.getRMApps().get(appId);
    if (rmApp == null) {
      LOG.error("Skip recovering container " + container
          + " for unknown application.");
      killOrphanContainerOnNode(nm, container);
      continue;
    }

    // Unmanaged AM recovery is addressed in YARN-1815
    if (rmApp.getApplicationSubmissionContext().getUnmanagedAM()) {
      LOG.info("Skip recovering container " + container + " for unmanaged AM."
          + rmApp.getApplicationId());
      killOrphanContainerOnNode(nm, container);
      continue;
    }

    SchedulerApplication<T> schedulerApp = applications.get(appId);
    if (schedulerApp == null) {
      LOG.info("Skip recovering container  " + container
          + " for unknown SchedulerApplication. Application current state is "
          + rmApp.getState());
      killOrphanContainerOnNode(nm, container);
      continue;
    }

    LOG.info("Recovering container " + container);
    SchedulerApplicationAttempt schedulerAttempt =
        schedulerApp.getCurrentAppAttempt();

    if (!rmApp.getApplicationSubmissionContext()
      .getKeepContainersAcrossApplicationAttempts()) {
      // Do not recover containers for stopped attempt or previous attempt.
      if (schedulerAttempt.isStopped()
          || !schedulerAttempt.getApplicationAttemptId().equals(
            container.getContainerId().getApplicationAttemptId())) {
        LOG.info("Skip recovering container " + container
            + " for already stopped attempt.");
        killOrphanContainerOnNode(nm, container);
        continue;
      }
    }

    // create container
    RMContainer rmContainer = recoverAndCreateContainer(container, nm);

    // recover RMContainer
    rmContainer.handle(new RMContainerRecoverEvent(container.getContainerId(),
      container));

    // recover scheduler node
    SchedulerNode schedulerNode = nodes.get(nm.getNodeID());
    schedulerNode.recoverContainer(rmContainer);

    // recover queue: update headroom etc.
    Queue queue = schedulerAttempt.getQueue();
    queue.recoverContainer(clusterResource, schedulerAttempt, rmContainer);

    // recover scheduler attempt
    schedulerAttempt.recoverContainer(schedulerNode, rmContainer);
          
    // set master container for the current running AMContainer for this
    // attempt.
    RMAppAttempt appAttempt = rmApp.getCurrentAppAttempt();
    if (appAttempt != null) {
      Container masterContainer = appAttempt.getMasterContainer();

      // Mark current running AMContainer's RMContainer based on the master
      // container ID stored in AppAttempt.
      if (masterContainer != null
          && masterContainer.getId().equals(rmContainer.getContainerId())) {
        ((RMContainerImpl)rmContainer).setAMContainer(true);
      }
    }

    synchronized (schedulerAttempt) {
      Set<ContainerId> releases = schedulerAttempt.getPendingRelease();
      if (releases.contains(container.getContainerId())) {
        // release the container
        rmContainer.handle(new RMContainerFinishedEvent(container
          .getContainerId(), SchedulerUtils.createAbnormalContainerStatus(
          container.getContainerId(), SchedulerUtils.RELEASED_CONTAINER),
          RMContainerEventType.RELEASED));
        releases.remove(container.getContainerId());
        LOG.info(container.getContainerId() + " is released by application.");
      }
    }
  }
}
 

synchronized public RMContainer allocate(NodeType type, FSSchedulerNode node,
    Priority priority, ResourceRequest request,
    Container container) {
  // Update allowed locality level
  NodeType allowed = allowedLocalityLevel.get(priority);
  if (allowed != null) {
    if (allowed.equals(NodeType.OFF_SWITCH) &&
        (type.equals(NodeType.NODE_LOCAL) ||
            type.equals(NodeType.RACK_LOCAL))) {
      this.resetAllowedLocalityLevel(priority, type);
    }
    else if (allowed.equals(NodeType.RACK_LOCAL) &&
        type.equals(NodeType.NODE_LOCAL)) {
      this.resetAllowedLocalityLevel(priority, type);
    }
  }

  // Required sanity check - AM can call 'allocate' to update resource 
  // request without locking the scheduler, hence we need to check
  if (getTotalRequiredResources(priority) <= 0) {
    return null;
  }
  
  // Create RMContainer
  RMContainer rmContainer = new RMContainerImpl(container, 
      getApplicationAttemptId(), node.getNodeID(),
      appSchedulingInfo.getUser(), rmContext);

  // Add it to allContainers list.
  newlyAllocatedContainers.add(rmContainer);
  liveContainers.put(container.getId(), rmContainer);    

  // Update consumption and track allocations
  List<ResourceRequest> resourceRequestList = appSchedulingInfo.allocate(
      type, node, priority, request, container);
  this.attemptResourceUsage.incUsed(container.getResource());

  // Update resource requests related to "request" and store in RMContainer
  ((RMContainerImpl) rmContainer).setResourceRequests(resourceRequestList);

  // Inform the container
  rmContainer.handle(
      new RMContainerEvent(container.getId(), RMContainerEventType.START));

  if (LOG.isDebugEnabled()) {
    LOG.debug("allocate: applicationAttemptId=" 
        + container.getId().getApplicationAttemptId() 
        + " container=" + container.getId() + " host="
        + container.getNodeId().getHost() + " type=" + type);
  }
  RMAuditLogger.logSuccess(getUser(), 
      AuditConstants.ALLOC_CONTAINER, "SchedulerApp", 
      getApplicationId(), container.getId());
  
  return rmContainer;
}
 

@Override
public RMAppAttemptState transition(RMAppAttemptImpl appAttempt,
    RMAppAttemptEvent event) {
  // Acquire the AM container from the scheduler.
  Allocation amContainerAllocation =
      appAttempt.scheduler.allocate(appAttempt.applicationAttemptId,
        EMPTY_CONTAINER_REQUEST_LIST, EMPTY_CONTAINER_RELEASE_LIST, null,
        null);
  // There must be at least one container allocated, because a
  // CONTAINER_ALLOCATED is emitted after an RMContainer is constructed,
  // and is put in SchedulerApplication#newlyAllocatedContainers.

  // Note that YarnScheduler#allocate is not guaranteed to be able to
  // fetch it since container may not be fetchable for some reason like
  // DNS unavailable causing container token not generated. As such, we
  // return to the previous state and keep retry until am container is
  // fetched.
  if (amContainerAllocation.getContainers().size() == 0) {
    appAttempt.retryFetchingAMContainer(appAttempt);
    return RMAppAttemptState.SCHEDULED;
  }

  // Set the masterContainer
  appAttempt.setMasterContainer(amContainerAllocation.getContainers()
      .get(0));
  RMContainerImpl rmMasterContainer = (RMContainerImpl)appAttempt.scheduler
      .getRMContainer(appAttempt.getMasterContainer().getId());
  rmMasterContainer.setAMContainer(true);
  // The node set in NMTokenSecrentManager is used for marking whether the
  // NMToken has been issued for this node to the AM.
  // When AM container was allocated to RM itself, the node which allocates
  // this AM container was marked as the NMToken already sent. Thus,
  // clear this node set so that the following allocate requests from AM are
  // able to retrieve the corresponding NMToken.
  appAttempt.rmContext.getNMTokenSecretManager()
    .clearNodeSetForAttempt(appAttempt.applicationAttemptId);
  appAttempt.getSubmissionContext().setResource(
    appAttempt.getMasterContainer().getResource());
  appAttempt.storeAttempt();
  return RMAppAttemptState.ALLOCATED_SAVING;
}
 

public static void intercept(@This Object rmContainerImpl,
                             @Argument(0) RMContainerEvent rmContainerEvent) {
    try {
        ContainerId cID = rmContainerEvent.getContainerId();
        ApplicationAttemptId applicationAttemptId = cID.getApplicationAttemptId();
        String applicationId = applicationAttemptId.getApplicationId().toString();
        String attemptId = applicationAttemptId.toString();

        Header header = Header.newBuilder()
            .withId(applicationId)
            .withApplicationID(applicationId)
            .withAttemptID(attemptId)
            .withContainerID(cID.toString())
            .build();

        RMContainerImpl rmc = (RMContainerImpl) rmContainerImpl;

        ResourceManagerEventProtos.ContainerEvent.Builder eventBuilder = ResourceManagerEventProtos.ContainerEvent.newBuilder()
            .setType(rmContainerEvent.getType().name())
            .setState(rmc.getState().name())
            .setStartTime(rmc.getCreationTime())
            .setLogUrl(rmc.getLogURL());

        if (rmc.getContainer() != null && rmc.getContainer().getResource() != null) {
            eventBuilder.setVcoresReserved(rmc.getContainer().getResource().getVirtualCores());
            eventBuilder.setMemoryReserved(rmc.getContainer().getResource().getMemory());
        }

        if (rmc.getAllocatedNode() != null) {
            eventBuilder.setContainerHostname(rmc.getAllocatedNode().getHost());
        }

        ContainerStatus containerStatus = (ContainerStatus) getField().get(rmContainerImpl);
        if (containerStatus != null) {
            eventBuilder
                .setIsFinished(true)
                .setExitStatus(rmc.getContainerExitStatus())
                .setReason(rmc.getDiagnosticsInfo())
                .setFinishTime(rmc.getFinishTime());
        }

        eventHandler.accept(System.currentTimeMillis(), header, eventBuilder.build());
    } catch (Throwable ignored) {
    }
}
 

/**
 * Returns a new RMContainer corresponding to the RMNode and RMContext. The RMContainer is the 
 * ResourceManager's view of an application container per the Hadoop docs
 * 
 * @param node
 * @param context
 * @param appId
 * @param cores
 * @param memory
 * @return RMContainer
 */
public static RMContainer getRMContainer(RMNode node, RMContext context, int appId, int cores, int memory) {
  ContainerId containerId = ContainerId.newContainerId(ApplicationAttemptId.newInstance(
      ApplicationId.newInstance(123456789, 1), 1), appId);

  Container container = Container.newInstance(containerId, node.getNodeID(), node.getHttpAddress(),
      Resources.createResource(memory, cores), null, null);
  return new RMContainerImpl(container, containerId.getApplicationAttemptId(), node.getNodeID(), "user1", context);
}