org.apache.hadoop.yarn.server.resourcemanager.recovery.RMStateStore Java Examples

@Override
protected void serviceStop() throws Exception {

  DefaultMetricsSystem.shutdown();

  if (rmContext != null) {
    RMStateStore store = rmContext.getStateStore();
    try {
      store.close();
    } catch (Exception e) {
      LOG.error("Error closing store.", e);
    }
  }

  super.serviceStop();
}
 

private void recoverAppAttemptCredentials(Credentials appAttemptTokens,
    RMAppAttemptState state) {
  if (appAttemptTokens == null || state == RMAppAttemptState.FAILED
      || state == RMAppAttemptState.FINISHED
      || state == RMAppAttemptState.KILLED) {
    return;
  }

  if (UserGroupInformation.isSecurityEnabled()) {
    byte[] clientTokenMasterKeyBytes = appAttemptTokens.getSecretKey(
        RMStateStore.AM_CLIENT_TOKEN_MASTER_KEY_NAME);
    if (clientTokenMasterKeyBytes != null) {
      clientTokenMasterKey = rmContext.getClientToAMTokenSecretManager()
          .registerMasterKey(applicationAttemptId, clientTokenMasterKeyBytes);
    }
  }

  setAMRMToken(rmContext.getAMRMTokenSecretManager().createAndGetAMRMToken(
      applicationAttemptId));
}
 

private void recoverAppAttemptCredentials(Credentials appAttemptTokens,
    RMAppAttemptState state) {
  if (appAttemptTokens == null || state == RMAppAttemptState.FAILED
      || state == RMAppAttemptState.FINISHED
      || state == RMAppAttemptState.KILLED) {
    return;
  }

  if (UserGroupInformation.isSecurityEnabled()) {
    byte[] clientTokenMasterKeyBytes = appAttemptTokens.getSecretKey(
        RMStateStore.AM_CLIENT_TOKEN_MASTER_KEY_NAME);
    if (clientTokenMasterKeyBytes != null) {
      clientTokenMasterKey = rmContext.getClientToAMTokenSecretManager()
          .registerMasterKey(applicationAttemptId, clientTokenMasterKeyBytes);
    }
  }

  setAMRMToken(rmContext.getAMRMTokenSecretManager().createAndGetAMRMToken(
      applicationAttemptId));
}
 

@Override
protected void serviceStop() throws Exception {

  DefaultMetricsSystem.shutdown();

  if (rmContext != null) {
    RMStateStore store = rmContext.getStateStore();
    try {
      store.close();
    } catch (Exception e) {
      LOG.error("Error closing store.", e);
    }
  }

  super.serviceStop();
}
 

/**
 * Deletes the RMStateStore
 *
 * @param conf
 * @throws Exception
 */
private static void deleteRMStateStore(Configuration conf) throws Exception {
  RMStateStore rmStore = RMStateStoreFactory.getStore(conf);
  rmStore.init(conf);
  rmStore.start();
  try {
    LOG.info("Deleting ResourceManager state store...");
    rmStore.deleteStore();
    LOG.info("State store deleted");
  } finally {
    rmStore.stop();
  }
}
 

public MockRM(Configuration conf, RMStateStore store,
    boolean useNullRMNodeLabelsManager) {
  super();
  this.useNullRMNodeLabelsManager = useNullRMNodeLabelsManager;
  init(conf instanceof YarnConfiguration ? conf : new YarnConfiguration(conf));
  if(store != null) {
    setRMStateStore(store);
  }
  Logger rootLogger = LogManager.getRootLogger();
  rootLogger.setLevel(Level.DEBUG);
}
 

public MockRM(Configuration conf, RMStateStore store,
    boolean useNullRMNodeLabelsManager) {
  super();
  this.useNullRMNodeLabelsManager = useNullRMNodeLabelsManager;
  init(conf instanceof YarnConfiguration ? conf : new YarnConfiguration(conf));
  if(store != null) {
    setRMStateStore(store);
  }
  Logger rootLogger = LogManager.getRootLogger();
  rootLogger.setLevel(Level.DEBUG);
}
 

