@SuppressWarnings("unchecked")
  @Override
  protected synchronized void heartbeat() throws Exception {
    AllocateRequest allocateRequest =
        AllocateRequest.newInstance(this.lastResponseID,
          super.getApplicationProgress(), new ArrayList<ResourceRequest>(),
        new ArrayList<ContainerId>(), null);
    AllocateResponse allocateResponse = null;
    try {
      allocateResponse = scheduler.allocate(allocateRequest);
      // Reset retry count if no exception occurred.
      retrystartTime = System.currentTimeMillis();
    } catch (ApplicationAttemptNotFoundException e) {
      LOG.info("Event from RM: shutting down Application Master");
      // This can happen if the RM has been restarted. If it is in that state,
      // this application must clean itself up.
      eventHandler.handle(new JobEvent(this.getJob().getID(),
        JobEventType.JOB_AM_REBOOT));
      throw new YarnRuntimeException(
        "Resource Manager doesn't recognize AttemptId: "
            + this.getContext().getApplicationID(), e);
    } catch (ApplicationMasterNotRegisteredException e) {
      LOG.info("ApplicationMaster is out of sync with ResourceManager,"
          + " hence resync and send outstanding requests.");
      this.lastResponseID = 0;
      register();
    } catch (Exception e) {
      // This can happen when the connection to the RM has gone down. Keep
      // re-trying until the retryInterval has expired.
      if (System.currentTimeMillis() - retrystartTime >= retryInterval) {
        LOG.error("Could not contact RM after " + retryInterval + " milliseconds.");
        eventHandler.handle(new JobEvent(this.getJob().getID(),
                                         JobEventType.INTERNAL_ERROR));
        throw new YarnRuntimeException("Could not contact RM after " +
                                retryInterval + " milliseconds.");
      }
      // Throw this up to the caller, which may decide to ignore it and
      // continue to attempt to contact the RM.
      throw e;
    }

    if (allocateResponse != null) {
      this.lastResponseID = allocateResponse.getResponseId();
      Token token = allocateResponse.getAMRMToken();
      if (token != null) {
        updateAMRMToken(token);
      }
      Priority priorityFromResponse = Priority.newInstance(allocateResponse
          .getApplicationPriority().getPriority());

      // Update the job priority to Job directly.
      getJob().setJobPriority(priorityFromResponse);
    }
  }
@Test
  public void testAMRMTokenUpdate() throws Exception {
    Configuration conf = new Configuration();
    ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(
        ApplicationId.newInstance(1, 1), 1);
    AMRMTokenIdentifier oldTokenId = new AMRMTokenIdentifier(attemptId, 1);
    AMRMTokenIdentifier newTokenId = new AMRMTokenIdentifier(attemptId, 2);
    Token<AMRMTokenIdentifier> oldToken = new Token<AMRMTokenIdentifier>(
        oldTokenId.getBytes(), "oldpassword".getBytes(), oldTokenId.getKind(),
        new Text());
    Token<AMRMTokenIdentifier> newToken = new Token<AMRMTokenIdentifier>(
        newTokenId.getBytes(), "newpassword".getBytes(), newTokenId.getKind(),
        new Text());

    MockScheduler scheduler = new MockScheduler();
    scheduler.amToken = newToken;

    final LocalContainerAllocator lca =
        new StubbedLocalContainerAllocator(scheduler);
    lca.init(conf);
    lca.start();

    UserGroupInformation testUgi = UserGroupInformation.createUserForTesting(
        "someuser", new String[0]);
    testUgi.addToken(oldToken);
    testUgi.doAs(new PrivilegedExceptionAction<Void>() {
          @Override
          public Void run() throws Exception {
            lca.heartbeat();
            return null;
          }
    });
    lca.close();

    // verify there is only one AMRM token in the UGI and it matches the
    // updated token from the RM
    int tokenCount = 0;
    Token<? extends TokenIdentifier> ugiToken = null;
    for (Token<? extends TokenIdentifier> token : testUgi.getTokens()) {
      if (AMRMTokenIdentifier.KIND_NAME.equals(token.getKind())) {
        ugiToken = token;
        ++tokenCount;
      }
    }

    Assert.assertEquals("too many AMRM tokens", 1, tokenCount);
    Assert.assertArrayEquals("token identifier not updated",
        newToken.getIdentifier(), ugiToken.getIdentifier());
    Assert.assertArrayEquals("token password not updated",
        newToken.getPassword(), ugiToken.getPassword());
    Assert.assertEquals("AMRM token service not updated",
        new Text(ClientRMProxy.getAMRMTokenService(conf)),
        ugiToken.getService());
  }