@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 testRMConnectionRetry() throws Exception {
    // verify the connection exception is thrown
    // if we haven't exhausted the retry interval
    ApplicationMasterProtocol mockScheduler =
        mock(ApplicationMasterProtocol.class);
    when(mockScheduler.allocate(isA(AllocateRequest.class)))
      .thenThrow(RPCUtil.getRemoteException(new IOException("forcefail")));
    Configuration conf = new Configuration();
    LocalContainerAllocator lca =
        new StubbedLocalContainerAllocator(mockScheduler);
    lca.init(conf);
    lca.start();
    try {
      lca.heartbeat();
      Assert.fail("heartbeat was supposed to throw");
    } catch (YarnException e) {
      // YarnException is expected
    } finally {
      lca.stop();
    }

    // verify YarnRuntimeException is thrown when the retry interval has expired
    conf.setLong(MRJobConfig.MR_AM_TO_RM_WAIT_INTERVAL_MS, 0);
    lca = new StubbedLocalContainerAllocator(mockScheduler);
    lca.init(conf);
    lca.start();
    try {
      lca.heartbeat();
      Assert.fail("heartbeat was supposed to throw");
    } catch (YarnRuntimeException e) {
      // YarnRuntimeException is expected
    } finally {
      lca.stop();
    }
  }