Java Code Examples for org.apache.hadoop.yarn.conf.YarnConfiguration#setInt()

@Before
public void setup() throws IOException {
  failoverThread = null;
  keepRunning = true;
  conf = new YarnConfiguration();
  conf.setBoolean(YarnConfiguration.RM_HA_ENABLED, true);
  conf.setInt(YarnConfiguration.CLIENT_FAILOVER_MAX_ATTEMPTS, 5);
  conf.set(YarnConfiguration.RM_HA_IDS, RM1_NODE_ID + "," + RM2_NODE_ID);
  setRpcAddressForRM(RM1_NODE_ID, RM1_PORT_BASE);
  setRpcAddressForRM(RM2_NODE_ID, RM2_PORT_BASE);

  conf.setLong(YarnConfiguration.CLIENT_FAILOVER_SLEEPTIME_BASE_MS, 100L);

  conf.setBoolean(YarnConfiguration.YARN_MINICLUSTER_FIXED_PORTS, true);
  conf.setBoolean(YarnConfiguration.YARN_MINICLUSTER_USE_RPC, true);
}
 

@Override
public RMStateStore getRMStateStore() throws Exception {
  YarnConfiguration conf = new YarnConfiguration();
  conf.set(YarnConfiguration.FS_RM_STATE_STORE_URI,
      workingDirPathURI.toString());
  conf.set(YarnConfiguration.FS_RM_STATE_STORE_RETRY_POLICY_SPEC,
          "100,6000");
  conf.setInt(YarnConfiguration.FS_RM_STATE_STORE_NUM_RETRIES, 8);
  conf.setLong(YarnConfiguration.FS_RM_STATE_STORE_RETRY_INTERVAL_MS,
          900L);
  if (adminCheckEnable) {
    conf.setBoolean(
      YarnConfiguration.YARN_INTERMEDIATE_DATA_ENCRYPTION, true);
  }
  this.store = new TestFileSystemRMStore(conf);
  Assert.assertEquals(store.getNumRetries(), 8);
  Assert.assertEquals(store.getRetryInterval(), 900L);
  return store;
}
 

@Test  
public void testNonMinZeroResourcesSettings() throws IOException {
  scheduler = new FairScheduler();
  YarnConfiguration conf = new YarnConfiguration();
  conf.setInt(YarnConfiguration.RM_SCHEDULER_MINIMUM_ALLOCATION_MB, 256);
  conf.setInt(YarnConfiguration.RM_SCHEDULER_MINIMUM_ALLOCATION_VCORES, 1);
  conf.setInt(YarnConfiguration.RM_SCHEDULER_MINIMUM_ALLOCATION_GCORES, 1);
  conf.setInt(
    FairSchedulerConfiguration.RM_SCHEDULER_INCREMENT_ALLOCATION_MB, 512);
  conf.setInt(
    FairSchedulerConfiguration.RM_SCHEDULER_INCREMENT_ALLOCATION_VCORES, 2);
  conf.setInt(
    FairSchedulerConfiguration.RM_SCHEDULER_INCREMENT_ALLOCATION_GCORES, 1);
  scheduler.init(conf);
  scheduler.reinitialize(conf, null);
  Assert.assertEquals(256, scheduler.getMinimumResourceCapability().getMemory());
  Assert.assertEquals(1, scheduler.getMinimumResourceCapability().getVirtualCores());
  Assert.assertEquals(1, scheduler.getMinimumResourceCapability().getGpuCores());
  Assert.assertEquals(512, scheduler.getIncrementResourceCapability().getMemory());
  Assert.assertEquals(2, scheduler.getIncrementResourceCapability().getVirtualCores());
  Assert.assertEquals(1, scheduler.getIncrementResourceCapability().getGpuCores());
}
 