public RMContext mockRMContext(int n, long time) {
  final List<RMApp> apps = newRMApps(n, time, RMAppState.FINISHED);
  final ConcurrentMap<ApplicationId, RMApp> map = Maps.newConcurrentMap();
  for (RMApp app : apps) {
    map.put(app.getApplicationId(), app);
  }
  Dispatcher rmDispatcher = new AsyncDispatcher();
  ContainerAllocationExpirer containerAllocationExpirer = new ContainerAllocationExpirer(
      rmDispatcher);
  AMLivelinessMonitor amLivelinessMonitor = new AMLivelinessMonitor(
      rmDispatcher);
  AMLivelinessMonitor amFinishingMonitor = new AMLivelinessMonitor(
      rmDispatcher);
  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  RMContext context = new RMContextImpl(rmDispatcher,
      containerAllocationExpirer, amLivelinessMonitor, amFinishingMonitor,
      null, null, null, null, null, writer) {
    @Override
    public ConcurrentMap<ApplicationId, RMApp> getRMApps() {
      return map;
    }
  };
  ((RMContextImpl)context).setStateStore(mock(RMStateStore.class));
  metricsPublisher = mock(SystemMetricsPublisher.class);
  ((RMContextImpl)context).setSystemMetricsPublisher(metricsPublisher);
  return context;
}
 

public RMContext mockRMContext(int n, long time) {
  final List<RMApp> apps = newRMApps(n, time, RMAppState.FINISHED);
  final ConcurrentMap<ApplicationId, RMApp> map = Maps.newConcurrentMap();
  for (RMApp app : apps) {
    map.put(app.getApplicationId(), app);
  }
  Dispatcher rmDispatcher = new AsyncDispatcher();
  ContainerAllocationExpirer containerAllocationExpirer = new ContainerAllocationExpirer(
      rmDispatcher);
  AMLivelinessMonitor amLivelinessMonitor = new AMLivelinessMonitor(
      rmDispatcher);
  AMLivelinessMonitor amFinishingMonitor = new AMLivelinessMonitor(
      rmDispatcher);
  RMApplicationHistoryWriter writer = mock(RMApplicationHistoryWriter.class);
  RMContext context = new RMContextImpl(rmDispatcher,
      containerAllocationExpirer, amLivelinessMonitor, amFinishingMonitor,
      null, null, null, null, null, writer) {
    @Override
    public ConcurrentMap<ApplicationId, RMApp> getRMApps() {
      return map;
    }
  };
  ((RMContextImpl)context).setStateStore(mock(RMStateStore.class));
  metricsPublisher = mock(SystemMetricsPublisher.class);
  ((RMContextImpl)context).setSystemMetricsPublisher(metricsPublisher);
  return context;
}
 

/**
 * Deletes the RMStateStore
 *
 * @param conf
 * @throws Exception
 */
private static void deleteRMStateStore(Configuration conf) throws Exception {
  RMStateStore rmStore = RMStateStoreFactory.getStore(conf);
  rmStore.init(conf);
  rmStore.start();
  try {
    LOG.info("Deleting ResourceManager state store...");
    rmStore.deleteStore();
    LOG.info("State store deleted");
  } finally {
    rmStore.stop();
  }
}
 

@Override
protected void serviceStart() throws Exception {
  RMStateStore rmStore = rmContext.getStateStore();
  // The state store needs to start irrespective of recoveryEnabled as apps
  // need events to move to further states.
  rmStore.start();

  if(recoveryEnabled) {
    try {
      LOG.info("Recovery started");
      rmStore.checkVersion();
      if (rmContext.isWorkPreservingRecoveryEnabled()) {
        rmContext.setEpoch(rmStore.getAndIncrementEpoch());
      }
      RMState state = rmStore.loadState();
      recover(state);
      LOG.info("Recovery ended");
    } catch (Exception e) {
      // the Exception from loadState() needs to be handled for
      // HA and we need to give up master status if we got fenced
      LOG.error("Failed to load/recover state", e);
      throw e;
    }
  }

  super.serviceStart();
}
 

