FederationInterceptor extends AbstractRequestInterceptor { @Override public synchronized RegisterApplicationMasterResponse registerApplicationMaster(RegisterApplicationMasterRequest request) throws YarnException, IOException { synchronized (this.lastAllocateResponseLock) { this.lastAllocateResponse.setResponseId(0); } this.justRecovered = false; if (this.amRegistrationRequest != null) { if (!this.amRegistrationRequest.equals(request)) { throw new YarnException("AM should not call " + "registerApplicationMaster with a different request body"); } } else { this.amRegistrationRequest = request; if (getNMStateStore() != null) { try { RegisterApplicationMasterRequestPBImpl pb = (RegisterApplicationMasterRequestPBImpl) this.amRegistrationRequest; getNMStateStore().storeAMRMProxyAppContextEntry(this.attemptId, NMSS_REG_REQUEST_KEY, pb.getProto().toByteArray()); } catch (Exception e) { LOG.error("Error storing AMRMProxy application context entry for " + this.attemptId, e); } } } if (this.amRegistrationResponse != null) { return this.amRegistrationResponse; } this.amRegistrationResponse = this.homeRMRelayer.registerApplicationMaster(request); if (this.amRegistrationResponse .getContainersFromPreviousAttempts() != null) { cacheAllocatedContainers( this.amRegistrationResponse.getContainersFromPreviousAttempts(), this.homeSubClusterId); } ApplicationId appId = this.attemptId.getApplicationId(); reAttachUAMAndMergeRegisterResponse(this.amRegistrationResponse, appId); if (getNMStateStore() != null) { try { RegisterApplicationMasterResponsePBImpl pb = (RegisterApplicationMasterResponsePBImpl) this.amRegistrationResponse; getNMStateStore().storeAMRMProxyAppContextEntry(this.attemptId, NMSS_REG_RESPONSE_KEY, pb.getProto().toByteArray()); } catch (Exception e) { LOG.error("Error storing AMRMProxy application context entry for " + this.attemptId, e); } } String queue = this.amRegistrationResponse.getQueue(); if (queue == null) { LOG.warn("Received null queue for application " + appId + " from home subcluster. Will use default queue name " + YarnConfiguration.DEFAULT_QUEUE_NAME + " for getting AMRMProxyPolicy"); } else { LOG.info("Application " + appId + " belongs to queue " + queue); } try { this.policyInterpreter = FederationPolicyUtils.loadAMRMPolicy(queue, this.policyInterpreter, getConf(), this.federationFacade, this.homeSubClusterId); } catch (FederationPolicyInitializationException e) { throw new YarnRuntimeException(e); } return this.amRegistrationResponse; } 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 int getUnmanagedAMPoolSize(); @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 testTwoIdenticalRegisterRequest() throws Exception { RegisterApplicationMasterRequest registerReq = Records.newRecord(RegisterApplicationMasterRequest.class); registerReq.setHost(Integer.toString(testAppId)); registerReq.setRpcPort(0); registerReq.setTrackingUrl(""); for (int i = 0; i < 2; i++) { RegisterApplicationMasterResponse registerResponse = interceptor.registerApplicationMaster(registerReq); Assert.assertNotNull(registerResponse); lastResponseId = 0; } }
@Test public void testTwoDifferentRegisterRequest() throws Exception { RegisterApplicationMasterRequest registerReq = Records.newRecord(RegisterApplicationMasterRequest.class); registerReq.setHost(Integer.toString(testAppId)); registerReq.setRpcPort(0); registerReq.setTrackingUrl(""); RegisterApplicationMasterResponse registerResponse = interceptor.registerApplicationMaster(registerReq); Assert.assertNotNull(registerResponse); lastResponseId = 0; registerReq = Records.newRecord(RegisterApplicationMasterRequest.class); registerReq.setHost(Integer.toString(testAppId)); registerReq.setRpcPort(0); registerReq.setTrackingUrl("different"); try { registerResponse = interceptor.registerApplicationMaster(registerReq); lastResponseId = 0; Assert.fail("Should throw if a different request obj is used"); } catch (YarnException e) { } }