org.apache.hadoop.yarn.client.api.AMRMClient Java Examples

@Test (timeout = 10000)
public void testAMRMClientAsyncShutDown() throws Exception {
  Configuration conf = new Configuration();
  TestCallbackHandler callbackHandler = new TestCallbackHandler();
  @SuppressWarnings("unchecked")
  AMRMClient<ContainerRequest> client = mock(AMRMClientImpl.class);

  createAllocateResponse(new ArrayList<ContainerStatus>(),
    new ArrayList<Container>(), null);
  when(client.allocate(anyFloat())).thenThrow(
    new ApplicationAttemptNotFoundException("app not found, shut down"));

  AMRMClientAsync<ContainerRequest> asyncClient =
      AMRMClientAsync.createAMRMClientAsync(client, 10, callbackHandler);
  asyncClient.init(conf);
  asyncClient.start();

  asyncClient.registerApplicationMaster("localhost", 1234, null);

  Thread.sleep(50);

  verify(client, times(1)).allocate(anyFloat());
  asyncClient.stop();
}
 

private Collection<AMRMClient.ContainerRequest> getPendingRequests() {
	final List<? extends Collection<AMRMClient.ContainerRequest>> matchingRequests = resourceManagerClient.getMatchingRequests(
		RM_REQUEST_PRIORITY,
		ResourceRequest.ANY,
		getContainerResource());

	final Collection<AMRMClient.ContainerRequest> matchingContainerRequests;

	if (matchingRequests.isEmpty()) {
		matchingContainerRequests = Collections.emptyList();
	} else {
		final Collection<AMRMClient.ContainerRequest> collection = matchingRequests.get(0);
		matchingContainerRequests = new ArrayList<>(collection);
	}

	Preconditions.checkState(
		matchingContainerRequests.size() == numPendingContainerRequests,
		"The RMClient's and YarnResourceManagers internal state about the number of pending container requests has diverged. Number client's pending container requests %s != Number RM's pending container requests %s.", matchingContainerRequests.size(), numPendingContainerRequests);

	return matchingContainerRequests;
}
 

/**
 * Issue requests to AM RM Client again if previous container requests expired and were not allocated by Yarn
 * @param amRmClient
 * @param requestedResources
 * @param loopCounter
 * @param resourceRequestor
 * @param containerRequests
 * @param removedContainerRequests
 */
public void reissueContainerRequests(AMRMClient<ContainerRequest> amRmClient, Map<StreamingContainerAgent.ContainerStartRequest, MutablePair<Integer, ContainerRequest>> requestedResources, int loopCounter, ResourceRequestHandler resourceRequestor, List<ContainerRequest> containerRequests, List<ContainerRequest> removedContainerRequests)
{
  if (!requestedResources.isEmpty()) {
    for (Map.Entry<StreamingContainerAgent.ContainerStartRequest, MutablePair<Integer, ContainerRequest>> entry : requestedResources.entrySet()) {
      /*
       * Create container requests again if pending requests were not allocated by Yarn till timeout.
       */
      if ((loopCounter - entry.getValue().getKey()) > NUMBER_MISSED_HEARTBEATS) {
        StreamingContainerAgent.ContainerStartRequest csr = entry.getKey();
        LOG.debug("Request for container {} timed out. Re-requesting container", csr.container);
        removedContainerRequests.add(entry.getValue().getRight());
        ContainerRequest cr = resourceRequestor.createContainerRequest(csr, false);
        entry.getValue().setLeft(loopCounter);
        entry.getValue().setRight(cr);
        containerRequests.add(cr);
      }
    }
  }
}
 

private Collection<AMRMClient.ContainerRequest> getPendingRequests() {
	final List<? extends Collection<AMRMClient.ContainerRequest>> matchingRequests = resourceManagerClient.getMatchingRequests(
		RM_REQUEST_PRIORITY,
		ResourceRequest.ANY,
		getContainerResource());

	final Collection<AMRMClient.ContainerRequest> matchingContainerRequests;

	if (matchingRequests.isEmpty()) {
		matchingContainerRequests = Collections.emptyList();
	} else {
		final Collection<AMRMClient.ContainerRequest> collection = matchingRequests.get(0);
		matchingContainerRequests = new ArrayList<>(collection);
	}

	Preconditions.checkState(
		matchingContainerRequests.size() == numPendingContainerRequests,
		"The RMClient's and YarnResourceManagers internal state about the number of pending container requests has diverged. Number client's pending container requests %s != Number RM's pending container requests %s.", matchingContainerRequests.size(), numPendingContainerRequests);

	return matchingContainerRequests;
}
 

