FederationInterceptor extends AbstractRequestInterceptor { @VisibleForTesting protected void mergeAllocateResponse(AllocateResponse homeResponse, AllocateResponse otherResponse, SubClusterId otherRMAddress) { if (otherResponse.getAMRMToken() != null) { if (otherRMAddress.equals(this.homeSubClusterId)) { homeResponse.setAMRMToken(otherResponse.getAMRMToken()); } else { throw new YarnRuntimeException( "amrmToken from UAM " + otherRMAddress + " should be null here"); } } if (!isNullOrEmpty(otherResponse.getAllocatedContainers())) { if (!isNullOrEmpty(homeResponse.getAllocatedContainers())) { homeResponse.getAllocatedContainers() .addAll(otherResponse.getAllocatedContainers()); } else { homeResponse .setAllocatedContainers(otherResponse.getAllocatedContainers()); } } if (!isNullOrEmpty(otherResponse.getCompletedContainersStatuses())) { if (!isNullOrEmpty(homeResponse.getCompletedContainersStatuses())) { homeResponse.getCompletedContainersStatuses() .addAll(otherResponse.getCompletedContainersStatuses()); } else { homeResponse.setCompletedContainersStatuses( otherResponse.getCompletedContainersStatuses()); } } if (!isNullOrEmpty(otherResponse.getUpdatedNodes())) { if (!isNullOrEmpty(homeResponse.getUpdatedNodes())) { homeResponse.getUpdatedNodes().addAll(otherResponse.getUpdatedNodes()); } else { homeResponse.setUpdatedNodes(otherResponse.getUpdatedNodes()); } } if (otherResponse.getApplicationPriority() != null) { homeResponse.setApplicationPriority( otherResponse.getApplicationPriority()); } homeResponse.setNumClusterNodes( homeResponse.getNumClusterNodes() + otherResponse.getNumClusterNodes()); PreemptionMessage homePreempMessage = homeResponse.getPreemptionMessage(); PreemptionMessage otherPreempMessage = otherResponse.getPreemptionMessage(); if (homePreempMessage == null && otherPreempMessage != null) { homeResponse.setPreemptionMessage(otherPreempMessage); } if (homePreempMessage != null && otherPreempMessage != null) { PreemptionContract par1 = homePreempMessage.getContract(); PreemptionContract par2 = otherPreempMessage.getContract(); if (par1 == null && par2 != null) { homePreempMessage.setContract(par2); } if (par1 != null && par2 != null) { par1.getResourceRequest().addAll(par2.getResourceRequest()); par1.getContainers().addAll(par2.getContainers()); } StrictPreemptionContract spar1 = homePreempMessage.getStrictContract(); StrictPreemptionContract spar2 = otherPreempMessage.getStrictContract(); if (spar1 == null && spar2 != null) { homePreempMessage.setStrictContract(spar2); } if (spar1 != null && spar2 != null) { spar1.getContainers().addAll(spar2.getContainers()); } } if (!isNullOrEmpty(otherResponse.getNMTokens())) { if (!isNullOrEmpty(homeResponse.getNMTokens())) { homeResponse.getNMTokens().addAll(otherResponse.getNMTokens()); } else { homeResponse.setNMTokens(otherResponse.getNMTokens()); } } if (!isNullOrEmpty(otherResponse.getUpdatedContainers())) { if (!isNullOrEmpty(homeResponse.getUpdatedContainers())) { homeResponse.getUpdatedContainers() .addAll(otherResponse.getUpdatedContainers()); } else { homeResponse.setUpdatedContainers(otherResponse.getUpdatedContainers()); } } if (!isNullOrEmpty(otherResponse.getUpdateErrors())) { if (!isNullOrEmpty(homeResponse.getUpdateErrors())) { homeResponse.getUpdateErrors().addAll(otherResponse.getUpdateErrors()); } else { homeResponse.setUpdateErrors(otherResponse.getUpdateErrors()); } } } FederationInterceptor(); @Override void init(AMRMProxyApplicationContext appContext); @Override void recover(Map<String, byte[]> recoveredDataMap); @Override synchronized RegisterApplicationMasterResponse registerApplicationMaster(RegisterApplicationMasterRequest request); @Override AllocateResponse allocate(AllocateRequest request); @Override FinishApplicationMasterResponse finishApplicationMaster(
      FinishApplicationMasterRequest request); @Override void setNextInterceptor(RequestInterceptor next); @Override void shutdown(); @VisibleForTesting Map<SubClusterId, List<AllocateResponse>> getAsyncResponseSink(); static boolean isNullOrEmpty(Collection<T> c); static boolean isNullOrEmpty(Map<T1, T2> c); static final String NMSS_CLASS_PREFIX; static final String NMSS_REG_REQUEST_KEY; static final String NMSS_REG_RESPONSE_KEY; static final String NMSS_SECONDARY_SC_PREFIX; static final String STRING_TO_BYTE_FORMAT; }
@Test public void testMergeAllocateResponse() { ContainerId cid = ContainerId.newContainerId(attemptId, 0); ContainerStatus cStatus = Records.newRecord(ContainerStatus.class); cStatus.setContainerId(cid); Container container = Container.newInstance(cid, null, null, null, null, null); AllocateResponse homeResponse = Records.newRecord(AllocateResponse.class); homeResponse.setAllocatedContainers(Collections.singletonList(container)); homeResponse.setCompletedContainersStatuses( Collections.singletonList(cStatus)); homeResponse.setUpdatedNodes( Collections.singletonList(Records.newRecord(NodeReport.class))); homeResponse.setNMTokens( Collections.singletonList(Records.newRecord(NMToken.class))); homeResponse.setUpdatedContainers( Collections.singletonList( Records.newRecord(UpdatedContainer.class))); homeResponse.setUpdateErrors(Collections .singletonList(Records.newRecord(UpdateContainerError.class))); homeResponse.setAvailableResources(Records.newRecord(Resource.class)); homeResponse.setPreemptionMessage(createDummyPreemptionMessage( ContainerId.newContainerId(attemptId, 0))); AllocateResponse response = Records.newRecord(AllocateResponse.class); response.setAllocatedContainers(Collections.singletonList(container)); response.setCompletedContainersStatuses(Collections.singletonList(cStatus)); response.setUpdatedNodes( Collections.singletonList(Records.newRecord(NodeReport.class))); response.setNMTokens( Collections.singletonList(Records.newRecord(NMToken.class))); response.setUpdatedContainers( Collections.singletonList( Records.newRecord(UpdatedContainer.class))); response.setUpdateErrors(Collections .singletonList(Records.newRecord(UpdateContainerError.class))); response.setAvailableResources(Records.newRecord(Resource.class)); response.setPreemptionMessage(createDummyPreemptionMessage( ContainerId.newContainerId(attemptId, 1))); interceptor.mergeAllocateResponse(homeResponse, response, SubClusterId.newInstance("SC-1")); Assert.assertEquals(2, homeResponse.getPreemptionMessage().getContract() .getContainers().size()); Assert.assertEquals(2, homeResponse.getAllocatedContainers().size()); Assert.assertEquals(2, homeResponse.getUpdatedNodes().size()); Assert.assertEquals(2, homeResponse.getCompletedContainersStatuses().size()); }