@Override
protected void serviceStart() throws Exception {
  RMStateStore rmStore = rmContext.getStateStore();
  // The state store needs to start irrespective of recoveryEnabled as apps
  // need events to move to further states.
  rmStore.start();

  if(recoveryEnabled) {
    try {
      LOG.info("Recovery started");
      rmStore.checkVersion();
      if (rmContext.isWorkPreservingRecoveryEnabled()) {
        rmContext.setEpoch(rmStore.getAndIncrementEpoch());
      }
      RMState state = rmStore.loadState();
      recover(state);
      LOG.info("Recovery ended");
    } catch (Exception e) {
      // the Exception from loadState() needs to be handled for
      // HA and we need to give up master status if we got fenced
      LOG.error("Failed to load/recover state", e);
      throw e;
    }
  }

  super.serviceStart();
}
 

@Override
public RMStateStore getStateStore() {
  return activeServiceContext.getStateStore();
}
 

@VisibleForTesting
public void setStateStore(RMStateStore store) {
  activeServiceContext.setStateStore(store);
}
 

public MockRM(Configuration conf, RMStateStore store) {
  this(conf, store, true);
}
 

public MyMockRM(Configuration conf, RMStateStore store) {
  super(conf, store);
}
 

/**
 *
 * @return a new MockRM that will be stopped at the end of the test.
 */
private MockRM createMockRM(YarnConfiguration conf, RMStateStore store) {
  MockRM rm = new MockRM(conf, store);
  rms.add(rm);
  return rm;
}
 

@Test (timeout = 60000)
public void testAppAttemptTokensRestoredOnRMRestart() throws Exception {
  conf.setInt(YarnConfiguration.RM_AM_MAX_ATTEMPTS, 2);
  conf.set(CommonConfigurationKeysPublic.HADOOP_SECURITY_AUTHENTICATION,
    "kerberos");
  UserGroupInformation.setConfiguration(conf);

  MemoryRMStateStore memStore = new MemoryRMStateStore();
  memStore.init(conf);
  RMState rmState = memStore.getState();

  Map<ApplicationId, ApplicationStateData> rmAppState =
      rmState.getApplicationState();
  MockRM rm1 = new TestSecurityMockRM(conf, memStore);
  rm1.start();
  MockNM nm1 =
      new MockNM("0.0.0.0:4321", 15120, rm1.getResourceTrackerService());
  nm1.registerNode();

  // submit an app
  RMApp app1 =
      rm1.submitApp(200, "name", "user",
        new HashMap<ApplicationAccessType, String>(), "default");

  // assert app info is saved
  ApplicationStateData appState = rmAppState.get(app1.getApplicationId());
  Assert.assertNotNull(appState);

  // Allocate the AM
  nm1.nodeHeartbeat(true);
  RMAppAttempt attempt1 = app1.getCurrentAppAttempt();
  ApplicationAttemptId attemptId1 = attempt1.getAppAttemptId();
  rm1.waitForState(attemptId1, RMAppAttemptState.ALLOCATED);

  // assert attempt info is saved
  ApplicationAttemptStateData attemptState = appState.getAttempt(attemptId1);
  Assert.assertNotNull(attemptState);
  Assert.assertEquals(BuilderUtils.newContainerId(attemptId1, 1),
    attemptState.getMasterContainer().getId());

  // the clientTokenMasterKey that are generated when
  // RMAppAttempt is created,
  byte[] clientTokenMasterKey =
      attempt1.getClientTokenMasterKey().getEncoded();

  // assert application credentials are saved
  Credentials savedCredentials = attemptState.getAppAttemptTokens();
  Assert.assertArrayEquals("client token master key not saved",
      clientTokenMasterKey, savedCredentials.getSecretKey(
          RMStateStore.AM_CLIENT_TOKEN_MASTER_KEY_NAME));

  // start new RM
  MockRM rm2 = new TestSecurityMockRM(conf, memStore);
  rm2.start();

  RMApp loadedApp1 =
      rm2.getRMContext().getRMApps().get(app1.getApplicationId());
  RMAppAttempt loadedAttempt1 = loadedApp1.getRMAppAttempt(attemptId1);

  // assert loaded attempt recovered
  Assert.assertNotNull(loadedAttempt1);

  // assert client token master key is recovered back to api-versioned
  // client token master key
  Assert.assertEquals("client token master key not restored",
      attempt1.getClientTokenMasterKey(),
      loadedAttempt1.getClientTokenMasterKey());

  // assert ClientTokenSecretManager also knows about the key
  Assert.assertArrayEquals(clientTokenMasterKey,
      rm2.getClientToAMTokenSecretManager().getMasterKey(attemptId1)
          .getEncoded());

  // assert AMRMTokenSecretManager also knows about the AMRMToken password
  Token<AMRMTokenIdentifier> amrmToken = loadedAttempt1.getAMRMToken();
  Assert.assertArrayEquals(amrmToken.getPassword(),
    rm2.getRMContext().getAMRMTokenSecretManager().retrievePassword(
      amrmToken.decodeIdentifier()));
}
 