/**
 * @throws Exception If failed.
 */
@Test
public void testClusterResource() throws Exception {
    rmMock.availableRes(new MockResource(1024, 2));

    appMaster.setRmClient(rmMock);
    appMaster.setNmClient(new NMMock());

    props.cpusPerNode(8);
    props.memoryPerNode(10240);
    props.instances(3);

    Thread thread = runAppMaster(appMaster);

    List<AMRMClient.ContainerRequest> contRequests = collectRequests(rmMock, 1, 1000);

    interruptedThread(thread);

    assertEquals(0, contRequests.size());
}
 

/**
 * Tests whether memory overhead prevents from allocating container.
 *
 * @throws Exception If failed.
 */
@Test
public void testMemoryOverHeadPreventAllocation() throws Exception {
    rmMock.availableRes(new MockResource(1024, 2));
    appMaster.setRmClient(rmMock);
    appMaster.setNmClient(new NMMock());

    props.cpusPerNode(2);
    props.memoryPerNode(1024);
    props.memoryOverHeadPerNode(512);
    props.instances(3);

    Thread thread = runAppMaster(appMaster);

    List<AMRMClient.ContainerRequest> contRequests = collectRequests(rmMock, 1, 1000);

    interruptedThread(thread);

    assertEquals(0, contRequests.size());
 }
 

@Test
public void testWaitFor() throws InterruptedException {
  AMRMClientImpl<ContainerRequest> amClient = null;
  CountDownSupplier countDownChecker = new CountDownSupplier();

  try {
    // start am rm client
    amClient =
        (AMRMClientImpl<ContainerRequest>) AMRMClient
            .<ContainerRequest> createAMRMClient();
    amClient.init(new YarnConfiguration());
    amClient.start();
    amClient.waitFor(countDownChecker, 1000);
    assertEquals(3, countDownChecker.counter);
  } finally {
    if (amClient != null) {
      amClient.stop();
    }
  }
}
 

/**
 * Tests whether memory overhead is allocated within container memory.
 *
 * @throws Exception If failed.
 */
@Test
public void testMemoryOverHeadAllocation() throws Exception {
    appMaster.setRmClient(rmMock);
    appMaster.setNmClient(new NMMock());

    props.cpusPerNode(2);
    props.memoryPerNode(1024);
    props.memoryOverHeadPerNode(512);
    props.instances(3);

    Thread thread = runAppMaster(appMaster);

    List<AMRMClient.ContainerRequest> contRequests = collectRequests(rmMock, 1, 1000);

    interruptedThread(thread);

    assertEquals(3, contRequests.size());

    for (AMRMClient.ContainerRequest req : contRequests) {
        assertEquals(2, req.getCapability().getVirtualCores());
        assertEquals(1024 + 512, req.getCapability().getMemory());
    }
}
 

private void requestContainer(Optional<String> preferredNode) {
  Priority priority = Records.newRecord(Priority.class);
  priority.setPriority(0);

  Resource capability = Records.newRecord(Resource.class);
  int maxMemoryCapacity = this.maxResourceCapacity.get().getMemory();
  capability.setMemory(this.requestedContainerMemoryMbs <= maxMemoryCapacity ?
      this.requestedContainerMemoryMbs : maxMemoryCapacity);
  int maxCoreCapacity = this.maxResourceCapacity.get().getVirtualCores();
  capability.setVirtualCores(this.requestedContainerCores <= maxCoreCapacity ?
      this.requestedContainerCores : maxCoreCapacity);

  String[] preferredNodes = preferredNode.isPresent() ? new String[] {preferredNode.get()} : null;
  this.amrmClientAsync.addContainerRequest(
      new AMRMClient.ContainerRequest(capability, preferredNodes, null, priority));
}
 

