@Override
  public AllocateResponse allocate(AllocateRequest allocateRequest)
      throws YarnException, IOException {
    AllocateResponse allocateResponse = null;
    long startTime = System.currentTimeMillis();
    synchronized (this) {
      if(this.shutdown){
        throw new YarnException("Allocate called after AMRMClientRelayer for "
            + "RM " + rmId + " shutdown.");
      }
      addNewAllocateRequest(allocateRequest);

      ArrayList<ResourceRequest> askList = new ArrayList<>(ask.size());
      for (ResourceRequest r : ask) {
        // create a copy of ResourceRequest as we might change it while the
        // RPC layer is using it to send info across
        askList.add(ResourceRequest.clone(r));
      }

      allocateRequest = AllocateRequest.newBuilder()
          .responseId(allocateRequest.getResponseId())
          .progress(allocateRequest.getProgress()).askList(askList)
          .releaseList(new ArrayList<>(this.release))
          .resourceBlacklistRequest(ResourceBlacklistRequest.newInstance(
              new ArrayList<>(this.blacklistAdditions),
              new ArrayList<>(this.blacklistRemovals)))
          .updateRequests(new ArrayList<>(this.change.values()))
          .schedulingRequests(new ArrayList<>(this.schedulingRequest))
          .build();

      if (this.resetResponseId != -1) {
        LOG.info("Override allocate responseId from "
            + allocateRequest.getResponseId() + " to " + this.resetResponseId
            + " for " + this.appId);
        allocateRequest.setResponseId(this.resetResponseId);
      }
    }

    // Do the actual allocate call
    try {
      allocateResponse = this.rmClient.allocate(allocateRequest);

      // Heartbeat succeeded, wipe out responseId overriding
      this.resetResponseId = -1;
    } catch (ApplicationMasterNotRegisteredException e) {
      // This is a retriable exception - we will re register and mke a
      // recursive call to retry
      LOG.warn("ApplicationMaster is out of sync with RM " + rmId
          + " for " + this.appId + ", hence resyncing.");

      this.metrics.incrRMMasterSlaveSwitch(this.rmId);

      synchronized (this) {
        // Add all remotePending data into to-send data structures
        for (ResourceRequestSet requestSet : this.remotePendingAsks
            .values()) {
          for (ResourceRequest rr : requestSet.getRRs()) {
            addResourceRequestToAsk(rr);
          }
        }
        this.release.addAll(this.remotePendingRelease);
        this.blacklistAdditions.addAll(this.remoteBlacklistedNodes);
        this.change.putAll(this.remotePendingChange);
        for (List<SchedulingRequest> reqs : this.remotePendingSchedRequest
            .values()) {
          this.schedulingRequest.addAll(reqs);
        }
      }

      // re-register with RM, then retry allocate recursively
      reRegisterApplicationMaster(this.amRegistrationRequest);
      // Reset responseId after re-register
      allocateRequest.setResponseId(0);
      allocateResponse = allocate(allocateRequest);
      return allocateResponse;
    } catch (Throwable t) {
      // Unexpected exception - rethrow and increment heart beat failure metric
      this.metrics.addHeartbeatFailure(this.rmId,
          System.currentTimeMillis() - startTime);

      // If RM is complaining about responseId out of sync, force reset next
      // time
      if (t instanceof InvalidApplicationMasterRequestException) {
        int responseId = AMRMClientUtils
            .parseExpectedResponseIdFromException(t.getMessage());
        if (responseId != -1) {
          this.resetResponseId = responseId;
          LOG.info("ResponseId out of sync with RM, expect " + responseId
              + " but " + allocateRequest.getResponseId() + " used by "
              + this.appId + ". Will override in the next allocate.");
        } else {
          LOG.warn("Failed to parse expected responseId out of exception for "
              + this.appId);
        }
      }

      throw t;
    }

    synchronized (this) {
      if (this.shutdown) {
        throw new YarnException("Allocate call succeeded for " + this.appId
            + " after AMRMClientRelayer for RM " + rmId + " shutdown.");
      }

      updateMetrics(allocateResponse, startTime);

      AMRMClientUtils.removeFromOutstandingSchedulingRequests(
          allocateResponse.getAllocatedContainers(),
          this.remotePendingSchedRequest);
      AMRMClientUtils.removeFromOutstandingSchedulingRequests(
          allocateResponse.getContainersFromPreviousAttempts(),
          this.remotePendingSchedRequest);

      this.ask.clear();
      this.release.clear();

      this.blacklistAdditions.clear();
      this.blacklistRemovals.clear();

      this.change.clear();
      this.schedulingRequest.clear();
      return allocateResponse;
    }
  }
@Test
  public void testResponseIdResync() throws YarnException, IOException {
    this.responseId = 10;

    AllocateResponse response = this.relayer.allocate(getAllocateRequest());
    Assert.assertEquals(this.responseId + 1, response.getResponseId());

    int expected = 5;
    this.mockAMS.setResponseIdReset(expected);

    try {
      this.relayer.allocate(getAllocateRequest());
      Assert.fail("Expecting exception from RM");
    } catch (InvalidApplicationMasterRequestException e) {
      // Expected exception
    }

    // Verify that the responseId is overridden
    response = this.relayer.allocate(getAllocateRequest());
    Assert.assertEquals(expected + 1, response.getResponseId());

    // Verify it is no longer overriden
    this.responseId = response.getResponseId();
    response = this.relayer.allocate(getAllocateRequest());
    Assert.assertEquals(this.responseId + 1, response.getResponseId());
  }