public TestSecurityMockRM(Configuration conf, RMStateStore store) {
  super(conf, store);
}
 

@BeforeClass
public static void setup() throws InterruptedException, IOException {
  RMStateStore store = RMStateStoreFactory.getStore(conf);
  conf.setBoolean(YarnConfiguration.YARN_ACL_ENABLE, true);
  AccessControlList adminACL = new AccessControlList("");
  adminACL.addGroup(SUPER_GROUP);
  conf.set(YarnConfiguration.YARN_ADMIN_ACL, adminACL.getAclString());
  resourceManager = new MockRM(conf) {

    @Override
    protected QueueACLsManager createQueueACLsManager(
        ResourceScheduler scheduler,
        Configuration conf) {
      QueueACLsManager mockQueueACLsManager = mock(QueueACLsManager.class);
      when(mockQueueACLsManager.checkAccess(any(UserGroupInformation.class),
          any(QueueACL.class), anyString())).thenAnswer(new Answer() {
        public Object answer(InvocationOnMock invocation) {
          return isQueueUser;
        }
      });
      return mockQueueACLsManager;
    }

    protected ClientRMService createClientRMService() {
      return new ClientRMService(getRMContext(), this.scheduler,
          this.rmAppManager, this.applicationACLsManager,
          this.queueACLsManager, null);
    };
  };
  new Thread() {
    public void run() {
      UserGroupInformation.createUserForTesting(ENEMY, new String[] {});
      UserGroupInformation.createUserForTesting(FRIEND,
          new String[] { FRIENDLY_GROUP });
      UserGroupInformation.createUserForTesting(SUPER_USER,
          new String[] { SUPER_GROUP });
      resourceManager.start();
    };
  }.start();
  int waitCount = 0;
  while (resourceManager.getServiceState() == STATE.INITED
      && waitCount++ < 60) {
    LOG.info("Waiting for RM to start...");
    Thread.sleep(1500);
  }
  if (resourceManager.getServiceState() != STATE.STARTED) {
    // RM could have failed.
    throw new IOException(
        "ResourceManager failed to start. Final state is "
            + resourceManager.getServiceState());
  }

  UserGroupInformation owner = UserGroupInformation
      .createRemoteUser(APP_OWNER);
  rmClient = owner.doAs(new PrivilegedExceptionAction<ApplicationClientProtocol>() {
    @Override
    public ApplicationClientProtocol run() throws Exception {
      return (ApplicationClientProtocol) rpc.getProxy(ApplicationClientProtocol.class,
          rmAddress, conf);
    }
  });
}
 