@Override
public void addSpecSupervisor(int number, int container_memory, int container_vcorecount, List<String> racks, List<String> hosts) throws TException {

    LOG.info("number:" + number + "; memory:" + container_memory + "; vcore:" + container_vcorecount + "; racks:" + JstormYarnUtils.join(racks, ",", false));
    LOG.info("hosts:" + JstormYarnUtils.join(hosts, ",", false));

    if (container_memory > jstormMasterContext.maxMemory) {
        container_vcorecount = jstormMasterContext.maxMemory;
    }
    if (container_vcorecount > jstormMasterContext.maxVcores) {
        container_vcorecount = jstormMasterContext.maxVcores;
    }
    for (int i = 0; i < number; i++) {
        AMRMClient.ContainerRequest containerAsk = jstormMaster.setupContainerAskForRM(container_memory, container_vcorecount, 0, racks.toArray(new String[0]), hosts.toArray(new String[0]));
        try {
            requestQueue.put(containerAsk);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        amRMClient.addContainerRequest(containerAsk);
        jstormMasterContext.numRequestedContainers.getAndIncrement();
    }
}
 

/**
 *
 * Remove a previously submitted resource request. The previous container request may have
 * been submitted. Even after the remove request, a Callback implementation must
 * be prepared to receive an allocation for the previous request. This is merely a best effort cancellation.
 *
 * @param request the request to be cancelled
 */
@Override
public void cancelResourceRequest(SamzaResourceRequest request) {
  String processorId = request.getProcessorId();
  String preferredHost = request.getPreferredHost();
  String requestId = request.getRequestId();
  log.info("Cancelling resource request for Processor ID: {} on host: {} with Request ID: {}",
      processorId, preferredHost, requestId);
  //ensure that removal and cancellation are done atomically.
  synchronized (lock) {
    AMRMClient.ContainerRequest containerRequest = requestsMap.get(request);
    if (containerRequest == null) {
      log.info("Resource request for Processor ID: {} on host: {} with Request ID: {} already cancelled.",
          processorId, preferredHost, requestId);
      return;
    }
    requestsMap.remove(request);
    amClient.removeContainerRequest(containerRequest);
  }
}
 

@Test(timeout=10000)
public void testMatchingRequestsForTopPriority() {
  String[] hosts = { "host1" };
  String[] racks = { "rack1" };
  AMRMClient.ContainerRequest req1 = new AMRMClient.ContainerRequest(
    Resource.newInstance(2048, 1), hosts, racks,
    Priority.newInstance(1));
  AMRMClient.ContainerRequest req2 = new AMRMClient.ContainerRequest(
    Resource.newInstance(1024, 1), hosts, racks,
    Priority.newInstance(2));
  AMRMClient.ContainerRequest req3 = new AMRMClient.ContainerRequest(
    Resource.newInstance(1024, 1), hosts, racks,
    Priority.newInstance(3));
  amrmClient.addContainerRequest(req1);
  amrmClient.addContainerRequest(req2);
  amrmClient.addContainerRequest(req3);

  Assert.assertTrue(amrmClient.getMatchingRequestsForTopPriority("host1",
    Resource.newInstance(1024, 1)).isEmpty());

  List<? extends Collection<AMRMClient.ContainerRequest>> ret =
    amrmClient.getMatchingRequestsForTopPriority("host1",
      Resource.newInstance(2048, 1));
  Assert.assertFalse(ret.isEmpty());
  Assert.assertEquals(req1, ret.get(0).iterator().next());

  amrmClient.removeContainerRequest(req1);

  ret = amrmClient.getMatchingRequestsForTopPriority("host1",
      Resource.newInstance(1024, 1));
  Assert.assertFalse(ret.isEmpty());
  Assert.assertEquals(req2, ret.get(0).iterator().next());
}
 

@Private
@VisibleForTesting
public AMRMClientAsyncImpl(AMRMClient<T> client, int intervalMs,
    CallbackHandler callbackHandler) {
  super(client, intervalMs, callbackHandler);
  heartbeatThread = new HeartbeatThread();
  handlerThread = new CallbackHandlerThread();
  responseQueue = new LinkedBlockingQueue<AllocateResponse>();
  keepRunning = true;
  savedException = null;
}
 

@Private
@VisibleForTesting
protected AMRMClientAsync(AMRMClient<T> client, int intervalMs,
    CallbackHandler callbackHandler) {
  super(AMRMClientAsync.class.getName());
  this.client = client;
  this.heartbeatIntervalMs.set(intervalMs);
  this.handler = callbackHandler;
}
 

private void runNonKerberized(Configuration conf) throws Exception {
    // Original user has the YARN tokens
    UserGroupInformation original = UserGroupInformation.getCurrentUser();

    String user = System.getProperty("splice.spark.yarn.user", "hbase");
    LOG.info("Login with user");
    ugi = UserGroupInformation.createRemoteUser(user);
    Collection<Token<? extends TokenIdentifier>> tokens = UserGroupInformation.getCurrentUser().getCredentials().getAllTokens();
    for (Token<? extends TokenIdentifier> token : tokens) {
        LOG.debug("Token kind is " + token.getKind().toString()
                + " and the token's service name is " + token.getService());
        if (AMRMTokenIdentifier.KIND_NAME.equals(token.getKind())) {
            ugi.addToken(token);
        }
    }

    // Transfer tokens from original user to the one we'll use from now on
    SparkHadoopUtil.get().transferCredentials(original, ugi);

    UserGroupInformation.isSecurityEnabled();
    if (mode == Mode.YARN) {
        rmClient = ugi.doAs(new PrivilegedExceptionAction<AMRMClientAsync<AMRMClient.ContainerRequest>>() {
            @Override
            public AMRMClientAsync<AMRMClient.ContainerRequest> run() throws Exception {
                return initClient(conf);
            }
        });
        LOG.info("Registered with Resource Manager");
    }
}
 

@SuppressWarnings("unchecked")
@Test(timeout=10000)
public void testMatchingRequestsForTopPriority() {
  String[] hosts = { "host1" };
  String[] racks = { "rack1" };
  AMRMClient.ContainerRequest req1 = new AMRMClient.ContainerRequest(
    Resource.newInstance(2048, 1), hosts, racks,
    Priority.newInstance(1));
  AMRMClient.ContainerRequest req2 = new AMRMClient.ContainerRequest(
    Resource.newInstance(1024, 1), hosts, racks,
    Priority.newInstance(2));
  AMRMClient.ContainerRequest req3 = new AMRMClient.ContainerRequest(
    Resource.newInstance(1024, 1), hosts, racks,
    Priority.newInstance(3));
  amrmClient.addContainerRequest(req1);
  amrmClient.addContainerRequest(req2);
  amrmClient.addContainerRequest(req3);

  Assert.assertTrue(amrmClient.getMatchingRequestsForTopPriority("host1",
    Resource.newInstance(1024, 1)).isEmpty());

  List<? extends Collection<AMRMClient.ContainerRequest>> ret =
    amrmClient.getMatchingRequestsForTopPriority("host1",
      Resource.newInstance(2048, 1));
  Assert.assertFalse(ret.isEmpty());
  Assert.assertEquals(req1, ret.get(0).iterator().next());

  amrmClient.removeContainerRequest(req1);

  ret = amrmClient.getMatchingRequestsForTopPriority("host1",
      Resource.newInstance(1024, 1));
  Assert.assertFalse(ret.isEmpty());
  Assert.assertEquals(req2, ret.get(0).iterator().next());
}
 

private void runHeartBeatThrowOutException(Exception ex) throws Exception{
  Configuration conf = new Configuration();
  TestCallbackHandler callbackHandler = new TestCallbackHandler();
  @SuppressWarnings("unchecked")
  AMRMClient<ContainerRequest> client = mock(AMRMClientImpl.class);
  when(client.allocate(anyFloat())).thenThrow(ex);

  AMRMClientAsync<ContainerRequest> asyncClient =
      AMRMClientAsync.createAMRMClientAsync(client, 20, callbackHandler);
  asyncClient.init(conf);
  asyncClient.start();
  
  synchronized (callbackHandler.notifier) {
    asyncClient.registerApplicationMaster("localhost", 1234, null);
    while(callbackHandler.savedException == null) {
      try {
        callbackHandler.notifier.wait();
      } catch (InterruptedException e) {
        e.printStackTrace();
      }
    }
  }
  Assert.assertTrue(callbackHandler.savedException.getMessage().contains(
      ex.getMessage()));
  
  asyncClient.stop();
  // stopping should have joined all threads and completed all callbacks
  Assert.assertTrue(callbackHandler.callbackCount == 0);
}
 

private AMRMClientAsync<AMRMClient.ContainerRequest> initClient(Configuration conf) throws YarnException, IOException {
    AMRMClientAsync.CallbackHandler allocListener = new AMRMClientAsync.CallbackHandler() {
        @Override
        public void onContainersCompleted(List<ContainerStatus> statuses) {
        }

        @Override
        public void onContainersAllocated(List<Container> containers) {
        }

        @Override
        public void onShutdownRequest() {
            LOG.warn("Shutting down");
            end.set(true);
        }

        @Override
        public void onNodesUpdated(List<NodeReport> updatedNodes) {
        }

        @Override
        public float getProgress() {
            return 0;
        }

        @Override
        public void onError(Throwable e) {
            LOG.error("Unexpected error", e);
            end.set(true);
        }
    };
    AMRMClientAsync<AMRMClient.ContainerRequest> rmClient = AMRMClientAsync.createAMRMClientAsync(1000, allocListener);
    rmClient.init(conf);
    rmClient.start();

    // Register with ResourceManager
    rmClient.registerApplicationMaster(Utils.localHostName(), 0, "");

    return rmClient;
}
 

@Test (timeout = 10000)
public void testAMRMClientAsyncShutDownWithWaitFor() throws Exception {
  Configuration conf = new Configuration();
  final TestCallbackHandler callbackHandler = new TestCallbackHandler();
  @SuppressWarnings("unchecked")
  AMRMClient<ContainerRequest> client = mock(AMRMClientImpl.class);
  when(client.allocate(anyFloat())).thenThrow(
    new ApplicationAttemptNotFoundException("app not found, shut down"));

  AMRMClientAsync<ContainerRequest> asyncClient =
      AMRMClientAsync.createAMRMClientAsync(client, 10, callbackHandler);
  asyncClient.init(conf);
  asyncClient.start();

  Supplier<Boolean> checker = new Supplier<Boolean>() {
    @Override
    public Boolean get() {
      return callbackHandler.reboot;
    }
  };

  asyncClient.registerApplicationMaster("localhost", 1234, null);
  asyncClient.waitFor(checker);

  asyncClient.stop();
  // stopping should have joined all threads and completed all callbacks
  Assert.assertTrue(callbackHandler.callbackCount == 0);

  verify(client, times(1)).allocate(anyFloat());
  asyncClient.stop();
}
 

@Test (timeout = 5000)
public void testCallAMRMClientAsyncStopFromCallbackHandler()
    throws YarnException, IOException, InterruptedException {
  Configuration conf = new Configuration();
  TestCallbackHandler2 callbackHandler = new TestCallbackHandler2();
  @SuppressWarnings("unchecked")
  AMRMClient<ContainerRequest> client = mock(AMRMClientImpl.class);

  List<ContainerStatus> completed = Arrays.asList(
      ContainerStatus.newInstance(newContainerId(0, 0, 0, 0),
          ContainerState.COMPLETE, "", 0));
  final AllocateResponse response = createAllocateResponse(completed,
      new ArrayList<Container>(), null);

  when(client.allocate(anyFloat())).thenReturn(response);

  AMRMClientAsync<ContainerRequest> asyncClient =
      AMRMClientAsync.createAMRMClientAsync(client, 20, callbackHandler);
  callbackHandler.asynClient = asyncClient;
  asyncClient.init(conf);
  asyncClient.start();

  synchronized (callbackHandler.notifier) {
    asyncClient.registerApplicationMaster("localhost", 1234, null);
    while(callbackHandler.notify == false) {
      try {
        callbackHandler.notifier.wait();
      } catch (InterruptedException e) {
        e.printStackTrace();
      }
    }
  }
}
 

@Test (timeout=60000)
public void testAMRMClient() throws YarnException, IOException {
  AMRMClient<ContainerRequest> amClient = null;
  try {
    // start am rm client
    amClient = AMRMClient.<ContainerRequest>createAMRMClient();

    //setting an instance NMTokenCache
    amClient.setNMTokenCache(new NMTokenCache());
    //asserting we are not using the singleton instance cache
    Assert.assertNotSame(NMTokenCache.getSingleton(), 
        amClient.getNMTokenCache());

    amClient.init(conf);
    amClient.start();

    amClient.registerApplicationMaster("Host", 10000, "");

    testAllocation((AMRMClientImpl<ContainerRequest>)amClient);

    amClient.unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED,
        null, null);

  } finally {
    if (amClient != null && amClient.getServiceState() == STATE.STARTED) {
      amClient.stop();
    }
  }
}
 