@Before
public void setup() throws IOException {
  conf = new YarnConfiguration();
  conf.setBoolean(YarnConfiguration.RM_HA_ENABLED, true);
  conf.setBoolean(YarnConfiguration.AUTO_FAILOVER_ENABLED, true);
  conf.setBoolean(YarnConfiguration.AUTO_FAILOVER_EMBEDDED, true);
  conf.set(YarnConfiguration.RM_CLUSTER_ID, "yarn-test-cluster");
  conf.set(YarnConfiguration.RM_ZK_ADDRESS, hostPort);
  conf.setInt(YarnConfiguration.RM_ZK_TIMEOUT_MS, 2000);

  conf.set(YarnConfiguration.RM_HA_IDS, RM1_NODE_ID + "," + RM2_NODE_ID);
  conf.set(YarnConfiguration.RM_HA_ID, RM1_NODE_ID);
  setRpcAddressForRM(RM1_NODE_ID, RM1_PORT_BASE);
  setRpcAddressForRM(RM2_NODE_ID, RM2_PORT_BASE);

  conf.setLong(YarnConfiguration.CLIENT_FAILOVER_SLEEPTIME_BASE_MS, 100L);

  callbackCalled = new AtomicBoolean(false);
}
 

@Before
public void setup() throws IOException {
  conf = new YarnConfiguration();
  conf.setBoolean(YarnConfiguration.RM_HA_ENABLED, true);
  conf.setBoolean(YarnConfiguration.AUTO_FAILOVER_ENABLED, true);
  conf.setBoolean(YarnConfiguration.AUTO_FAILOVER_EMBEDDED, true);
  conf.set(YarnConfiguration.RM_CLUSTER_ID, "yarn-test-cluster");
  conf.set(YarnConfiguration.RM_ZK_ADDRESS, hostPort);
  conf.setInt(YarnConfiguration.RM_ZK_TIMEOUT_MS, 2000);

  conf.set(YarnConfiguration.RM_HA_IDS, RM1_NODE_ID + "," + RM2_NODE_ID);
  conf.set(YarnConfiguration.RM_HA_ID, RM1_NODE_ID);
  setRpcAddressForRM(RM1_NODE_ID, RM1_PORT_BASE);
  setRpcAddressForRM(RM2_NODE_ID, RM2_PORT_BASE);

  conf.setLong(YarnConfiguration.CLIENT_FAILOVER_SLEEPTIME_BASE_MS, 100L);

  callbackCalled = new AtomicBoolean(false);
}
 

@BeforeClass
public static void setup() throws Exception {
  YarnConfiguration conf = new YarnConfiguration();
  conf.setBoolean(YarnConfiguration.TIMELINE_SERVICE_ENABLED, true);
  conf.setBoolean(YarnConfiguration.RM_SYSTEM_METRICS_PUBLISHER_ENABLED, true);
  conf.setClass(YarnConfiguration.TIMELINE_SERVICE_STORE,
      MemoryTimelineStore.class, TimelineStore.class);
  conf.setClass(YarnConfiguration.TIMELINE_SERVICE_STATE_STORE_CLASS,
      MemoryTimelineStateStore.class, TimelineStateStore.class);
  conf.setInt(
      YarnConfiguration.RM_SYSTEM_METRICS_PUBLISHER_DISPATCHER_POOL_SIZE,
      2);

  timelineServer = new ApplicationHistoryServer();
  timelineServer.init(conf);
  timelineServer.start();
  store = timelineServer.getTimelineStore();

  metricsPublisher = new SystemMetricsPublisher();
  metricsPublisher.init(conf);
  metricsPublisher.start();
}
 

@Test
public void testMaximimumAllocationMemory() throws Exception {
  final int node1MaxMemory = 15 * 1024;
  final int node2MaxMemory = 5 * 1024;
  final int node3MaxMemory = 6 * 1024;
  final int configuredMaxMemory = 10 * 1024;
  configureScheduler();
  YarnConfiguration conf = getConf();
  conf.setInt(YarnConfiguration.RM_SCHEDULER_MAXIMUM_ALLOCATION_MB,
      configuredMaxMemory);
  conf.setLong(
      YarnConfiguration.RM_WORK_PRESERVING_RECOVERY_SCHEDULING_WAIT_MS,
      1000 * 1000);
  MockRM rm = new MockRM(conf);
  try {
    rm.start();
    testMaximumAllocationMemoryHelper(
        (AbstractYarnScheduler) rm.getResourceScheduler(),
        node1MaxMemory, node2MaxMemory, node3MaxMemory,
        configuredMaxMemory, configuredMaxMemory, configuredMaxMemory,
        configuredMaxMemory, configuredMaxMemory, configuredMaxMemory);
  } finally {
    rm.stop();
  }

  conf.setLong(
      YarnConfiguration.RM_WORK_PRESERVING_RECOVERY_SCHEDULING_WAIT_MS,
      0);
  rm = new MockRM(conf);
  try {
    rm.start();
    testMaximumAllocationMemoryHelper(
        (AbstractYarnScheduler) rm.getResourceScheduler(),
        node1MaxMemory, node2MaxMemory, node3MaxMemory,
        configuredMaxMemory, configuredMaxMemory, configuredMaxMemory,
        node2MaxMemory, node3MaxMemory, node2MaxMemory);
  } finally {
    rm.stop();
  }
}
 