private void rememberTargetTransitionsAndStoreState(RMAppAttemptEvent event,
    Object transitionToDo, RMAppAttemptState targetFinalState,
    RMAppAttemptState stateToBeStored) {

  rememberTargetTransitions(event, transitionToDo, targetFinalState);
  stateBeforeFinalSaving = getState();

  // As of today, finalState, diagnostics, final-tracking-url and
  // finalAppStatus are the only things that we store into the StateStore
  // AFTER the initial saving on app-attempt-start
  // These fields can be visible from outside only after they are saved in
  // StateStore
  String diags = null;

  // don't leave the tracking URL pointing to a non-existent AM
  setTrackingUrlToRMAppPage(stateToBeStored);
  String finalTrackingUrl = getOriginalTrackingUrl();
  FinalApplicationStatus finalStatus = null;
  int exitStatus = ContainerExitStatus.INVALID;
  switch (event.getType()) {
  case LAUNCH_FAILED:
    RMAppAttemptLaunchFailedEvent launchFaileEvent =
        (RMAppAttemptLaunchFailedEvent) event;
    diags = launchFaileEvent.getMessage();
    break;
  case REGISTERED:
    diags = getUnexpectedAMRegisteredDiagnostics();
    break;
  case UNREGISTERED:
    RMAppAttemptUnregistrationEvent unregisterEvent =
        (RMAppAttemptUnregistrationEvent) event;
    diags = unregisterEvent.getDiagnostics();
    // reset finalTrackingUrl to url sent by am
    finalTrackingUrl = sanitizeTrackingUrl(unregisterEvent.getFinalTrackingUrl());
    finalStatus = unregisterEvent.getFinalApplicationStatus();
    break;
  case CONTAINER_FINISHED:
    RMAppAttemptContainerFinishedEvent finishEvent =
        (RMAppAttemptContainerFinishedEvent) event;
    diags = getAMContainerCrashedDiagnostics(finishEvent);
    exitStatus = finishEvent.getContainerStatus().getExitStatus();
    break;
  case KILL:
    break;
  case EXPIRE:
    diags = getAMExpiredDiagnostics(event);
    break;
  default:
    break;
  }
  AggregateAppResourceUsage resUsage =
      this.attemptMetrics.getAggregateAppResourceUsage();
  RMStateStore rmStore = rmContext.getStateStore();
  setFinishTime(System.currentTimeMillis());

  ApplicationAttemptStateData attemptState =
      ApplicationAttemptStateData.newInstance(
          applicationAttemptId,  getMasterContainer(),
          rmStore.getCredentialsFromAppAttempt(this),
          startTime, stateToBeStored, finalTrackingUrl, diags,
          finalStatus, exitStatus,
        getFinishTime(), resUsage.getMemorySeconds(),
        resUsage.getVcoreSeconds());
  LOG.info("Updating application attempt " + applicationAttemptId
      + " with final state: " + targetedFinalState + ", and exit status: "
      + exitStatus);
  rmStore.updateApplicationAttemptState(attemptState);
}
 

@Before
public void setUp() throws Exception {
  conf = new YarnConfiguration();
  AuthenticationMethod authMethod = AuthenticationMethod.SIMPLE;
  if (isSecurityEnabled) {
    authMethod = AuthenticationMethod.KERBEROS;
  }
  SecurityUtil.setAuthenticationMethod(authMethod, conf);
  UserGroupInformation.setConfiguration(conf);

  rmDispatcher = new DrainDispatcher();
  ContainerAllocationExpirer containerAllocationExpirer = 
      mock(ContainerAllocationExpirer.class);
  AMLivelinessMonitor amLivelinessMonitor = mock(AMLivelinessMonitor.class);
  AMLivelinessMonitor amFinishingMonitor = mock(AMLivelinessMonitor.class);
  store = mock(RMStateStore.class);
  writer = mock(RMApplicationHistoryWriter.class);
  DelegationTokenRenewer renewer = mock(DelegationTokenRenewer.class);
  RMContext realRMContext = 
      new RMContextImpl(rmDispatcher,
        containerAllocationExpirer, amLivelinessMonitor, amFinishingMonitor,
        renewer, new AMRMTokenSecretManager(conf, this.rmContext),
        new RMContainerTokenSecretManager(conf),
        new NMTokenSecretManagerInRM(conf),
        new ClientToAMTokenSecretManagerInRM(),
        writer);
  ((RMContextImpl)realRMContext).setStateStore(store);
  publisher = mock(SystemMetricsPublisher.class);
  ((RMContextImpl)realRMContext).setSystemMetricsPublisher(publisher);

  this.rmContext = spy(realRMContext);

  ResourceScheduler resourceScheduler = mock(ResourceScheduler.class);
  doReturn(null).when(resourceScheduler)
            .getAppResourceUsageReport((ApplicationAttemptId)Matchers.any());
  doReturn(resourceScheduler).when(rmContext).getScheduler();

  rmDispatcher.register(RMAppAttemptEventType.class,
      new TestApplicationAttemptEventDispatcher(this.rmContext));

  rmDispatcher.register(RMAppEventType.class,
      new TestApplicationEventDispatcher(rmContext));
  
  rmDispatcher.register(RMAppManagerEventType.class,
      new TestApplicationManagerEventDispatcher());
  
  schedulerDispatcher = new TestSchedulerEventDispatcher();
  rmDispatcher.register(SchedulerEventType.class,
      schedulerDispatcher);
  
  rmDispatcher.init(conf);
  rmDispatcher.start();
}
 