@Override
protected AMRMClientAsync<AMRMClient.ContainerRequest> createAndStartResourceManagerClient(
		YarnConfiguration yarnConfiguration,
		int yarnHeartbeatIntervalMillis,
		@Nullable String webInterfaceUrl) {
	return testingYarnAMRMClientAsync;
}
 

public Hadoop21YarnAMClient(Configuration conf) {
  super(ApplicationConstants.Environment.CONTAINER_ID.name());

  this.amrmClient = AMRMClient.createAMRMClient();
  this.amrmClient.init(conf);
  this.nmClient = new Hadoop21YarnNMClient(conf);
}
 

@Test (timeout = 5000)
public void testCallAMRMClientAsyncStopFromCallbackHandlerWithWaitFor()
    throws YarnException, IOException, InterruptedException {
  Configuration conf = new Configuration();
  final TestCallbackHandler2 callbackHandler = new TestCallbackHandler2();
  @SuppressWarnings("unchecked")
  AMRMClient<ContainerRequest> client = mock(AMRMClientImpl.class);

  List<ContainerStatus> completed = Arrays.asList(
      ContainerStatus.newInstance(newContainerId(0, 0, 0, 0),
          ContainerState.COMPLETE, "", 0));
  final AllocateResponse response = createAllocateResponse(completed,
      new ArrayList<Container>(), null);

  when(client.allocate(anyFloat())).thenReturn(response);

  AMRMClientAsync<ContainerRequest> asyncClient =
      AMRMClientAsync.createAMRMClientAsync(client, 20, callbackHandler);
  callbackHandler.asynClient = asyncClient;
  asyncClient.init(conf);
  asyncClient.start();

  Supplier<Boolean> checker = new Supplier<Boolean>() {
    @Override
    public Boolean get() {
      return callbackHandler.notify;
    }
  };

  asyncClient.registerApplicationMaster("localhost", 1234, null);
  asyncClient.waitFor(checker);
  Assert.assertTrue(checker.get());
}
 