private YarnConfiguration createNMConfig() {
  YarnConfiguration conf = new YarnConfiguration();
  conf.setInt(YarnConfiguration.NM_PMEM_MB, 5 * 1024); // 5GB
  conf.set(YarnConfiguration.NM_ADDRESS, "127.0.0.1:12345");
  conf.set(YarnConfiguration.NM_LOCALIZER_ADDRESS, "127.0.0.1:12346");
  conf.set(YarnConfiguration.NM_LOG_DIRS, logsDir.getAbsolutePath());
  conf.set(YarnConfiguration.NM_LOCAL_DIRS, nmLocalDir.getAbsolutePath());
  conf.setLong(YarnConfiguration.NM_LOG_RETAIN_SECONDS, 1);
  return conf;
}
 

@Test
public void testIncrementFileCountForPath() {
  YarnConfiguration conf = new YarnConfiguration();
  conf.setInt(YarnConfiguration.NM_LOCAL_CACHE_MAX_FILES_PER_DIRECTORY,
      LocalCacheDirectoryManager.DIRECTORIES_PER_LEVEL + 2);
  LocalCacheDirectoryManager mgr = new LocalCacheDirectoryManager(conf);
  final String rootPath = "";
  mgr.incrementFileCountForPath(rootPath);
  Assert.assertEquals(rootPath, mgr.getRelativePathForLocalization());
  Assert.assertFalse("root dir should be full",
      rootPath.equals(mgr.getRelativePathForLocalization()));
  // finish filling the other directory
  mgr.getRelativePathForLocalization();
  // free up space in the root dir
  mgr.decrementFileCountForPath(rootPath);
  mgr.decrementFileCountForPath(rootPath);
  Assert.assertEquals(rootPath, mgr.getRelativePathForLocalization());
  Assert.assertEquals(rootPath, mgr.getRelativePathForLocalization());
  String otherDir = mgr.getRelativePathForLocalization();
  Assert.assertFalse("root dir should be full", otherDir.equals(rootPath));

  final String deepDir0 = "d/e/e/p/0";
  final String deepDir1 = "d/e/e/p/1";
  final String deepDir2 = "d/e/e/p/2";
  final String deepDir3 = "d/e/e/p/3";
  mgr.incrementFileCountForPath(deepDir0);
  Assert.assertEquals(otherDir, mgr.getRelativePathForLocalization());
  Assert.assertEquals(deepDir0, mgr.getRelativePathForLocalization());
  Assert.assertEquals("total dir count incorrect after increment",
      deepDir1, mgr.getRelativePathForLocalization());
  mgr.incrementFileCountForPath(deepDir2);
  mgr.incrementFileCountForPath(deepDir1);
  mgr.incrementFileCountForPath(deepDir2);
  Assert.assertEquals(deepDir3, mgr.getRelativePathForLocalization());
}
 

private YarnConfiguration getConfForCM(boolean pMemEnabled,
    boolean vMemEnabled, int nmPmem, float vMemToPMemRatio) {
  YarnConfiguration conf = new YarnConfiguration();
  conf.setInt(YarnConfiguration.NM_PMEM_MB, nmPmem);
  conf.setBoolean(YarnConfiguration.NM_PMEM_CHECK_ENABLED, pMemEnabled);
  conf.setBoolean(YarnConfiguration.NM_VMEM_CHECK_ENABLED, vMemEnabled);
  conf.setFloat(YarnConfiguration.NM_VMEM_PMEM_RATIO, vMemToPMemRatio);
  return conf;
}
 