@SuppressWarnings("deprecation")
@Before
public void setUp() throws Exception {
  AuthenticationMethod authMethod = AuthenticationMethod.SIMPLE;
  if (isSecurityEnabled) {
    authMethod = AuthenticationMethod.KERBEROS;
  }
  SecurityUtil.setAuthenticationMethod(authMethod, conf);
  UserGroupInformation.setConfiguration(conf);
  InlineDispatcher rmDispatcher = new InlineDispatcher();

  ContainerAllocationExpirer containerAllocationExpirer =
      mock(ContainerAllocationExpirer.class);
  amLivelinessMonitor = mock(AMLivelinessMonitor.class);
  amFinishingMonitor = mock(AMLivelinessMonitor.class);
  writer = mock(RMApplicationHistoryWriter.class);
  MasterKeyData masterKeyData = amRMTokenManager.createNewMasterKey();
  when(amRMTokenManager.getMasterKey()).thenReturn(masterKeyData);
  rmContext =
      new RMContextImpl(rmDispatcher,
        containerAllocationExpirer, amLivelinessMonitor, amFinishingMonitor,
        null, amRMTokenManager,
        new RMContainerTokenSecretManager(conf),
        nmTokenManager,
        clientToAMTokenManager,
        writer);
  
  store = mock(RMStateStore.class);
  ((RMContextImpl) rmContext).setStateStore(store);
  publisher = mock(SystemMetricsPublisher.class);
  ((RMContextImpl) rmContext).setSystemMetricsPublisher(publisher);
  
  scheduler = mock(YarnScheduler.class);
  masterService = mock(ApplicationMasterService.class);
  applicationMasterLauncher = mock(ApplicationMasterLauncher.class);
  
  rmDispatcher.register(RMAppAttemptEventType.class,
      new TestApplicationAttemptEventDispatcher());

  rmDispatcher.register(RMAppEventType.class,
      new TestApplicationEventDispatcher());
  
  rmDispatcher.register(SchedulerEventType.class, 
      new TestSchedulerEventDispatcher());
  
  rmDispatcher.register(AMLauncherEventType.class, 
      new TestAMLauncherEventDispatcher());

  rmnodeEventHandler = mock(RMNodeImpl.class);
  rmDispatcher.register(RMNodeEventType.class, rmnodeEventHandler);

  rmDispatcher.init(conf);
  rmDispatcher.start();
  

  ApplicationId applicationId = MockApps.newAppID(appId++);
  ApplicationAttemptId applicationAttemptId =
      ApplicationAttemptId.newInstance(applicationId, 0);

  resourceScheduler = mock(ResourceScheduler.class);

  ApplicationResourceUsageReport appResUsgRpt =
      mock(ApplicationResourceUsageReport.class);
  when(appResUsgRpt.getMemorySeconds()).thenReturn(0L);
  when(appResUsgRpt.getVcoreSeconds()).thenReturn(0L);
  when(resourceScheduler
      .getAppResourceUsageReport((ApplicationAttemptId)Matchers.any()))
   .thenReturn(appResUsgRpt);
  spyRMContext = spy(rmContext);
  Mockito.doReturn(resourceScheduler).when(spyRMContext).getScheduler();


  final String user = MockApps.newUserName();
  final String queue = MockApps.newQueue();
  submissionContext = mock(ApplicationSubmissionContext.class);
  when(submissionContext.getQueue()).thenReturn(queue);
  Resource resource = BuilderUtils.newResource(1536, 1);
  ContainerLaunchContext amContainerSpec =
      BuilderUtils.newContainerLaunchContext(null, null,
          null, null, null, null);
  when(submissionContext.getAMContainerSpec()).thenReturn(amContainerSpec);
  when(submissionContext.getResource()).thenReturn(resource);

  unmanagedAM = false;
  
  application = mock(RMAppImpl.class);
  applicationAttempt =
      new RMAppAttemptImpl(applicationAttemptId, spyRMContext, scheduler,
          masterService, submissionContext, new Configuration(), false,
          BuilderUtils.newResourceRequest(
              RMAppAttemptImpl.AM_CONTAINER_PRIORITY, ResourceRequest.ANY,
              submissionContext.getResource(), 1));

  when(application.getCurrentAppAttempt()).thenReturn(applicationAttempt);
  when(application.getApplicationId()).thenReturn(applicationId);
  spyRMContext.getRMApps().put(application.getApplicationId(), application);

  testAppAttemptNewState();
}
 

