AMRMClientRelayer implements ApplicationMasterProtocol { @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) { 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); } } try { allocateResponse = this.rmClient.allocate(allocateRequest); this.resetResponseId = -1; } catch (ApplicationMasterNotRegisteredException e) { LOG.warn("ApplicationMaster is out of sync with RM " + rmId + " for " + this.appId + ", hence resyncing."); this.metrics.incrRMMasterSlaveSwitch(this.rmId); synchronized (this) { 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); } } reRegisterApplicationMaster(this.amRegistrationRequest); allocateRequest.setResponseId(0); allocateResponse = allocate(allocateRequest); return allocateResponse; } catch (Throwable t) { this.metrics.addHeartbeatFailure(this.rmId, System.currentTimeMillis() - startTime); 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; } } AMRMClientRelayer(ApplicationMasterProtocol rmClient,
      ApplicationId appId, String rmId); void setAMRegistrationRequest(
      RegisterApplicationMasterRequest registerRequest); String getRMIdentifier(); void setRMClient(ApplicationMasterProtocol client); void shutdown(); @Override RegisterApplicationMasterResponse registerApplicationMaster(
      RegisterApplicationMasterRequest request); @Override FinishApplicationMasterResponse finishApplicationMaster(
      FinishApplicationMasterRequest request); @Override AllocateResponse allocate(AllocateRequest allocateRequest);  }
@Test public void testResendRequestsOnRMRestart() throws YarnException, IOException { ContainerId c1 = createContainerId(1); ContainerId c2 = createContainerId(2); ContainerId c3 = createContainerId(3); this.asks.add(createResourceRequest(0, "node1", 2048, 1, 1, ExecutionType.GUARANTEED, 1)); this.asks.add(createResourceRequest(0, "rack", 2048, 1, 1, ExecutionType.GUARANTEED, 1)); this.asks.add(createResourceRequest(0, ResourceRequest.ANY, 2048, 1, 1, ExecutionType.GUARANTEED, 2)); this.releases.add(c1); this.blacklistAdditions.add("node1"); this.blacklistRemoval.add("node0"); this.relayer.allocate(getAllocateRequest()); assertAsksAndReleases(3, 1); assertBlacklistAdditionsAndRemovals(1, 1); clearAllocateRequestLists(); this.relayer.allocate(getAllocateRequest()); assertAsksAndReleases(0, 0); assertBlacklistAdditionsAndRemovals(0, 0); clearAllocateRequestLists(); this.mockAMS.setFailoverFlag(); this.blacklistAdditions.add("node2"); this.releases.add(c2); this.relayer.allocate(getAllocateRequest()); assertAsksAndReleases(3, 2); assertBlacklistAdditionsAndRemovals(2, 0); clearAllocateRequestLists(); }
@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) { } response = this.relayer.allocate(getAllocateRequest()); Assert.assertEquals(expected + 1, response.getResponseId()); this.responseId = response.getResponseId(); response = this.relayer.allocate(getAllocateRequest()); Assert.assertEquals(this.responseId + 1, response.getResponseId()); }