@Test (timeout = 5000)
public void testCallAMRMClientAsyncStopFromCallbackHandler()
    throws YarnException, IOException, InterruptedException {
  Configuration conf = new Configuration();
  TestCallbackHandler2 callbackHandler = new TestCallbackHandler2();
  @SuppressWarnings("unchecked")
  AMRMClient<ContainerRequest> client = mock(AMRMClientImpl.class);

  List<ContainerStatus> completed = Arrays.asList(
      ContainerStatus.newInstance(newContainerId(0, 0, 0, 0),
          ContainerState.COMPLETE, "", 0));
  final AllocateResponse response = createAllocateResponse(completed,
      new ArrayList<Container>(), null);

  when(client.allocate(anyFloat())).thenReturn(response);

  AMRMClientAsync<ContainerRequest> asyncClient =
      AMRMClientAsync.createAMRMClientAsync(client, 20, callbackHandler);
  callbackHandler.asynClient = asyncClient;
  asyncClient.init(conf);
  asyncClient.start();

  synchronized (callbackHandler.notifier) {
    asyncClient.registerApplicationMaster("localhost", 1234, null);
    while(callbackHandler.notify == false) {
      try {
        callbackHandler.notifier.wait();
      } catch (InterruptedException e) {
        e.printStackTrace();
      }
    }
  }
}
 