@BeforeClass
public static void setup() throws Exception {
  // start minicluster
  conf = new YarnConfiguration();
  conf.setLong(
    YarnConfiguration.RM_AMRM_TOKEN_MASTER_KEY_ROLLING_INTERVAL_SECS,
    rolling_interval_sec);
  conf.setLong(YarnConfiguration.RM_AM_EXPIRY_INTERVAL_MS, am_expire_ms);
  conf.setInt(YarnConfiguration.RM_NM_HEARTBEAT_INTERVAL_MS, 100);
  conf.setLong(YarnConfiguration.NM_LOG_RETAIN_SECONDS, 1);
  yarnCluster = new MiniYARNCluster(TestAMRMClient.class.getName(), nodeCount, 1, 1);
  yarnCluster.init(conf);
  yarnCluster.start();

  // start rm client
  yarnClient = YarnClient.createYarnClient();
  yarnClient.init(conf);
  yarnClient.start();

  // get node info
  nodeReports = yarnClient.getNodeReports(NodeState.RUNNING);
  
  priority = Priority.newInstance(1);
  priority2 = Priority.newInstance(2);
  capability = Resource.newInstance(1024, 1, 1);

  node = nodeReports.get(0).getNodeId().getHost();
  rack = nodeReports.get(0).getRackName();
  nodes = new String[]{ node };
  racks = new String[]{ rack };
}
 

public YarnConfiguration createRegistryConfiguration() {
  YarnConfiguration conf = new YarnConfiguration();
  conf.setInt(RegistryConstants.KEY_REGISTRY_ZK_CONNECTION_TIMEOUT, 1000);
  conf.setInt(RegistryConstants.KEY_REGISTRY_ZK_RETRY_INTERVAL, 500);
  conf.setInt(RegistryConstants.KEY_REGISTRY_ZK_RETRY_TIMES, 10);
  conf.setInt(RegistryConstants.KEY_REGISTRY_ZK_RETRY_CEILING, 10);
  conf.set(RegistryConstants.KEY_REGISTRY_ZK_QUORUM,
      zookeeper.getConnectionString());
  return conf;
}
 

@Test (timeout = 20000)
public void testZKClientRetry() throws Exception {
  TestZKClient zkClientTester = new TestZKClient();
  final String path = "/test";
  YarnConfiguration conf = new YarnConfiguration();
  conf.setInt(YarnConfiguration.RM_ZK_TIMEOUT_MS, ZK_TIMEOUT_MS);
  conf.setLong(YarnConfiguration.RM_ZK_RETRY_INTERVAL_MS, 100);
  final ZKRMStateStore store =
      (ZKRMStateStore) zkClientTester.getRMStateStore(conf);
  TestDispatcher dispatcher = new TestDispatcher();
  store.setRMDispatcher(dispatcher);
  final AtomicBoolean assertionFailedInThread = new AtomicBoolean(false);

  stopServer();
  Thread clientThread = new Thread() {
    @Override
    public void run() {
      try {
        store.getDataWithRetries(path, true);
      } catch (Exception e) {
        e.printStackTrace();
        assertionFailedInThread.set(true);
      }
    }
  };
  Thread.sleep(2000);
  startServer();
  clientThread.join();
  Assert.assertFalse(assertionFailedInThread.get());
}
 

private YarnConfiguration createNMConfig() {
  YarnConfiguration conf = new YarnConfiguration();
  conf.setInt(YarnConfiguration.NM_PMEM_MB, 5*1024); // 5GB
  conf.set(YarnConfiguration.NM_ADDRESS, "127.0.0.1:12345");
  conf.set(YarnConfiguration.NM_LOCALIZER_ADDRESS, "127.0.0.1:12346");
  conf.set(YarnConfiguration.NM_LOG_DIRS, logsDir.getAbsolutePath());
  conf.set(YarnConfiguration.NM_REMOTE_APP_LOG_DIR,
    remoteLogsDir.getAbsolutePath());
  conf.set(YarnConfiguration.NM_LOCAL_DIRS, nmLocalDir.getAbsolutePath());
  conf.setLong(YarnConfiguration.NM_LOG_RETAIN_SECONDS, 1);
  return conf;
}
 