public MyResourceManager(Configuration conf, RMStateStore store) {
  super(conf, store);
}
 

public MyResourceManager2(Configuration conf, RMStateStore store) {
  super(conf, store);
}
 

public MyResourceManager(Configuration conf, RMStateStore store) {
  super(conf, store);
}
 

@Override
public void recover(RMStateStore.RMState rmState) throws Exception {
  scheduler.recover(rmState);
}
 

public void setStateStore(RMStateStore stateStore) {
  this.stateStore = stateStore;
}
 

@Override
public RMStateStore getStateStore() {
  return stateStore;
}
 

@BeforeClass
public static void setup() throws InterruptedException, IOException {
  RMStateStore store = RMStateStoreFactory.getStore(conf);
  conf.setBoolean(YarnConfiguration.YARN_ACL_ENABLE, true);
  AccessControlList adminACL = new AccessControlList("");
  adminACL.addGroup(SUPER_GROUP);
  conf.set(YarnConfiguration.YARN_ADMIN_ACL, adminACL.getAclString());
  resourceManager = new MockRM(conf) {

    @Override
    protected QueueACLsManager createQueueACLsManager(
        ResourceScheduler scheduler,
        Configuration conf) {
      QueueACLsManager mockQueueACLsManager = mock(QueueACLsManager.class);
      when(mockQueueACLsManager.checkAccess(any(UserGroupInformation.class),
          any(QueueACL.class), anyString())).thenAnswer(new Answer() {
        public Object answer(InvocationOnMock invocation) {
          return isQueueUser;
        }
      });
      return mockQueueACLsManager;
    }

    protected ClientRMService createClientRMService() {
      return new ClientRMService(getRMContext(), this.scheduler,
          this.rmAppManager, this.applicationACLsManager,
          this.queueACLsManager, null);
    };
  };
  new Thread() {
    public void run() {
      UserGroupInformation.createUserForTesting(ENEMY, new String[] {});
      UserGroupInformation.createUserForTesting(FRIEND,
          new String[] { FRIENDLY_GROUP });
      UserGroupInformation.createUserForTesting(SUPER_USER,
          new String[] { SUPER_GROUP });
      resourceManager.start();
    };
  }.start();
  int waitCount = 0;
  while (resourceManager.getServiceState() == STATE.INITED
      && waitCount++ < 60) {
    LOG.info("Waiting for RM to start...");
    Thread.sleep(1500);
  }
  if (resourceManager.getServiceState() != STATE.STARTED) {
    // RM could have failed.
    throw new IOException(
        "ResourceManager failed to start. Final state is "
            + resourceManager.getServiceState());
  }

  UserGroupInformation owner = UserGroupInformation
      .createRemoteUser(APP_OWNER);
  rmClient = owner.doAs(new PrivilegedExceptionAction<ApplicationClientProtocol>() {
    @Override
    public ApplicationClientProtocol run() throws Exception {
      return (ApplicationClientProtocol) rpc.getProxy(ApplicationClientProtocol.class,
          rmAddress, conf);
    }
  });
}