/**
 * Get the list of matching container requests for the specified resource memory and cores.
 */
public List<? extends Collection<AMRMClient.ContainerRequest>> getMatchingRequestsList(int memory, int cores) {
  Resource resource = Resource.newInstance(memory, cores);
  Priority priority = Priority.newInstance(0);

  return getAmrmClientAsync().getMatchingRequests(priority, ResourceRequest.ANY, resource);
}
 

/**
 * Tests that YarnResourceManager will not request more containers than needs during
 * callback from Yarn when container is Completed.
 */
@Test
public void testOnContainerCompleted() throws Exception {
	new Context() {{
		runTest(() -> {
			CompletableFuture<?> registerSlotRequestFuture = resourceManager.runInMainThread(() -> {
				rmServices.slotManager.registerSlotRequest(
					new SlotRequest(new JobID(), new AllocationID(), resourceProfile1, taskHost));
				return null;
			});

			// wait for the registerSlotRequest completion
			registerSlotRequestFuture.get();

			// Callback from YARN when container is allocated.
			Container testingContainer = mockContainer("container", 1234, 1, resourceManager.getContainerResource());

			doReturn(Collections.singletonList(Collections.singletonList(resourceManager.getContainerRequest())))
				.when(mockResourceManagerClient).getMatchingRequests(any(Priority.class), anyString(), any(Resource.class));

			resourceManager.onContainersAllocated(ImmutableList.of(testingContainer));
			verify(mockResourceManagerClient).addContainerRequest(any(AMRMClient.ContainerRequest.class));
			verify(mockResourceManagerClient).removeContainerRequest(any(AMRMClient.ContainerRequest.class));
			verify(mockNMClient).startContainer(eq(testingContainer), any(ContainerLaunchContext.class));

			// Callback from YARN when container is Completed, pending request can not be fulfilled by pending
			// containers, need to request new container.
			ContainerStatus testingContainerStatus = mockContainerStatus(testingContainer.getId());

			resourceManager.onContainersCompleted(ImmutableList.of(testingContainerStatus));
			verify(mockResourceManagerClient, times(2)).addContainerRequest(any(AMRMClient.ContainerRequest.class));

			// Callback from YARN when container is Completed happened before global fail, pending request
			// slot is already fulfilled by pending containers, no need to request new container.
			resourceManager.onContainersCompleted(ImmutableList.of(testingContainerStatus));
			verify(mockResourceManagerClient, times(2)).addContainerRequest(any(AMRMClient.ContainerRequest.class));
		});
	}};
}
 