@Test
public void testMaximimumAllocationMemory() throws Exception {
  final int node1MaxMemory = 15 * 1024;
  final int node2MaxMemory = 5 * 1024;
  final int node3MaxMemory = 6 * 1024;
  final int configuredMaxMemory = 10 * 1024;
  configureScheduler();
  YarnConfiguration conf = getConf();
  conf.setInt(YarnConfiguration.RM_SCHEDULER_MAXIMUM_ALLOCATION_MB,
      configuredMaxMemory);
  conf.setLong(
      YarnConfiguration.RM_WORK_PRESERVING_RECOVERY_SCHEDULING_WAIT_MS,
      1000 * 1000);
  MockRM rm = new MockRM(conf);
  try {
    rm.start();
    testMaximumAllocationMemoryHelper(
        (AbstractYarnScheduler) rm.getResourceScheduler(),
        node1MaxMemory, node2MaxMemory, node3MaxMemory,
        configuredMaxMemory, configuredMaxMemory, configuredMaxMemory,
        configuredMaxMemory, configuredMaxMemory, configuredMaxMemory);
  } finally {
    rm.stop();
  }

  conf.setLong(
      YarnConfiguration.RM_WORK_PRESERVING_RECOVERY_SCHEDULING_WAIT_MS,
      0);
  rm = new MockRM(conf);
  try {
    rm.start();
    testMaximumAllocationMemoryHelper(
        (AbstractYarnScheduler) rm.getResourceScheduler(),
        node1MaxMemory, node2MaxMemory, node3MaxMemory,
        configuredMaxMemory, configuredMaxMemory, configuredMaxMemory,
        node2MaxMemory, node3MaxMemory, node2MaxMemory);
  } finally {
    rm.stop();
  }
}
 

private void configForTesting(String classPathRoot) throws URISyntaxException {
    yarnSiteConfigURL = Thread.currentThread().getContextClassLoader().getResource("yarn-site.xml");
    if (yarnSiteConfigURL == null) {
        throw new RuntimeException("Could not find 'yarn-site.xml' file in classpath");
    } else {
        LOG.info("Found 'yarn-site.xml' at "+ yarnSiteConfigURL.toURI().toString());
    }

    conf = new YarnConfiguration();
    conf.set(FileSystem.FS_DEFAULT_NAME_KEY, "file:///");

    keytab = classPathRoot.substring(0, classPathRoot.lastIndexOf('/'))+"/splice.keytab";
    if (secure) {
        conf.set("hadoop.security.authentication", "kerberos");
        conf.set("yarn.resourcemanager.principal", "yarn/[email protected]");
        conf.set("yarn.resourcemanager.keytab", keytab);
        conf.set("yarn.nodemanager.principal", "yarn/[email protected]");
        conf.set("yarn.nodemanager.keytab", keytab);
    }
    conf.setDouble("yarn.nodemanager.resource.io-spindles",2.0);
    conf.set("fs.default.name", "file:///");
    conf.set("yarn.nodemanager.container-executor.class","org.apache.hadoop.yarn.server.nodemanager.DefaultContainerExecutor");
    System.setProperty("zookeeper.sasl.client", "false");
    System.setProperty("zookeeper.sasl.serverconfig", "fake");

    conf.setInt(YarnConfiguration.RM_NM_HEARTBEAT_INTERVAL_MS, DEFAULT_HEARTBEAT_INTERVAL);
    conf.setInt(YarnConfiguration.RM_SCHEDULER_MINIMUM_ALLOCATION_MB, 128);
    conf.setClass(YarnConfiguration.RM_SCHEDULER, FifoScheduler.class, ResourceScheduler.class);
    conf.set("yarn.application.classpath", new File(yarnSiteConfigURL.getPath()).getParent());
}
 

@Test(timeout = 10000)
public void testResetTimer() throws Exception {
  YarnConfiguration conf = new YarnConfiguration();
  UserGroupInformation.setConfiguration(conf);
  conf.set(YarnConfiguration.RECOVERY_ENABLED, "true");
  conf.set(YarnConfiguration.RM_STORE, MemoryRMStateStore.class.getName());
  conf.setBoolean(YarnConfiguration.RM_WORK_PRESERVING_RECOVERY_ENABLED, true);
  conf.setInt(YarnConfiguration.RM_AM_EXPIRY_INTERVAL_MS, 6000);
  final ControlledClock clock = new ControlledClock(new SystemClock());
  clock.setTime(0);
  MemoryRMStateStore memStore = new MemoryRMStateStore() {
    @Override
    public synchronized RMState loadState() throws Exception {
      clock.setTime(8000);
      return super.loadState();
    }
  };
  memStore.init(conf);
  final ApplicationAttemptId attemptId = mock(ApplicationAttemptId.class);
  final Dispatcher dispatcher = mock(Dispatcher.class);
  final boolean[] expired = new boolean[]{false};
  final AMLivelinessMonitor monitor = new AMLivelinessMonitor(
      dispatcher, clock) {
    @Override
    protected void expire(ApplicationAttemptId id) {
      Assert.assertEquals(id, attemptId);
      expired[0] = true;
    }
  };
  monitor.register(attemptId);
  MockRM rm = new MockRM(conf, memStore) {
    @Override
    protected AMLivelinessMonitor createAMLivelinessMonitor() {
      return monitor;
    }
  };
  rm.start();
  // make sure that monitor has started
  while (monitor.getServiceState() != Service.STATE.STARTED) {
    Thread.sleep(100);
  }
  // expired[0] would be set to true without resetTimer
  Assert.assertFalse(expired[0]);
  rm.stop();
}
 

@Test(timeout = 50000)
public void testRMRestartOrFailoverNotCountedForAMFailures()
    throws Exception {
  YarnConfiguration conf = new YarnConfiguration();
  conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class,
    ResourceScheduler.class);
  conf.setBoolean(YarnConfiguration.RECOVERY_ENABLED, true);
  conf.setBoolean(YarnConfiguration.RM_WORK_PRESERVING_RECOVERY_ENABLED, false);

  conf.set(YarnConfiguration.RM_STORE, MemoryRMStateStore.class.getName());
  // explicitly set max-am-retry count as 1.
  conf.setInt(YarnConfiguration.RM_AM_MAX_ATTEMPTS, 1);
  MemoryRMStateStore memStore = new MemoryRMStateStore();
  memStore.init(conf);

  MockRM rm1 = new MockRM(conf, memStore);
  rm1.start();
  MockNM nm1 =
      new MockNM("127.0.0.1:1234", 8000, rm1.getResourceTrackerService());
  nm1.registerNode();
  RMApp app1 = rm1.submitApp(200);
  // AM should be restarted even though max-am-attempt is 1.
  MockAM am1 = MockRM.launchAndRegisterAM(app1, rm1, nm1);
  RMAppAttempt attempt1 = app1.getCurrentAppAttempt();
  Assert.assertTrue(((RMAppAttemptImpl) attempt1).mayBeLastAttempt());

  // Restart rm.
  MockRM rm2 = new MockRM(conf, memStore);
  rm2.start();
  ApplicationStateData appState =
      memStore.getState().getApplicationState().get(app1.getApplicationId());
  // re-register the NM
  nm1.setResourceTrackerService(rm2.getResourceTrackerService());
  NMContainerStatus status = Records.newRecord(NMContainerStatus.class);
  status
    .setContainerExitStatus(ContainerExitStatus.KILLED_BY_RESOURCEMANAGER);
  status.setContainerId(attempt1.getMasterContainer().getId());
  status.setContainerState(ContainerState.COMPLETE);
  status.setDiagnostics("");
  nm1.registerNode(Collections.singletonList(status), null);

  rm2.waitForState(attempt1.getAppAttemptId(), RMAppAttemptState.FAILED);
  Assert.assertEquals(ContainerExitStatus.KILLED_BY_RESOURCEMANAGER,
    appState.getAttempt(am1.getApplicationAttemptId())
      .getAMContainerExitStatus());
  // Will automatically start a new AppAttempt in rm2
  rm2.waitForState(app1.getApplicationId(), RMAppState.ACCEPTED);
  MockAM am2 =
      rm2.waitForNewAMToLaunchAndRegister(app1.getApplicationId(), 2, nm1);
  MockRM.finishAMAndVerifyAppState(app1, rm2, nm1, am2);
  RMAppAttempt attempt3 =
      rm2.getRMContext().getRMApps().get(app1.getApplicationId())
        .getCurrentAppAttempt();
  Assert.assertTrue(attempt3.shouldCountTowardsMaxAttemptRetry());
  Assert.assertEquals(ContainerExitStatus.INVALID,
    appState.getAttempt(am2.getApplicationAttemptId())
      .getAMContainerExitStatus());

  rm1.stop();
  rm2.stop();
}
 