private void runHeartBeatThrowOutException(Exception ex) throws Exception{
  Configuration conf = new Configuration();
  TestCallbackHandler callbackHandler = new TestCallbackHandler();
  @SuppressWarnings("unchecked")
  AMRMClient<ContainerRequest> client = mock(AMRMClientImpl.class);
  when(client.allocate(anyFloat())).thenThrow(ex);

  AMRMClientAsync<ContainerRequest> asyncClient =
      AMRMClientAsync.createAMRMClientAsync(client, 20, callbackHandler);
  asyncClient.init(conf);
  asyncClient.start();
  
  synchronized (callbackHandler.notifier) {
    asyncClient.registerApplicationMaster("localhost", 1234, null);
    while(callbackHandler.savedException == null) {
      try {
        callbackHandler.notifier.wait();
      } catch (InterruptedException e) {
        e.printStackTrace();
      }
    }
  }
  Assert.assertTrue(callbackHandler.savedException.getMessage().contains(
      ex.getMessage()));
  
  asyncClient.stop();
  // stopping should have joined all threads and completed all callbacks
  Assert.assertTrue(callbackHandler.callbackCount == 0);
}
 

@Nonnull
@VisibleForTesting
AMRMClient.ContainerRequest getContainerRequest() {
	return new AMRMClient.ContainerRequest(
		getContainerResource(),
		null,
		null,
		RM_REQUEST_PRIORITY);
}
 

protected AMRMClientAsync<AMRMClient.ContainerRequest> createAndStartResourceManagerClient(
		YarnConfiguration yarnConfiguration,
		int yarnHeartbeatIntervalMillis,
		@Nullable String webInterfaceUrl) throws Exception {
	AMRMClientAsync<AMRMClient.ContainerRequest> resourceManagerClient = AMRMClientAsync.createAMRMClientAsync(
		yarnHeartbeatIntervalMillis,
		this);

	resourceManagerClient.init(yarnConfiguration);
	resourceManagerClient.start();

	//TODO: change akka address to tcp host and port, the getAddress() interface should return a standard tcp address
	Tuple2<String, Integer> hostPort = parseHostPort(getAddress());

	final int restPort;

	if (webInterfaceUrl != null) {
		final int lastColon = webInterfaceUrl.lastIndexOf(':');

		if (lastColon == -1) {
			restPort = -1;
		} else {
			restPort = Integer.valueOf(webInterfaceUrl.substring(lastColon + 1));
		}
	} else {
		restPort = -1;
	}

	final RegisterApplicationMasterResponse registerApplicationMasterResponse =
		resourceManagerClient.registerApplicationMaster(hostPort.f0, restPort, webInterfaceUrl);
	getContainersFromPreviousAttempts(registerApplicationMasterResponse);

	return resourceManagerClient;
}