@Test(timeout = 10000)
public void testResetTimer() throws Exception {
  YarnConfiguration conf = new YarnConfiguration();
  UserGroupInformation.setConfiguration(conf);
  conf.set(YarnConfiguration.RECOVERY_ENABLED, "true");
  conf.set(YarnConfiguration.RM_STORE, MemoryRMStateStore.class.getName());
  conf.setBoolean(YarnConfiguration.RM_WORK_PRESERVING_RECOVERY_ENABLED, true);
  conf.setInt(YarnConfiguration.RM_AM_EXPIRY_INTERVAL_MS, 6000);
  final ControlledClock clock = new ControlledClock(new SystemClock());
  clock.setTime(0);
  MemoryRMStateStore memStore = new MemoryRMStateStore() {
    @Override
    public synchronized RMState loadState() throws Exception {
      clock.setTime(8000);
      return super.loadState();
    }
  };
  memStore.init(conf);
  final ApplicationAttemptId attemptId = mock(ApplicationAttemptId.class);
  final Dispatcher dispatcher = mock(Dispatcher.class);
  final boolean[] expired = new boolean[]{false};
  final AMLivelinessMonitor monitor = new AMLivelinessMonitor(
      dispatcher, clock) {
    @Override
    protected void expire(ApplicationAttemptId id) {
      Assert.assertEquals(id, attemptId);
      expired[0] = true;
    }
  };
  monitor.register(attemptId);
  MockRM rm = new MockRM(conf, memStore) {
    @Override
    protected AMLivelinessMonitor createAMLivelinessMonitor() {
      return monitor;
    }
  };
  rm.start();
  // make sure that monitor has started
  while (monitor.getServiceState() != Service.STATE.STARTED) {
    Thread.sleep(100);
  }
  // expired[0] would be set to true without resetTimer
  Assert.assertFalse(expired[0]);
  rm.stop();
}
 

@Test
public void testUpdateMaxAllocationUsesTotal() throws IOException {
  final int configuredMaxVCores = 20;
  final int configuredMaxMemory = 10 * 1024;
  Resource configuredMaximumResource = Resource.newInstance
      (configuredMaxMemory, configuredMaxVCores);

  configureScheduler();
  YarnConfiguration conf = getConf();
  conf.setInt(YarnConfiguration.RM_SCHEDULER_MAXIMUM_ALLOCATION_VCORES,
      configuredMaxVCores);
  conf.setInt(YarnConfiguration.RM_SCHEDULER_MAXIMUM_ALLOCATION_MB,
      configuredMaxMemory);
  conf.setLong(
      YarnConfiguration.RM_WORK_PRESERVING_RECOVERY_SCHEDULING_WAIT_MS,
      0);

  MockRM rm = new MockRM(conf);
  try {
    rm.start();
    AbstractYarnScheduler scheduler = (AbstractYarnScheduler) rm
        .getResourceScheduler();

    Resource emptyResource = Resource.newInstance(0, 0);
    Resource fullResource1 = Resource.newInstance(1024, 5);
    Resource fullResource2 = Resource.newInstance(2048, 10);

    SchedulerNode mockNode1 = mock(SchedulerNode.class);
    when(mockNode1.getNodeID()).thenReturn(NodeId.newInstance("foo", 8080));
    when(mockNode1.getAvailableResource()).thenReturn(emptyResource);
    when(mockNode1.getTotalResource()).thenReturn(fullResource1);

    SchedulerNode mockNode2 = mock(SchedulerNode.class);
    when(mockNode1.getNodeID()).thenReturn(NodeId.newInstance("bar", 8081));
    when(mockNode2.getAvailableResource()).thenReturn(emptyResource);
    when(mockNode2.getTotalResource()).thenReturn(fullResource2);

    verifyMaximumResourceCapability(configuredMaximumResource, scheduler);

    scheduler.nodes = new HashMap<NodeId, SchedulerNode>();

    scheduler.nodes.put(mockNode1.getNodeID(), mockNode1);
    scheduler.updateMaximumAllocation(mockNode1, true);
    verifyMaximumResourceCapability(fullResource1, scheduler);

    scheduler.nodes.put(mockNode2.getNodeID(), mockNode2);
    scheduler.updateMaximumAllocation(mockNode2, true);
    verifyMaximumResourceCapability(fullResource2, scheduler);

    scheduler.nodes.remove(mockNode2.getNodeID());
    scheduler.updateMaximumAllocation(mockNode2, false);
    verifyMaximumResourceCapability(fullResource1, scheduler);

    scheduler.nodes.remove(mockNode1.getNodeID());
    scheduler.updateMaximumAllocation(mockNode1, false);
    verifyMaximumResourceCapability(configuredMaximumResource, scheduler);
  } finally {
    rm.stop();
  }
}