org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.TestCapacityScheduler.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.TestCapacityScheduler.java

Source

/**
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements.  See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership.  The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License.  You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity;

import com.google.common.base.Supplier;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.security.PrivilegedAction;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.BrokenBarrierException;
import java.util.concurrent.CyclicBarrier;

import io.hops.util.DBUtility;
import io.hops.util.RMStorageFactory;
import io.hops.util.YarnAPIStorageFactory;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.net.NetworkTopology;
import org.apache.hadoop.security.Credentials;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.security.token.Token;
import org.apache.hadoop.security.token.TokenIdentifier;
import org.apache.hadoop.yarn.LocalConfigurationProvider;
import org.apache.hadoop.yarn.api.ApplicationMasterProtocol;
import org.apache.hadoop.yarn.api.protocolrecords.AllocateRequest;
import org.apache.hadoop.yarn.api.protocolrecords.AllocateResponse;
import org.apache.hadoop.yarn.api.protocolrecords.RegisterApplicationMasterRequest;
import org.apache.hadoop.yarn.api.records.ApplicationAccessType;
import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
import org.apache.hadoop.yarn.api.records.ApplicationId;
import org.apache.hadoop.yarn.api.records.ApplicationResourceUsageReport;
import org.apache.hadoop.yarn.api.records.ApplicationSubmissionContext;
import org.apache.hadoop.yarn.api.records.Container;
import org.apache.hadoop.yarn.api.records.ContainerId;
import org.apache.hadoop.yarn.api.records.ContainerState;
import org.apache.hadoop.yarn.api.records.ContainerStatus;
import org.apache.hadoop.yarn.api.records.ContainerUpdateType;
import org.apache.hadoop.yarn.api.records.NodeId;
import org.apache.hadoop.yarn.api.records.NodeState;
import org.apache.hadoop.yarn.api.records.Priority;
import org.apache.hadoop.yarn.api.records.QueueInfo;
import org.apache.hadoop.yarn.api.records.QueueState;
import org.apache.hadoop.yarn.api.records.QueueUserACLInfo;
import org.apache.hadoop.yarn.api.records.Resource;
import org.apache.hadoop.yarn.api.records.ResourceOption;
import org.apache.hadoop.yarn.api.records.ResourceRequest;
import org.apache.hadoop.yarn.api.records.UpdateContainerRequest;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.hadoop.yarn.event.AsyncDispatcher;
import org.apache.hadoop.yarn.event.Dispatcher;
import org.apache.hadoop.yarn.event.Event;
import org.apache.hadoop.yarn.event.EventHandler;
import org.apache.hadoop.yarn.exceptions.YarnException;
import org.apache.hadoop.yarn.exceptions.YarnRuntimeException;
import org.apache.hadoop.yarn.factories.RecordFactory;
import org.apache.hadoop.yarn.factory.providers.RecordFactoryProvider;
import org.apache.hadoop.yarn.ipc.YarnRPC;
import org.apache.hadoop.yarn.server.api.protocolrecords.UpdateNodeResourceRequest;
import org.apache.hadoop.yarn.server.resourcemanager.AdminService;
import org.apache.hadoop.yarn.server.resourcemanager.Application;
import org.apache.hadoop.yarn.server.resourcemanager.MockAM;
import org.apache.hadoop.yarn.server.resourcemanager.MockNM;
import org.apache.hadoop.yarn.server.resourcemanager.MockNodes;
import org.apache.hadoop.yarn.server.resourcemanager.MockRM;
import org.apache.hadoop.yarn.server.resourcemanager.NodeManager;
import org.apache.hadoop.yarn.server.resourcemanager.RMContext;
import org.apache.hadoop.yarn.server.resourcemanager.RMContextImpl;
import org.apache.hadoop.yarn.server.resourcemanager.ResourceManager;
import org.apache.hadoop.yarn.server.resourcemanager.Task;
import org.apache.hadoop.yarn.server.resourcemanager.TestAMAuthorization.MockRMWithAMS;
import org.apache.hadoop.yarn.server.resourcemanager.TestAMAuthorization.MyContainerManager;
import org.apache.hadoop.yarn.server.resourcemanager.nodelabels.NullRMNodeLabelsManager;
import org.apache.hadoop.yarn.server.resourcemanager.nodelabels.RMNodeLabelsManager;
import org.apache.hadoop.yarn.server.resourcemanager.recovery.MemoryRMStateStore;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMApp;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMAppImpl;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMAppMetrics;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMAppState;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.RMAppAttempt;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.RMAppAttemptImpl;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.RMAppAttemptMetrics;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.RMAppAttemptState;
import org.apache.hadoop.yarn.server.resourcemanager.rmcontainer.RMContainer;
import org.apache.hadoop.yarn.server.resourcemanager.rmcontainer.RMContainerEvent;
import org.apache.hadoop.yarn.server.resourcemanager.rmcontainer.RMContainerEventType;
import org.apache.hadoop.yarn.server.resourcemanager.rmcontainer.RMContainerState;
import org.apache.hadoop.yarn.server.resourcemanager.rmnode.RMNode;
import org.apache.hadoop.yarn.server.resourcemanager.rmnode.RMNodeResourceUpdateEvent;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.AbstractYarnScheduler;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.Allocation;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.event.ContainerExpiredSchedulerEvent;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.ResourceScheduler;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.SchedulerApplication;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.SchedulerApplicationAttempt;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.SchedulerNode;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.SchedulerNodeReport;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.TestSchedulerUtils;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.YarnScheduler;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.common.fica.FiCaSchedulerApp;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.common.fica.FiCaSchedulerNode;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.event.AppAddedSchedulerEvent;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.event.AppAttemptAddedSchedulerEvent;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.event.AppAttemptRemovedSchedulerEvent;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.event.NodeAddedSchedulerEvent;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.event.NodeRemovedSchedulerEvent;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.event.NodeUpdateSchedulerEvent;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.event.SchedulerEvent;
import org.apache.hadoop.yarn.server.resourcemanager.scheduler.policy.FairOrderingPolicy;
import org.apache.hadoop.yarn.server.resourcemanager.security.ClientToAMTokenSecretManagerInRM;
import org.apache.hadoop.yarn.server.resourcemanager.security.NMTokenSecretManagerInRM;
import org.apache.hadoop.yarn.server.resourcemanager.security.RMContainerTokenSecretManager;
import org.apache.hadoop.yarn.server.resourcemanager.webapp.dao.CapacitySchedulerInfo;
import org.apache.hadoop.yarn.server.resourcemanager.webapp.dao.CapacitySchedulerLeafQueueInfo;
import org.apache.hadoop.yarn.server.resourcemanager.webapp.dao.CapacitySchedulerQueueInfo;
import org.apache.hadoop.yarn.server.resourcemanager.webapp.dao.CapacitySchedulerQueueInfoList;
import org.apache.hadoop.yarn.server.utils.BuilderUtils;
import org.apache.hadoop.yarn.util.resource.DefaultResourceCalculator;
import org.apache.hadoop.yarn.util.resource.DominantResourceCalculator;
import org.apache.hadoop.yarn.util.resource.DominantResourceCalculatorGPU;
import org.apache.hadoop.yarn.util.resource.Resources;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.mockito.Mockito;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Sets;
import org.apache.hadoop.test.GenericTestUtils;

public class TestCapacityScheduler {
    private static final Log LOG = LogFactory.getLog(TestCapacityScheduler.class);
    private final int GB = 1024;

    private static final String A = CapacitySchedulerConfiguration.ROOT + ".a";
    private static final String B = CapacitySchedulerConfiguration.ROOT + ".b";
    private static final String A1 = A + ".a1";
    private static final String A2 = A + ".a2";
    private static final String B1 = B + ".b1";
    private static final String B2 = B + ".b2";
    private static final String B3 = B + ".b3";
    private static float A_CAPACITY = 10.5f;
    private static float B_CAPACITY = 89.5f;
    private static final String P1 = CapacitySchedulerConfiguration.ROOT + ".p1";
    private static final String P2 = CapacitySchedulerConfiguration.ROOT + ".p2";
    private static final String X1 = P1 + ".x1";
    private static final String X2 = P1 + ".x2";
    private static final String Y1 = P2 + ".y1";
    private static final String Y2 = P2 + ".y2";
    private static float A1_CAPACITY = 30;
    private static float A2_CAPACITY = 70;
    private static float B1_CAPACITY = 79.2f;
    private static float B2_CAPACITY = 0.8f;
    private static float B3_CAPACITY = 20;

    private ResourceManager resourceManager = null;
    private RMContext mockContext;

    @Before
    public void setUp() throws Exception {
        resourceManager = new ResourceManager() {
            @Override
            protected RMNodeLabelsManager createNodeLabelManager() {
                RMNodeLabelsManager mgr = new NullRMNodeLabelsManager();
                mgr.init(getConfig());
                return mgr;
            }
        };
        CapacitySchedulerConfiguration csConf = new CapacitySchedulerConfiguration();
        setupQueueConfiguration(csConf);
        YarnConfiguration conf = new YarnConfiguration(csConf);
        conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);

        RMStorageFactory.setConfiguration(conf);
        YarnAPIStorageFactory.setConfiguration(conf);
        DBUtility.InitializeDB();
        resourceManager.init(conf);
        resourceManager.getRMContext().getContainerTokenSecretManager().rollMasterKey();
        resourceManager.getRMContext().getNMTokenSecretManager().rollMasterKey();
        ((AsyncDispatcher) resourceManager.getRMContext().getDispatcher()).start();
        mockContext = mock(RMContext.class);
        when(mockContext.getConfigurationProvider()).thenReturn(new LocalConfigurationProvider());
    }

    @After
    public void tearDown() throws Exception {
        if (resourceManager != null) {
            resourceManager.stop();
        }
    }

    @Test(timeout = 30000)
    public void testConfValidation() throws Exception {
        ResourceScheduler scheduler = new CapacityScheduler();
        scheduler.setRMContext(resourceManager.getRMContext());
        Configuration conf = new YarnConfiguration();
        conf.setInt(YarnConfiguration.RM_SCHEDULER_MINIMUM_ALLOCATION_MB, 2048);
        conf.setInt(YarnConfiguration.RM_SCHEDULER_MAXIMUM_ALLOCATION_MB, 1024);
        try {
            scheduler.reinitialize(conf, mockContext);
            fail("Exception is expected because the min memory allocation is"
                    + " larger than the max memory allocation.");
        } catch (YarnRuntimeException e) {
            // Exception is expected.
            assertTrue("The thrown exception is not the expected one.",
                    e.getMessage().startsWith("Invalid resource scheduler memory"));
        }

        conf = new YarnConfiguration();
        conf.setInt(YarnConfiguration.RM_SCHEDULER_MINIMUM_ALLOCATION_VCORES, 2);
        conf.setInt(YarnConfiguration.RM_SCHEDULER_MAXIMUM_ALLOCATION_VCORES, 1);
        try {
            scheduler.reinitialize(conf, mockContext);
            fail("Exception is expected because the min vcores allocation is"
                    + " larger than the max vcores allocation.");
        } catch (YarnRuntimeException e) {
            // Exception is expected.
            assertTrue("The thrown exception is not the expected one.",
                    e.getMessage().startsWith("Invalid resource scheduler vcores"));
        }

        conf = new YarnConfiguration();
        conf.setInt(YarnConfiguration.RM_SCHEDULER_MINIMUM_ALLOCATION_GPUS, 2);
        conf.setInt(YarnConfiguration.RM_SCHEDULER_MAXIMUM_ALLOCATION_GPUS, 1);
        try {
            scheduler.reinitialize(conf, mockContext);
            fail("Exception is expected because the min gpu allocation is"
                    + " larger than the max gpu allocation.");
        } catch (YarnRuntimeException e) {
            // Exception is expected.
            assertTrue("The thrown exception is not the expected one.",
                    e.getMessage().startsWith("Invalid resource scheduler GPUs"));
        }
    }

    private org.apache.hadoop.yarn.server.resourcemanager.NodeManager registerNode(String hostName,
            int containerManagerPort, int httpPort, String rackName, Resource capability)
            throws IOException, YarnException {
        org.apache.hadoop.yarn.server.resourcemanager.NodeManager nm = new org.apache.hadoop.yarn.server.resourcemanager.NodeManager(
                hostName, containerManagerPort, httpPort, rackName, capability, resourceManager);
        NodeAddedSchedulerEvent nodeAddEvent1 = new NodeAddedSchedulerEvent(
                resourceManager.getRMContext().getRMNodes().get(nm.getNodeId()));
        resourceManager.getResourceScheduler().handle(nodeAddEvent1);
        return nm;
    }

    @Test
    public void testCapacityScheduler() throws Exception {

        LOG.info("--- START: testCapacityScheduler ---");

        // Register node1
        String host_0 = "host_0";
        org.apache.hadoop.yarn.server.resourcemanager.NodeManager nm_0 = registerNode(host_0, 1234, 2345,
                NetworkTopology.DEFAULT_RACK, Resources.createResource(4 * GB, 1));

        // Register node2
        String host_1 = "host_1";
        org.apache.hadoop.yarn.server.resourcemanager.NodeManager nm_1 = registerNode(host_1, 1234, 2345,
                NetworkTopology.DEFAULT_RACK, Resources.createResource(2 * GB, 1));

        // ResourceRequest priorities
        Priority priority_0 = org.apache.hadoop.yarn.server.resourcemanager.resource.Priority.create(0);
        Priority priority_1 = org.apache.hadoop.yarn.server.resourcemanager.resource.Priority.create(1);

        // Submit an application
        Application application_0 = new Application("user_0", "a1", resourceManager);
        application_0.submit();

        application_0.addNodeManager(host_0, 1234, nm_0);
        application_0.addNodeManager(host_1, 1234, nm_1);

        Resource capability_0_0 = Resources.createResource(1 * GB, 1);
        application_0.addResourceRequestSpec(priority_1, capability_0_0);

        Resource capability_0_1 = Resources.createResource(2 * GB, 1);
        application_0.addResourceRequestSpec(priority_0, capability_0_1);

        Task task_0_0 = new Task(application_0, priority_1, new String[] { host_0, host_1 });
        application_0.addTask(task_0_0);

        // Submit another application
        Application application_1 = new Application("user_1", "b2", resourceManager);
        application_1.submit();

        application_1.addNodeManager(host_0, 1234, nm_0);
        application_1.addNodeManager(host_1, 1234, nm_1);

        Resource capability_1_0 = Resources.createResource(3 * GB, 1);
        application_1.addResourceRequestSpec(priority_1, capability_1_0);

        Resource capability_1_1 = Resources.createResource(2 * GB, 1);
        application_1.addResourceRequestSpec(priority_0, capability_1_1);

        Task task_1_0 = new Task(application_1, priority_1, new String[] { host_0, host_1 });
        application_1.addTask(task_1_0);

        // Send resource requests to the scheduler
        application_0.schedule();
        application_1.schedule();

        // Send a heartbeat to kick the tires on the Scheduler
        LOG.info("Kick!");

        // task_0_0 and task_1_0 allocated, used=4G
        nodeUpdate(nm_0);

        // nothing allocated
        nodeUpdate(nm_1);

        // Get allocations from the scheduler
        application_0.schedule(); // task_0_0 
        checkApplicationResourceUsage(1 * GB, application_0);

        application_1.schedule(); // task_1_0
        checkApplicationResourceUsage(3 * GB, application_1);

        checkNodeResourceUsage(4 * GB, nm_0); // task_0_0 (1G) and task_1_0 (3G)
        checkNodeResourceUsage(0 * GB, nm_1); // no tasks, 2G available

        LOG.info("Adding new tasks...");

        Task task_1_1 = new Task(application_1, priority_0, new String[] { ResourceRequest.ANY });
        application_1.addTask(task_1_1);

        application_1.schedule();

        Task task_0_1 = new Task(application_0, priority_0, new String[] { host_0, host_1 });
        application_0.addTask(task_0_1);

        application_0.schedule();

        // Send a heartbeat to kick the tires on the Scheduler
        LOG.info("Sending hb from " + nm_0.getHostName());
        // nothing new, used=4G
        nodeUpdate(nm_0);

        LOG.info("Sending hb from " + nm_1.getHostName());
        // task_0_1 is prefer as locality, used=2G
        nodeUpdate(nm_1);

        // Get allocations from the scheduler
        LOG.info("Trying to allocate...");
        application_0.schedule();
        checkApplicationResourceUsage(1 * GB, application_0);

        application_1.schedule();
        checkApplicationResourceUsage(5 * GB, application_1);

        nodeUpdate(nm_0);
        nodeUpdate(nm_1);

        checkNodeResourceUsage(4 * GB, nm_0);
        checkNodeResourceUsage(2 * GB, nm_1);

        LOG.info("--- END: testCapacityScheduler ---");
    }

    private void nodeUpdate(org.apache.hadoop.yarn.server.resourcemanager.NodeManager nm) {
        RMNode node = resourceManager.getRMContext().getRMNodes().get(nm.getNodeId());
        // Send a heartbeat to kick the tires on the Scheduler
        NodeUpdateSchedulerEvent nodeUpdate = new NodeUpdateSchedulerEvent(node);
        resourceManager.getResourceScheduler().handle(nodeUpdate);
    }

    private CapacitySchedulerConfiguration setupQueueConfiguration(CapacitySchedulerConfiguration conf) {

        // Define top-level queues
        conf.setQueues(CapacitySchedulerConfiguration.ROOT, new String[] { "a", "b" });

        conf.setCapacity(A, A_CAPACITY);
        conf.setCapacity(B, B_CAPACITY);

        // Define 2nd-level queues
        conf.setQueues(A, new String[] { "a1", "a2" });
        conf.setCapacity(A1, A1_CAPACITY);
        conf.setUserLimitFactor(A1, 100.0f);
        conf.setCapacity(A2, A2_CAPACITY);
        conf.setUserLimitFactor(A2, 100.0f);

        conf.setQueues(B, new String[] { "b1", "b2", "b3" });
        conf.setCapacity(B1, B1_CAPACITY);
        conf.setUserLimitFactor(B1, 100.0f);
        conf.setCapacity(B2, B2_CAPACITY);
        conf.setUserLimitFactor(B2, 100.0f);
        conf.setCapacity(B3, B3_CAPACITY);
        conf.setUserLimitFactor(B3, 100.0f);

        LOG.info("Setup top-level queues a and b");
        return conf;
    }

    private CapacitySchedulerConfiguration setupBlockedQueueConfiguration(CapacitySchedulerConfiguration conf) {

        // Define top-level queues
        conf.setQueues(CapacitySchedulerConfiguration.ROOT, new String[] { "a", "b" });

        conf.setCapacity(A, 80f);
        conf.setCapacity(B, 20f);
        conf.setUserLimitFactor(A, 100);
        conf.setUserLimitFactor(B, 100);
        conf.setMaximumCapacity(A, 100);
        conf.setMaximumCapacity(B, 100);
        LOG.info("Setup top-level queues a and b");
        return conf;
    }

    private CapacitySchedulerConfiguration setupOtherBlockedQueueConfiguration(
            CapacitySchedulerConfiguration conf) {

        // Define top-level queues
        conf.setQueues(CapacitySchedulerConfiguration.ROOT, new String[] { "p1", "p2" });

        conf.setCapacity(P1, 50f);
        conf.setMaximumCapacity(P1, 50f);
        conf.setCapacity(P2, 50f);
        conf.setMaximumCapacity(P2, 100f);
        // Define 2nd-level queues
        conf.setQueues(P1, new String[] { "x1", "x2" });
        conf.setCapacity(X1, 80f);
        conf.setMaximumCapacity(X1, 100f);
        conf.setUserLimitFactor(X1, 2f);
        conf.setCapacity(X2, 20f);
        conf.setMaximumCapacity(X2, 100f);
        conf.setUserLimitFactor(X2, 2f);

        conf.setQueues(P2, new String[] { "y1", "y2" });
        conf.setCapacity(Y1, 80f);
        conf.setUserLimitFactor(Y1, 2f);
        conf.setCapacity(Y2, 20f);
        conf.setUserLimitFactor(Y2, 2f);
        return conf;
    }

    @Test
    public void testMaximumCapacitySetup() {
        float delta = 0.0000001f;
        CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
        assertEquals(CapacitySchedulerConfiguration.MAXIMUM_CAPACITY_VALUE,
                conf.getNonLabeledQueueMaximumCapacity(A), delta);
        conf.setMaximumCapacity(A, 50.0f);
        assertEquals(50.0f, conf.getNonLabeledQueueMaximumCapacity(A), delta);
        conf.setMaximumCapacity(A, -1);
        assertEquals(CapacitySchedulerConfiguration.MAXIMUM_CAPACITY_VALUE,
                conf.getNonLabeledQueueMaximumCapacity(A), delta);
    }

    @Test
    public void testRefreshQueues() throws Exception {
        CapacityScheduler cs = new CapacityScheduler();
        CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
        RMContextImpl rmContext = new RMContextImpl(null, null, null, null, null, null,
                new RMContainerTokenSecretManager(conf), new NMTokenSecretManagerInRM(conf),
                new ClientToAMTokenSecretManagerInRM(), null);
        setupQueueConfiguration(conf);
        cs.setConf(new YarnConfiguration());
        cs.setRMContext(resourceManager.getRMContext());
        cs.init(conf);
        cs.start();
        cs.reinitialize(conf, rmContext);
        checkQueueCapacities(cs, A_CAPACITY, B_CAPACITY);

        conf.setCapacity(A, 80f);
        conf.setCapacity(B, 20f);
        cs.reinitialize(conf, mockContext);
        checkQueueCapacities(cs, 80f, 20f);
        cs.stop();
    }

    void checkQueueCapacities(CapacityScheduler cs, float capacityA, float capacityB) {
        CSQueue rootQueue = cs.getRootQueue();
        CSQueue queueA = findQueue(rootQueue, A);
        CSQueue queueB = findQueue(rootQueue, B);
        CSQueue queueA1 = findQueue(queueA, A1);
        CSQueue queueA2 = findQueue(queueA, A2);
        CSQueue queueB1 = findQueue(queueB, B1);
        CSQueue queueB2 = findQueue(queueB, B2);
        CSQueue queueB3 = findQueue(queueB, B3);

        float capA = capacityA / 100.0f;
        float capB = capacityB / 100.0f;

        checkQueueCapacity(queueA, capA, capA, 1.0f, 1.0f);
        checkQueueCapacity(queueB, capB, capB, 1.0f, 1.0f);
        checkQueueCapacity(queueA1, A1_CAPACITY / 100.0f, (A1_CAPACITY / 100.0f) * capA, 1.0f, 1.0f);
        checkQueueCapacity(queueA2, A2_CAPACITY / 100.0f, (A2_CAPACITY / 100.0f) * capA, 1.0f, 1.0f);
        checkQueueCapacity(queueB1, B1_CAPACITY / 100.0f, (B1_CAPACITY / 100.0f) * capB, 1.0f, 1.0f);
        checkQueueCapacity(queueB2, B2_CAPACITY / 100.0f, (B2_CAPACITY / 100.0f) * capB, 1.0f, 1.0f);
        checkQueueCapacity(queueB3, B3_CAPACITY / 100.0f, (B3_CAPACITY / 100.0f) * capB, 1.0f, 1.0f);
    }

    private void checkQueueCapacity(CSQueue q, float expectedCapacity, float expectedAbsCapacity,
            float expectedMaxCapacity, float expectedAbsMaxCapacity) {
        final float epsilon = 1e-5f;
        assertEquals("capacity", expectedCapacity, q.getCapacity(), epsilon);
        assertEquals("absolute capacity", expectedAbsCapacity, q.getAbsoluteCapacity(), epsilon);
        assertEquals("maximum capacity", expectedMaxCapacity, q.getMaximumCapacity(), epsilon);
        assertEquals("absolute maximum capacity", expectedAbsMaxCapacity, q.getAbsoluteMaximumCapacity(), epsilon);
    }

    private CSQueue findQueue(CSQueue root, String queuePath) {
        if (root.getQueuePath().equals(queuePath)) {
            return root;
        }

        List<CSQueue> childQueues = root.getChildQueues();
        if (childQueues != null) {
            for (CSQueue q : childQueues) {
                if (queuePath.startsWith(q.getQueuePath())) {
                    CSQueue result = findQueue(q, queuePath);
                    if (result != null) {
                        return result;
                    }
                }
            }
        }

        return null;
    }

    private void checkApplicationResourceUsage(int expected, Application application) {
        Assert.assertEquals(expected, application.getUsedResources().getMemorySize());
    }

    private void checkNodeResourceUsage(int expected,
            org.apache.hadoop.yarn.server.resourcemanager.NodeManager node) {
        Assert.assertEquals(expected, node.getUsed().getMemorySize());
        node.checkResourceUsage();
    }

    /** Test that parseQueue throws an exception when two leaf queues have the
     *  same name
    * @throws IOException
     */
    @Test(expected = IOException.class)
    public void testParseQueue() throws IOException {
        CapacityScheduler cs = new CapacityScheduler();
        cs.setConf(new YarnConfiguration());
        cs.setRMContext(resourceManager.getRMContext());
        CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
        setupQueueConfiguration(conf);
        cs.init(conf);
        cs.start();

        conf.setQueues(CapacitySchedulerConfiguration.ROOT + ".a.a1", new String[] { "b1" });
        conf.setCapacity(CapacitySchedulerConfiguration.ROOT + ".a.a1.b1", 100.0f);
        conf.setUserLimitFactor(CapacitySchedulerConfiguration.ROOT + ".a.a1.b1", 100.0f);

        cs.reinitialize(conf,
                new RMContextImpl(null, null, null, null, null, null, new RMContainerTokenSecretManager(conf),
                        new NMTokenSecretManagerInRM(conf), new ClientToAMTokenSecretManagerInRM(), null));
    }

    @Test
    public void testReconnectedNode() throws Exception {
        CapacitySchedulerConfiguration csConf = new CapacitySchedulerConfiguration();
        setupQueueConfiguration(csConf);
        CapacityScheduler cs = new CapacityScheduler();
        cs.setConf(new YarnConfiguration());

        cs.setRMContext(resourceManager.getRMContext());
        cs.init(csConf);
        cs.start();
        cs.reinitialize(csConf,
                new RMContextImpl(null, null, null, null, null, null, new RMContainerTokenSecretManager(csConf),
                        new NMTokenSecretManagerInRM(csConf), new ClientToAMTokenSecretManagerInRM(), null));

        RMNode n1 = MockNodes.newNodeInfo(0, MockNodes.newResource(4 * GB), 1);
        RMNode n2 = MockNodes.newNodeInfo(0, MockNodes.newResource(2 * GB), 2);

        cs.handle(new NodeAddedSchedulerEvent(n1));
        cs.handle(new NodeAddedSchedulerEvent(n2));

        Assert.assertEquals(6 * GB, cs.getClusterResource().getMemorySize());

        // reconnect n1 with downgraded memory
        n1 = MockNodes.newNodeInfo(0, MockNodes.newResource(2 * GB), 1);
        cs.handle(new NodeRemovedSchedulerEvent(n1));
        cs.handle(new NodeAddedSchedulerEvent(n1));

        Assert.assertEquals(4 * GB, cs.getClusterResource().getMemorySize());
        cs.stop();
    }

    @Test
    public void testRefreshQueuesWithNewQueue() throws Exception {
        CapacityScheduler cs = new CapacityScheduler();
        CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
        setupQueueConfiguration(conf);
        cs.setConf(new YarnConfiguration());
        cs.setRMContext(resourceManager.getRMContext());
        cs.init(conf);
        cs.start();
        cs.reinitialize(conf,
                new RMContextImpl(null, null, null, null, null, null, new RMContainerTokenSecretManager(conf),
                        new NMTokenSecretManagerInRM(conf), new ClientToAMTokenSecretManagerInRM(), null));
        checkQueueCapacities(cs, A_CAPACITY, B_CAPACITY);

        // Add a new queue b4
        String B4 = B + ".b4";
        float B4_CAPACITY = 10;

        B3_CAPACITY -= B4_CAPACITY;
        try {
            conf.setCapacity(A, 80f);
            conf.setCapacity(B, 20f);
            conf.setQueues(B, new String[] { "b1", "b2", "b3", "b4" });
            conf.setCapacity(B1, B1_CAPACITY);
            conf.setCapacity(B2, B2_CAPACITY);
            conf.setCapacity(B3, B3_CAPACITY);
            conf.setCapacity(B4, B4_CAPACITY);
            cs.reinitialize(conf, mockContext);
            checkQueueCapacities(cs, 80f, 20f);

            // Verify parent for B4
            CSQueue rootQueue = cs.getRootQueue();
            CSQueue queueB = findQueue(rootQueue, B);
            CSQueue queueB4 = findQueue(queueB, B4);

            assertEquals(queueB, queueB4.getParent());
        } finally {
            B3_CAPACITY += B4_CAPACITY;
            cs.stop();
        }
    }

    @Test
    public void testCapacitySchedulerInfo() throws Exception {
        QueueInfo queueInfo = resourceManager.getResourceScheduler().getQueueInfo("a", true, true);
        Assert.assertEquals(queueInfo.getQueueName(), "a");
        Assert.assertEquals(queueInfo.getChildQueues().size(), 2);

        List<QueueUserACLInfo> userACLInfo = resourceManager.getResourceScheduler().getQueueUserAclInfo();
        Assert.assertNotNull(userACLInfo);
        for (QueueUserACLInfo queueUserACLInfo : userACLInfo) {
            Assert.assertEquals(getQueueCount(userACLInfo, queueUserACLInfo.getQueueName()), 1);
        }

    }

    private int getQueueCount(List<QueueUserACLInfo> queueInformation, String queueName) {
        int result = 0;
        for (QueueUserACLInfo queueUserACLInfo : queueInformation) {
            if (queueName.equals(queueUserACLInfo.getQueueName())) {
                result++;
            }
        }
        return result;
    }

    @Test
    public void testBlackListNodes() throws Exception {
        Configuration conf = new Configuration();
        conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);
        MockRM rm = new MockRM(conf);
        rm.start();
        CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();

        String host = "127.0.0.1";
        RMNode node = MockNodes.newNodeInfo(0, MockNodes.newResource(4 * GB), 1, host);
        cs.handle(new NodeAddedSchedulerEvent(node));

        ApplicationId appId = BuilderUtils.newApplicationId(100, 1);
        ApplicationAttemptId appAttemptId = BuilderUtils.newApplicationAttemptId(appId, 1);

        RMAppAttemptMetrics attemptMetric = new RMAppAttemptMetrics(appAttemptId, rm.getRMContext());
        RMAppImpl app = mock(RMAppImpl.class);
        when(app.getApplicationId()).thenReturn(appId);
        RMAppAttemptImpl attempt = mock(RMAppAttemptImpl.class);
        Container container = mock(Container.class);
        when(attempt.getMasterContainer()).thenReturn(container);
        ApplicationSubmissionContext submissionContext = mock(ApplicationSubmissionContext.class);
        when(attempt.getSubmissionContext()).thenReturn(submissionContext);
        when(attempt.getAppAttemptId()).thenReturn(appAttemptId);
        when(attempt.getRMAppAttemptMetrics()).thenReturn(attemptMetric);
        when(app.getCurrentAppAttempt()).thenReturn(attempt);

        rm.getRMContext().getRMApps().put(appId, app);

        SchedulerEvent addAppEvent = new AppAddedSchedulerEvent(appId, "default", "user");
        cs.handle(addAppEvent);
        SchedulerEvent addAttemptEvent = new AppAttemptAddedSchedulerEvent(appAttemptId, false);
        cs.handle(addAttemptEvent);

        // Verify the blacklist can be updated independent of requesting containers
        cs.allocate(appAttemptId, Collections.<ResourceRequest>emptyList(), Collections.<ContainerId>emptyList(),
                Collections.singletonList(host), null, null, null);
        Assert.assertTrue(cs.getApplicationAttempt(appAttemptId).isPlaceBlacklisted(host));
        cs.allocate(appAttemptId, Collections.<ResourceRequest>emptyList(), Collections.<ContainerId>emptyList(),
                null, Collections.singletonList(host), null, null);
        Assert.assertFalse(cs.getApplicationAttempt(appAttemptId).isPlaceBlacklisted(host));
        rm.stop();
    }

    @Test
    public void testAllocateReorder() throws Exception {

        //Confirm that allocation (resource request) alone will trigger a change in
        //application ordering where appropriate

        Configuration conf = new Configuration();
        conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);
        MockRM rm = new MockRM(conf);
        rm.start();
        CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();

        LeafQueue q = (LeafQueue) cs.getQueue("default");
        Assert.assertNotNull(q);

        FairOrderingPolicy fop = new FairOrderingPolicy();
        fop.setSizeBasedWeight(true);
        q.setOrderingPolicy(fop);

        String host = "127.0.0.1";
        RMNode node = MockNodes.newNodeInfo(0, MockNodes.newResource(4 * GB), 1, host);
        cs.handle(new NodeAddedSchedulerEvent(node));

        //add app begin
        ApplicationId appId1 = BuilderUtils.newApplicationId(100, 1);
        ApplicationAttemptId appAttemptId1 = BuilderUtils.newApplicationAttemptId(appId1, 1);

        RMAppAttemptMetrics attemptMetric1 = new RMAppAttemptMetrics(appAttemptId1, rm.getRMContext());
        RMAppImpl app1 = mock(RMAppImpl.class);
        when(app1.getApplicationId()).thenReturn(appId1);
        RMAppAttemptImpl attempt1 = mock(RMAppAttemptImpl.class);
        Container container = mock(Container.class);
        when(attempt1.getMasterContainer()).thenReturn(container);
        ApplicationSubmissionContext submissionContext = mock(ApplicationSubmissionContext.class);
        when(attempt1.getSubmissionContext()).thenReturn(submissionContext);
        when(attempt1.getAppAttemptId()).thenReturn(appAttemptId1);
        when(attempt1.getRMAppAttemptMetrics()).thenReturn(attemptMetric1);
        when(app1.getCurrentAppAttempt()).thenReturn(attempt1);

        rm.getRMContext().getRMApps().put(appId1, app1);

        SchedulerEvent addAppEvent1 = new AppAddedSchedulerEvent(appId1, "default", "user");
        cs.handle(addAppEvent1);
        SchedulerEvent addAttemptEvent1 = new AppAttemptAddedSchedulerEvent(appAttemptId1, false);
        cs.handle(addAttemptEvent1);
        //add app end

        //add app begin
        ApplicationId appId2 = BuilderUtils.newApplicationId(100, 2);
        ApplicationAttemptId appAttemptId2 = BuilderUtils.newApplicationAttemptId(appId2, 1);

        RMAppAttemptMetrics attemptMetric2 = new RMAppAttemptMetrics(appAttemptId2, rm.getRMContext());
        RMAppImpl app2 = mock(RMAppImpl.class);
        when(app2.getApplicationId()).thenReturn(appId2);
        RMAppAttemptImpl attempt2 = mock(RMAppAttemptImpl.class);
        when(attempt2.getMasterContainer()).thenReturn(container);
        when(attempt2.getSubmissionContext()).thenReturn(submissionContext);
        when(attempt2.getAppAttemptId()).thenReturn(appAttemptId2);
        when(attempt2.getRMAppAttemptMetrics()).thenReturn(attemptMetric2);
        when(app2.getCurrentAppAttempt()).thenReturn(attempt2);

        rm.getRMContext().getRMApps().put(appId2, app2);

        SchedulerEvent addAppEvent2 = new AppAddedSchedulerEvent(appId2, "default", "user");
        cs.handle(addAppEvent2);
        SchedulerEvent addAttemptEvent2 = new AppAttemptAddedSchedulerEvent(appAttemptId2, false);
        cs.handle(addAttemptEvent2);
        //add app end

        RecordFactory recordFactory = RecordFactoryProvider.getRecordFactory(null);

        Priority priority = TestUtils.createMockPriority(1);
        ResourceRequest r1 = TestUtils.createResourceRequest(ResourceRequest.ANY, 1 * GB, 1, true, priority,
                recordFactory);

        //This will allocate for app1
        cs.allocate(appAttemptId1, Collections.<ResourceRequest>singletonList(r1),
                Collections.<ContainerId>emptyList(), null, null, null, null);

        //And this will result in container assignment for app1
        CapacityScheduler.schedule(cs);

        //Verify that app1 is still first in assignment order
        //This happens because app2 has no demand/a magnitude of NaN, which
        //results in app1 and app2 being equal in the fairness comparison and
        //failling back to fifo (start) ordering
        assertEquals(q.getOrderingPolicy().getAssignmentIterator().next().getId(), appId1.toString());

        //Now, allocate for app2 (this would be the first/AM allocation)
        ResourceRequest r2 = TestUtils.createResourceRequest(ResourceRequest.ANY, 1 * GB, 1, true, priority,
                recordFactory);
        cs.allocate(appAttemptId2, Collections.<ResourceRequest>singletonList(r2),
                Collections.<ContainerId>emptyList(), null, null, null, null);

        //In this case we do not perform container assignment because we want to
        //verify re-ordering based on the allocation alone

        //Now, the first app for assignment is app2
        assertEquals(q.getOrderingPolicy().getAssignmentIterator().next().getId(), appId2.toString());

        rm.stop();
    }

    @Test
    public void testResourceOverCommit() throws Exception {
        int waitCount;
        Configuration conf = new Configuration();
        conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);
        MockRM rm = new MockRM(conf);
        rm.start();

        MockNM nm1 = rm.registerNode("127.0.0.1:1234", 4 * GB);
        RMApp app1 = rm.submitApp(2048);
        // kick the scheduling, 2 GB given to AM1, remaining 2GB on nm1
        nm1.nodeHeartbeat(true);
        RMAppAttempt attempt1 = app1.getCurrentAppAttempt();
        MockAM am1 = rm.sendAMLaunched(attempt1.getAppAttemptId());
        am1.registerAppAttempt();
        SchedulerNodeReport report_nm1 = rm.getResourceScheduler().getNodeReport(nm1.getNodeId());
        // check node report, 2 GB used and 2 GB available
        Assert.assertEquals(2 * GB, report_nm1.getUsedResource().getMemorySize());
        Assert.assertEquals(2 * GB, report_nm1.getAvailableResource().getMemorySize());

        // add request for containers
        am1.addRequests(new String[] { "127.0.0.1", "127.0.0.2" }, 2 * GB, 1, 1);
        AllocateResponse alloc1Response = am1.schedule(); // send the request

        // kick the scheduler, 2 GB given to AM1, resource remaining 0
        nm1.nodeHeartbeat(true);
        while (alloc1Response.getAllocatedContainers().size() < 1) {
            LOG.info("Waiting for containers to be created for app 1...");
            Thread.sleep(100);
            alloc1Response = am1.schedule();
        }

        List<Container> allocated1 = alloc1Response.getAllocatedContainers();
        Assert.assertEquals(1, allocated1.size());
        Assert.assertEquals(2 * GB, allocated1.get(0).getResource().getMemorySize());
        Assert.assertEquals(nm1.getNodeId(), allocated1.get(0).getNodeId());

        report_nm1 = rm.getResourceScheduler().getNodeReport(nm1.getNodeId());
        // check node report, 4 GB used and 0 GB available
        Assert.assertEquals(0, report_nm1.getAvailableResource().getMemorySize());
        Assert.assertEquals(4 * GB, report_nm1.getUsedResource().getMemorySize());

        // check container is assigned with 2 GB.
        Container c1 = allocated1.get(0);
        Assert.assertEquals(2 * GB, c1.getResource().getMemorySize());

        // update node resource to 2 GB, so resource is over-consumed.
        Map<NodeId, ResourceOption> nodeResourceMap = new HashMap<NodeId, ResourceOption>();
        nodeResourceMap.put(nm1.getNodeId(), ResourceOption.newInstance(Resource.newInstance(2 * GB, 1), -1));
        UpdateNodeResourceRequest request = UpdateNodeResourceRequest.newInstance(nodeResourceMap);
        AdminService as = ((MockRM) rm).getAdminService();
        as.updateNodeResource(request);

        waitCount = 0;
        while (waitCount++ != 20) {
            report_nm1 = rm.getResourceScheduler().getNodeReport(nm1.getNodeId());
            if (report_nm1.getAvailableResource().getMemorySize() != 0) {
                break;
            }
            LOG.info("Waiting for RMNodeResourceUpdateEvent to be handled... Tried " + waitCount
                    + " times already..");
            Thread.sleep(1000);
        }
        // Now, the used resource is still 4 GB, and available resource is minus value.
        report_nm1 = rm.getResourceScheduler().getNodeReport(nm1.getNodeId());
        Assert.assertEquals(4 * GB, report_nm1.getUsedResource().getMemorySize());
        Assert.assertEquals(-2 * GB, report_nm1.getAvailableResource().getMemorySize());

        // Check container can complete successfully in case of resource over-commitment.
        ContainerStatus containerStatus = BuilderUtils.newContainerStatus(c1.getId(), ContainerState.COMPLETE, "",
                0, c1.getResource());
        nm1.containerStatus(containerStatus);
        waitCount = 0;
        while (attempt1.getJustFinishedContainers().size() < 1 && waitCount++ != 20) {
            LOG.info("Waiting for containers to be finished for app 1... Tried " + waitCount + " times already..");
            Thread.sleep(100);
        }
        Assert.assertEquals(1, attempt1.getJustFinishedContainers().size());
        Assert.assertEquals(1, am1.schedule().getCompletedContainersStatuses().size());
        report_nm1 = rm.getResourceScheduler().getNodeReport(nm1.getNodeId());
        Assert.assertEquals(2 * GB, report_nm1.getUsedResource().getMemorySize());
        // As container return 2 GB back, the available resource becomes 0 again.
        Assert.assertEquals(0 * GB, report_nm1.getAvailableResource().getMemorySize());

        // Verify no NPE is trigger in schedule after resource is updated.
        am1.addRequests(new String[] { "127.0.0.1", "127.0.0.2" }, 3 * GB, 1, 1);
        alloc1Response = am1.schedule();
        Assert.assertEquals("Shouldn't have enough resource to allocate containers", 0,
                alloc1Response.getAllocatedContainers().size());
        int times = 0;
        // try 10 times as scheduling is async process.
        while (alloc1Response.getAllocatedContainers().size() < 1 && times++ < 10) {
            LOG.info("Waiting for containers to be allocated for app 1... Tried " + times + " times already..");
            Thread.sleep(100);
        }
        Assert.assertEquals("Shouldn't have enough resource to allocate containers", 0,
                alloc1Response.getAllocatedContainers().size());
        rm.stop();
    }

    @Test
    public void testGetAppsInQueue() throws Exception {
        Application application_0 = new Application("user_0", "a1", resourceManager);
        application_0.submit();

        Application application_1 = new Application("user_0", "a2", resourceManager);
        application_1.submit();

        Application application_2 = new Application("user_0", "b2", resourceManager);
        application_2.submit();

        ResourceScheduler scheduler = resourceManager.getResourceScheduler();

        List<ApplicationAttemptId> appsInA1 = scheduler.getAppsInQueue("a1");
        assertEquals(1, appsInA1.size());

        List<ApplicationAttemptId> appsInA = scheduler.getAppsInQueue("a");
        assertTrue(appsInA.contains(application_0.getApplicationAttemptId()));
        assertTrue(appsInA.contains(application_1.getApplicationAttemptId()));
        assertEquals(2, appsInA.size());

        List<ApplicationAttemptId> appsInRoot = scheduler.getAppsInQueue("root");
        assertTrue(appsInRoot.contains(application_0.getApplicationAttemptId()));
        assertTrue(appsInRoot.contains(application_1.getApplicationAttemptId()));
        assertTrue(appsInRoot.contains(application_2.getApplicationAttemptId()));
        assertEquals(3, appsInRoot.size());

        Assert.assertNull(scheduler.getAppsInQueue("nonexistentqueue"));
    }

    @Test
    public void testAddAndRemoveAppFromCapacityScheduler() throws Exception {
        CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
        setupQueueConfiguration(conf);
        conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);
        MockRM rm = new MockRM(conf);
        @SuppressWarnings("unchecked")
        AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode> cs = (AbstractYarnScheduler<SchedulerApplicationAttempt, SchedulerNode>) rm
                .getResourceScheduler();

        SchedulerApplication<SchedulerApplicationAttempt> app = TestSchedulerUtils
                .verifyAppAddedAndRemovedFromScheduler(cs.getSchedulerApplications(), cs, "a1");
        Assert.assertEquals("a1", app.getQueue().getQueueName());
    }

    @Test
    public void testAsyncScheduling() throws Exception {
        Configuration conf = new Configuration();
        conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);
        MockRM rm = new MockRM(conf);
        rm.start();
        CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();

        final int NODES = 100;

        // Register nodes
        for (int i = 0; i < NODES; ++i) {
            String host = "192.168.1." + i;
            RMNode node = MockNodes.newNodeInfo(0, MockNodes.newResource(4 * GB), 1, host);
            cs.handle(new NodeAddedSchedulerEvent(node));
        }

        // Now directly exercise the scheduling loop
        for (int i = 0; i < NODES; ++i) {
            CapacityScheduler.schedule(cs);
        }
    }

    private void waitForAppPreemptionInfo(RMApp app, Resource preempted, int numAMPreempted, int numTaskPreempted,
            Resource currentAttemptPreempted, boolean currentAttemptAMPreempted, int numLatestAttemptTaskPreempted)
            throws InterruptedException {
        while (true) {
            RMAppMetrics appPM = app.getRMAppMetrics();
            RMAppAttemptMetrics attemptPM = app.getCurrentAppAttempt().getRMAppAttemptMetrics();

            if (appPM.getResourcePreempted().equals(preempted)
                    && appPM.getNumAMContainersPreempted() == numAMPreempted
                    && appPM.getNumNonAMContainersPreempted() == numTaskPreempted
                    && attemptPM.getResourcePreempted().equals(currentAttemptPreempted)
                    && app.getCurrentAppAttempt().getRMAppAttemptMetrics()
                            .getIsPreempted() == currentAttemptAMPreempted
                    && attemptPM.getNumNonAMContainersPreempted() == numLatestAttemptTaskPreempted) {
                return;
            }
            Thread.sleep(500);
        }
    }

    private void waitForNewAttemptCreated(RMApp app, ApplicationAttemptId previousAttemptId)
            throws InterruptedException {
        while (app.getCurrentAppAttempt().equals(previousAttemptId)) {
            Thread.sleep(500);
        }
    }

    @Test(timeout = 30000)
    public void testAllocateDoesNotBlockOnSchedulerLock() throws Exception {
        final YarnConfiguration conf = new YarnConfiguration();
        conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);
        MyContainerManager containerManager = new MyContainerManager();
        final MockRMWithAMS rm = new MockRMWithAMS(conf, containerManager);
        rm.start();

        MockNM nm1 = rm.registerNode("localhost:1234", 5120);

        Map<ApplicationAccessType, String> acls = new HashMap<ApplicationAccessType, String>(2);
        acls.put(ApplicationAccessType.VIEW_APP, "*");
        RMApp app = rm.submitApp(1024, "appname", "appuser", acls);

        nm1.nodeHeartbeat(true);

        RMAppAttempt attempt = app.getCurrentAppAttempt();
        ApplicationAttemptId applicationAttemptId = attempt.getAppAttemptId();
        int msecToWait = 10000;
        int msecToSleep = 100;
        while (attempt.getAppAttemptState() != RMAppAttemptState.LAUNCHED && msecToWait > 0) {
            LOG.info("Waiting for AppAttempt to reach LAUNCHED state. " + "Current state is "
                    + attempt.getAppAttemptState());
            Thread.sleep(msecToSleep);
            msecToWait -= msecToSleep;
        }
        Assert.assertEquals(attempt.getAppAttemptState(), RMAppAttemptState.LAUNCHED);

        // Create a client to the RM.
        final YarnRPC rpc = YarnRPC.create(conf);

        UserGroupInformation currentUser = UserGroupInformation.createRemoteUser(applicationAttemptId.toString());
        Credentials credentials = containerManager.getContainerCredentials();
        final InetSocketAddress rmBindAddress = rm.getApplicationMasterService().getBindAddress();
        Token<? extends TokenIdentifier> amRMToken = MockRMWithAMS.setupAndReturnAMRMToken(rmBindAddress,
                credentials.getAllTokens());
        currentUser.addToken(amRMToken);
        ApplicationMasterProtocol client = currentUser.doAs(new PrivilegedAction<ApplicationMasterProtocol>() {
            @Override
            public ApplicationMasterProtocol run() {
                return (ApplicationMasterProtocol) rpc.getProxy(ApplicationMasterProtocol.class, rmBindAddress,
                        conf);
            }
        });

        RegisterApplicationMasterRequest request = RegisterApplicationMasterRequest.newInstance("localhost", 12345,
                "");
        client.registerApplicationMaster(request);

        // Allocate a container
        List<ResourceRequest> asks = Collections.singletonList(
                ResourceRequest.newInstance(Priority.newInstance(1), "*", Resources.createResource(2 * GB), 1));
        AllocateRequest allocateRequest = AllocateRequest.newInstance(0, 0.0f, asks, null, null);
        client.allocate(allocateRequest);

        // Make sure the container is allocated in RM
        nm1.nodeHeartbeat(true);
        ContainerId containerId2 = ContainerId.newContainerId(applicationAttemptId, 2);
        Assert.assertTrue(rm.waitForState(nm1, containerId2, RMContainerState.ALLOCATED, 10 * 1000));

        // Acquire the container
        allocateRequest = AllocateRequest.newInstance(1, 0.0f, null, null, null);
        client.allocate(allocateRequest);

        // Launch the container
        final CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();
        RMContainer rmContainer = cs.getRMContainer(containerId2);
        rmContainer.handle(new RMContainerEvent(containerId2, RMContainerEventType.LAUNCHED));

        // grab the scheduler lock from another thread
        // and verify an allocate call in this thread doesn't block on it
        final CyclicBarrier barrier = new CyclicBarrier(2);
        Thread otherThread = new Thread(new Runnable() {
            @Override
            public void run() {
                synchronized (cs) {
                    try {
                        barrier.await();
                        barrier.await();
                    } catch (InterruptedException | BrokenBarrierException e) {
                        e.printStackTrace();
                    }
                }
            }
        });
        otherThread.start();
        barrier.await();
        List<ContainerId> release = Collections.singletonList(containerId2);
        allocateRequest = AllocateRequest.newInstance(2, 0.0f, null, release, null);
        client.allocate(allocateRequest);
        barrier.await();
        otherThread.join();

        rm.stop();
    }

    @Test
    public void testNumClusterNodes() throws Exception {
        YarnConfiguration conf = new YarnConfiguration();
        CapacityScheduler cs = new CapacityScheduler();
        cs.setConf(conf);
        RMContext rmContext = TestUtils.getMockRMContext();
        cs.setRMContext(rmContext);
        CapacitySchedulerConfiguration csConf = new CapacitySchedulerConfiguration();
        setupQueueConfiguration(csConf);
        cs.init(csConf);
        cs.start();
        assertEquals(0, cs.getNumClusterNodes());

        RMNode n1 = MockNodes.newNodeInfo(0, MockNodes.newResource(4 * GB), 1);
        RMNode n2 = MockNodes.newNodeInfo(0, MockNodes.newResource(2 * GB), 2);
        cs.handle(new NodeAddedSchedulerEvent(n1));
        cs.handle(new NodeAddedSchedulerEvent(n2));
        assertEquals(2, cs.getNumClusterNodes());

        cs.handle(new NodeRemovedSchedulerEvent(n1));
        assertEquals(1, cs.getNumClusterNodes());
        cs.handle(new NodeAddedSchedulerEvent(n1));
        assertEquals(2, cs.getNumClusterNodes());
        cs.handle(new NodeRemovedSchedulerEvent(n2));
        cs.handle(new NodeRemovedSchedulerEvent(n1));
        assertEquals(0, cs.getNumClusterNodes());

        cs.stop();
    }

    @Test(timeout = 120000)
    public void testPreemptionInfo() throws Exception {
        Configuration conf = new Configuration();
        conf.setInt(YarnConfiguration.RM_AM_MAX_ATTEMPTS, 3);
        conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);
        int CONTAINER_MEMORY = 1024; // start RM
        MockRM rm1 = new MockRM(conf);
        rm1.start();

        // get scheduler
        CapacityScheduler cs = (CapacityScheduler) rm1.getResourceScheduler();

        // start NM
        MockNM nm1 = new MockNM("127.0.0.1:1234", 15120, rm1.getResourceTrackerService());
        nm1.registerNode();

        // create app and launch the AM
        RMApp app0 = rm1.submitApp(CONTAINER_MEMORY);
        MockAM am0 = MockRM.launchAM(app0, rm1, nm1);
        am0.registerAppAttempt();

        // get scheduler app
        FiCaSchedulerApp schedulerAppAttempt = cs.getSchedulerApplications().get(app0.getApplicationId())
                .getCurrentAppAttempt();

        // allocate some containers and launch them
        List<Container> allocatedContainers = am0.allocateAndWaitForContainers(3, CONTAINER_MEMORY, nm1);

        // kill the 3 containers
        for (Container c : allocatedContainers) {
            cs.markContainerForKillable(schedulerAppAttempt.getRMContainer(c.getId()));
        }

        // check values
        waitForAppPreemptionInfo(app0, Resource.newInstance(CONTAINER_MEMORY * 3, 3), 0, 3,
                Resource.newInstance(CONTAINER_MEMORY * 3, 3), false, 3);

        // kill app0-attempt0 AM container
        cs.markContainerForKillable(
                schedulerAppAttempt.getRMContainer(app0.getCurrentAppAttempt().getMasterContainer().getId()));

        // wait for app0 failed
        waitForNewAttemptCreated(app0, am0.getApplicationAttemptId());

        // check values
        waitForAppPreemptionInfo(app0, Resource.newInstance(CONTAINER_MEMORY * 4, 4), 1, 3,
                Resource.newInstance(0, 0), false, 0);

        // launch app0-attempt1
        MockAM am1 = MockRM.launchAM(app0, rm1, nm1);
        am1.registerAppAttempt();

        schedulerAppAttempt = cs.getSchedulerApplications().get(app0.getApplicationId()).getCurrentAppAttempt();

        // allocate some containers and launch them
        allocatedContainers = am1.allocateAndWaitForContainers(3, CONTAINER_MEMORY, nm1);
        for (Container c : allocatedContainers) {
            cs.markContainerForKillable(schedulerAppAttempt.getRMContainer(c.getId()));
        }

        // check values
        waitForAppPreemptionInfo(app0, Resource.newInstance(CONTAINER_MEMORY * 7, 7), 1, 6,
                Resource.newInstance(CONTAINER_MEMORY * 3, 3), false, 3);

        rm1.stop();
    }

    @Test(timeout = 30000)
    public void testRecoverRequestAfterPreemption() throws Exception {
        Configuration conf = new Configuration();
        conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);
        MockRM rm1 = new MockRM(conf);
        rm1.start();
        MockNM nm1 = rm1.registerNode("127.0.0.1:1234", 8000);
        RMApp app1 = rm1.submitApp(1024);
        MockAM am1 = MockRM.launchAndRegisterAM(app1, rm1, nm1);
        CapacityScheduler cs = (CapacityScheduler) rm1.getResourceScheduler();

        // request a container.
        am1.allocate("127.0.0.1", 1024, 1, new ArrayList<ContainerId>());
        ContainerId containerId1 = ContainerId.newContainerId(am1.getApplicationAttemptId(), 2);
        rm1.waitForState(nm1, containerId1, RMContainerState.ALLOCATED);

        RMContainer rmContainer = cs.getRMContainer(containerId1);
        List<ResourceRequest> requests = rmContainer.getResourceRequests();
        FiCaSchedulerApp app = cs.getApplicationAttempt(am1.getApplicationAttemptId());

        FiCaSchedulerNode node = cs.getNode(rmContainer.getAllocatedNode());
        for (ResourceRequest request : requests) {
            // Skip the OffRack and RackLocal resource requests.
            if (request.getResourceName().equals(node.getRackName())
                    || request.getResourceName().equals(ResourceRequest.ANY)) {
                continue;
            }

            // Already the node local resource request is cleared from RM after
            // allocation.
            Assert.assertNull(app.getResourceRequest(request.getPriority(), request.getResourceName()));
        }

        // Call killContainer to preempt the container
        cs.markContainerForKillable(rmContainer);

        Assert.assertEquals(3, requests.size());
        for (ResourceRequest request : requests) {
            // Resource request must have added back in RM after preempt event
            // handling.
            Assert.assertEquals(1,
                    app.getResourceRequest(request.getPriority(), request.getResourceName()).getNumContainers());
        }

        // New container will be allocated and will move to ALLOCATED state
        ContainerId containerId2 = ContainerId.newContainerId(am1.getApplicationAttemptId(), 3);
        rm1.waitForState(nm1, containerId2, RMContainerState.ALLOCATED);

        // allocate container
        List<Container> containers = am1.allocate(new ArrayList<ResourceRequest>(), new ArrayList<ContainerId>())
                .getAllocatedContainers();

        // Now with updated ResourceRequest, a container is allocated for AM.
        Assert.assertTrue(containers.size() == 1);
    }

    private MockRM setUpMove() {
        CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
        return setUpMove(conf);
    }

    private MockRM setUpMove(Configuration config) {
        CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration(config);
        setupQueueConfiguration(conf);
        conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);
        MockRM rm = new MockRM(conf);
        rm.start();
        return rm;
    }

    @Test
    public void testMoveAppBasic() throws Exception {
        MockRM rm = setUpMove();
        AbstractYarnScheduler scheduler = (AbstractYarnScheduler) rm.getResourceScheduler();

        // submit an app
        RMApp app = rm.submitApp(GB, "test-move-1", "user_0", null, "a1");
        ApplicationAttemptId appAttemptId = rm.getApplicationReport(app.getApplicationId())
                .getCurrentApplicationAttemptId();

        // check preconditions
        List<ApplicationAttemptId> appsInA1 = scheduler.getAppsInQueue("a1");
        assertEquals(1, appsInA1.size());
        String queue = scheduler.getApplicationAttempt(appsInA1.get(0)).getQueue().getQueueName();
        Assert.assertTrue(queue.equals("a1"));

        List<ApplicationAttemptId> appsInA = scheduler.getAppsInQueue("a");
        assertTrue(appsInA.contains(appAttemptId));
        assertEquals(1, appsInA.size());

        List<ApplicationAttemptId> appsInRoot = scheduler.getAppsInQueue("root");
        assertTrue(appsInRoot.contains(appAttemptId));
        assertEquals(1, appsInRoot.size());

        List<ApplicationAttemptId> appsInB1 = scheduler.getAppsInQueue("b1");
        assertTrue(appsInB1.isEmpty());

        List<ApplicationAttemptId> appsInB = scheduler.getAppsInQueue("b");
        assertTrue(appsInB.isEmpty());

        // now move the app
        scheduler.moveApplication(app.getApplicationId(), "b1");

        // check postconditions
        appsInB1 = scheduler.getAppsInQueue("b1");
        assertEquals(1, appsInB1.size());
        queue = scheduler.getApplicationAttempt(appsInB1.get(0)).getQueue().getQueueName();
        Assert.assertTrue(queue.equals("b1"));

        appsInB = scheduler.getAppsInQueue("b");
        assertTrue(appsInB.contains(appAttemptId));
        assertEquals(1, appsInB.size());

        appsInRoot = scheduler.getAppsInQueue("root");
        assertTrue(appsInRoot.contains(appAttemptId));
        assertEquals(1, appsInRoot.size());

        appsInA1 = scheduler.getAppsInQueue("a1");
        assertTrue(appsInA1.isEmpty());

        appsInA = scheduler.getAppsInQueue("a");
        assertTrue(appsInA.isEmpty());

        rm.stop();
    }

    @Test
    public void testMoveAppSameParent() throws Exception {
        MockRM rm = setUpMove();
        AbstractYarnScheduler scheduler = (AbstractYarnScheduler) rm.getResourceScheduler();

        // submit an app
        RMApp app = rm.submitApp(GB, "test-move-1", "user_0", null, "a1");
        ApplicationAttemptId appAttemptId = rm.getApplicationReport(app.getApplicationId())
                .getCurrentApplicationAttemptId();

        // check preconditions
        List<ApplicationAttemptId> appsInA1 = scheduler.getAppsInQueue("a1");
        assertEquals(1, appsInA1.size());
        String queue = scheduler.getApplicationAttempt(appsInA1.get(0)).getQueue().getQueueName();
        Assert.assertTrue(queue.equals("a1"));

        List<ApplicationAttemptId> appsInA = scheduler.getAppsInQueue("a");
        assertTrue(appsInA.contains(appAttemptId));
        assertEquals(1, appsInA.size());

        List<ApplicationAttemptId> appsInRoot = scheduler.getAppsInQueue("root");
        assertTrue(appsInRoot.contains(appAttemptId));
        assertEquals(1, appsInRoot.size());

        List<ApplicationAttemptId> appsInA2 = scheduler.getAppsInQueue("a2");
        assertTrue(appsInA2.isEmpty());

        // now move the app
        scheduler.moveApplication(app.getApplicationId(), "a2");

        // check postconditions
        appsInA2 = scheduler.getAppsInQueue("a2");
        assertEquals(1, appsInA2.size());
        queue = scheduler.getApplicationAttempt(appsInA2.get(0)).getQueue().getQueueName();
        Assert.assertTrue(queue.equals("a2"));

        appsInA1 = scheduler.getAppsInQueue("a1");
        assertTrue(appsInA1.isEmpty());

        appsInA = scheduler.getAppsInQueue("a");
        assertTrue(appsInA.contains(appAttemptId));
        assertEquals(1, appsInA.size());

        appsInRoot = scheduler.getAppsInQueue("root");
        assertTrue(appsInRoot.contains(appAttemptId));
        assertEquals(1, appsInRoot.size());

        rm.stop();
    }

    @Test
    public void testMoveAppForMoveToQueueWithFreeCap() throws Exception {

        ResourceScheduler scheduler = resourceManager.getResourceScheduler();
        // Register node1
        String host_0 = "host_0";
        NodeManager nm_0 = registerNode(host_0, 1234, 2345, NetworkTopology.DEFAULT_RACK,
                Resources.createResource(4 * GB, 1));

        // Register node2
        String host_1 = "host_1";
        NodeManager nm_1 = registerNode(host_1, 1234, 2345, NetworkTopology.DEFAULT_RACK,
                Resources.createResource(2 * GB, 1));

        // ResourceRequest priorities
        Priority priority_0 = org.apache.hadoop.yarn.server.resourcemanager.resource.Priority.create(0);
        Priority priority_1 = org.apache.hadoop.yarn.server.resourcemanager.resource.Priority.create(1);

        // Submit application_0
        Application application_0 = new Application("user_0", "a1", resourceManager);
        application_0.submit(); // app + app attempt event sent to scheduler

        application_0.addNodeManager(host_0, 1234, nm_0);
        application_0.addNodeManager(host_1, 1234, nm_1);

        Resource capability_0_0 = Resources.createResource(1 * GB, 1);
        application_0.addResourceRequestSpec(priority_1, capability_0_0);

        Resource capability_0_1 = Resources.createResource(2 * GB, 1);
        application_0.addResourceRequestSpec(priority_0, capability_0_1);

        Task task_0_0 = new Task(application_0, priority_1, new String[] { host_0, host_1 });
        application_0.addTask(task_0_0);

        // Submit application_1
        Application application_1 = new Application("user_1", "b2", resourceManager);
        application_1.submit(); // app + app attempt event sent to scheduler

        application_1.addNodeManager(host_0, 1234, nm_0);
        application_1.addNodeManager(host_1, 1234, nm_1);

        Resource capability_1_0 = Resources.createResource(1 * GB, 1);
        application_1.addResourceRequestSpec(priority_1, capability_1_0);

        Resource capability_1_1 = Resources.createResource(2 * GB, 1);
        application_1.addResourceRequestSpec(priority_0, capability_1_1);

        Task task_1_0 = new Task(application_1, priority_1, new String[] { host_0, host_1 });
        application_1.addTask(task_1_0);

        // Send resource requests to the scheduler
        application_0.schedule(); // allocate
        application_1.schedule(); // allocate

        // task_0_0 task_1_0 allocated, used=2G
        nodeUpdate(nm_0);

        // nothing allocated
        nodeUpdate(nm_1);

        // Get allocations from the scheduler
        application_0.schedule(); // task_0_0
        checkApplicationResourceUsage(1 * GB, application_0);

        application_1.schedule(); // task_1_0
        checkApplicationResourceUsage(1 * GB, application_1);

        checkNodeResourceUsage(2 * GB, nm_0); // task_0_0 (1G) and task_1_0 (1G) 2G
                                              // available
        checkNodeResourceUsage(0 * GB, nm_1); // no tasks, 2G available

        // move app from a1(30% cap of total 10.5% cap) to b1(79,2% cap of 89,5%
        // total cap)
        scheduler.moveApplication(application_0.getApplicationId(), "b1");

        // 2GB 1C
        Task task_1_1 = new Task(application_1, priority_0, new String[] { ResourceRequest.ANY });
        application_1.addTask(task_1_1);

        application_1.schedule();

        // 2GB 1C
        Task task_0_1 = new Task(application_0, priority_0, new String[] { host_0, host_1 });
        application_0.addTask(task_0_1);

        application_0.schedule();

        // prev 2G used free 2G
        nodeUpdate(nm_0);

        // prev 0G used free 2G
        nodeUpdate(nm_1);

        // Get allocations from the scheduler
        application_1.schedule();
        checkApplicationResourceUsage(3 * GB, application_1);

        // Get allocations from the scheduler
        application_0.schedule();
        checkApplicationResourceUsage(3 * GB, application_0);

        checkNodeResourceUsage(4 * GB, nm_0);
        checkNodeResourceUsage(2 * GB, nm_1);

    }

    @Test
    public void testMoveAppSuccess() throws Exception {

        ResourceScheduler scheduler = resourceManager.getResourceScheduler();

        // Register node1
        String host_0 = "host_0";
        NodeManager nm_0 = registerNode(host_0, 1234, 2345, NetworkTopology.DEFAULT_RACK,
                Resources.createResource(5 * GB, 1));

        // Register node2
        String host_1 = "host_1";
        NodeManager nm_1 = registerNode(host_1, 1234, 2345, NetworkTopology.DEFAULT_RACK,
                Resources.createResource(5 * GB, 1));

        // ResourceRequest priorities
        Priority priority_0 = org.apache.hadoop.yarn.server.resourcemanager.resource.Priority.create(0);
        Priority priority_1 = org.apache.hadoop.yarn.server.resourcemanager.resource.Priority.create(1);

        // Submit application_0
        Application application_0 = new Application("user_0", "a1", resourceManager);
        application_0.submit(); // app + app attempt event sent to scheduler

        application_0.addNodeManager(host_0, 1234, nm_0);
        application_0.addNodeManager(host_1, 1234, nm_1);

        Resource capability_0_0 = Resources.createResource(3 * GB, 1);
        application_0.addResourceRequestSpec(priority_1, capability_0_0);

        Resource capability_0_1 = Resources.createResource(2 * GB, 1);
        application_0.addResourceRequestSpec(priority_0, capability_0_1);

        Task task_0_0 = new Task(application_0, priority_1, new String[] { host_0, host_1 });
        application_0.addTask(task_0_0);

        // Submit application_1
        Application application_1 = new Application("user_1", "b2", resourceManager);
        application_1.submit(); // app + app attempt event sent to scheduler

        application_1.addNodeManager(host_0, 1234, nm_0);
        application_1.addNodeManager(host_1, 1234, nm_1);

        Resource capability_1_0 = Resources.createResource(1 * GB, 1);
        application_1.addResourceRequestSpec(priority_1, capability_1_0);

        Resource capability_1_1 = Resources.createResource(2 * GB, 1);
        application_1.addResourceRequestSpec(priority_0, capability_1_1);

        Task task_1_0 = new Task(application_1, priority_1, new String[] { host_0, host_1 });
        application_1.addTask(task_1_0);

        // Send resource requests to the scheduler
        application_0.schedule(); // allocate
        application_1.schedule(); // allocate

        // b2 can only run 1 app at a time
        scheduler.moveApplication(application_0.getApplicationId(), "b2");

        nodeUpdate(nm_0);

        nodeUpdate(nm_1);

        // Get allocations from the scheduler
        application_0.schedule(); // task_0_0
        checkApplicationResourceUsage(0 * GB, application_0);

        application_1.schedule(); // task_1_0
        checkApplicationResourceUsage(1 * GB, application_1);

        // task_1_0 (1G) application_0 moved to b2 with max running app 1 so it is
        // not scheduled
        checkNodeResourceUsage(1 * GB, nm_0);
        checkNodeResourceUsage(0 * GB, nm_1);

        // lets move application_0 to a queue where it can run
        scheduler.moveApplication(application_0.getApplicationId(), "a2");
        application_0.schedule();

        nodeUpdate(nm_1);

        // Get allocations from the scheduler
        application_0.schedule(); // task_0_0
        checkApplicationResourceUsage(3 * GB, application_0);

        checkNodeResourceUsage(1 * GB, nm_0);
        checkNodeResourceUsage(3 * GB, nm_1);

    }

    @Test(expected = YarnException.class)
    public void testMoveAppViolateQueueState() throws Exception {
        resourceManager = new ResourceManager() {
            @Override
            protected RMNodeLabelsManager createNodeLabelManager() {
                RMNodeLabelsManager mgr = new NullRMNodeLabelsManager();
                mgr.init(getConfig());
                return mgr;
            }
        };
        CapacitySchedulerConfiguration csConf = new CapacitySchedulerConfiguration();
        setupQueueConfiguration(csConf);
        StringBuilder qState = new StringBuilder();
        qState.append(CapacitySchedulerConfiguration.PREFIX).append(B).append(CapacitySchedulerConfiguration.DOT)
                .append(CapacitySchedulerConfiguration.STATE);
        csConf.set(qState.toString(), QueueState.STOPPED.name());
        YarnConfiguration conf = new YarnConfiguration(csConf);
        conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);
        resourceManager.init(conf);
        resourceManager.getRMContext().getContainerTokenSecretManager().rollMasterKey();
        resourceManager.getRMContext().getNMTokenSecretManager().rollMasterKey();
        ((AsyncDispatcher) resourceManager.getRMContext().getDispatcher()).start();
        mockContext = mock(RMContext.class);
        when(mockContext.getConfigurationProvider()).thenReturn(new LocalConfigurationProvider());

        ResourceScheduler scheduler = resourceManager.getResourceScheduler();

        // Register node1
        String host_0 = "host_0";
        NodeManager nm_0 = registerNode(host_0, 1234, 2345, NetworkTopology.DEFAULT_RACK,
                Resources.createResource(6 * GB, 1));

        // ResourceRequest priorities
        Priority priority_0 = org.apache.hadoop.yarn.server.resourcemanager.resource.Priority.create(0);
        Priority priority_1 = org.apache.hadoop.yarn.server.resourcemanager.resource.Priority.create(1);

        // Submit application_0
        Application application_0 = new Application("user_0", "a1", resourceManager);
        application_0.submit(); // app + app attempt event sent to scheduler

        application_0.addNodeManager(host_0, 1234, nm_0);

        Resource capability_0_0 = Resources.createResource(3 * GB, 1);
        application_0.addResourceRequestSpec(priority_1, capability_0_0);

        Resource capability_0_1 = Resources.createResource(2 * GB, 1);
        application_0.addResourceRequestSpec(priority_0, capability_0_1);

        Task task_0_0 = new Task(application_0, priority_1, new String[] { host_0 });
        application_0.addTask(task_0_0);

        // Send resource requests to the scheduler
        application_0.schedule(); // allocate

        // task_0_0 allocated
        nodeUpdate(nm_0);

        // Get allocations from the scheduler
        application_0.schedule(); // task_0_0
        checkApplicationResourceUsage(3 * GB, application_0);

        checkNodeResourceUsage(3 * GB, nm_0);
        // b2 queue contains 3GB consumption app,
        // add another 3GB will hit max capacity limit on queue b
        scheduler.moveApplication(application_0.getApplicationId(), "b1");

    }

    @Test
    public void testMoveAppQueueMetricsCheck() throws Exception {
        ResourceScheduler scheduler = resourceManager.getResourceScheduler();

        // Register node1
        String host_0 = "host_0";
        NodeManager nm_0 = registerNode(host_0, 1234, 2345, NetworkTopology.DEFAULT_RACK,
                Resources.createResource(5 * GB, 1));

        // Register node2
        String host_1 = "host_1";
        NodeManager nm_1 = registerNode(host_1, 1234, 2345, NetworkTopology.DEFAULT_RACK,
                Resources.createResource(5 * GB, 1));

        // ResourceRequest priorities
        Priority priority_0 = org.apache.hadoop.yarn.server.resourcemanager.resource.Priority.create(0);
        Priority priority_1 = org.apache.hadoop.yarn.server.resourcemanager.resource.Priority.create(1);

        // Submit application_0
        Application application_0 = new Application("user_0", "a1", resourceManager);
        application_0.submit(); // app + app attempt event sent to scheduler

        application_0.addNodeManager(host_0, 1234, nm_0);
        application_0.addNodeManager(host_1, 1234, nm_1);

        Resource capability_0_0 = Resources.createResource(3 * GB, 1);
        application_0.addResourceRequestSpec(priority_1, capability_0_0);

        Resource capability_0_1 = Resources.createResource(2 * GB, 1);
        application_0.addResourceRequestSpec(priority_0, capability_0_1);

        Task task_0_0 = new Task(application_0, priority_1, new String[] { host_0, host_1 });
        application_0.addTask(task_0_0);

        // Submit application_1
        Application application_1 = new Application("user_1", "b2", resourceManager);
        application_1.submit(); // app + app attempt event sent to scheduler

        application_1.addNodeManager(host_0, 1234, nm_0);
        application_1.addNodeManager(host_1, 1234, nm_1);

        Resource capability_1_0 = Resources.createResource(1 * GB, 1);
        application_1.addResourceRequestSpec(priority_1, capability_1_0);

        Resource capability_1_1 = Resources.createResource(2 * GB, 1);
        application_1.addResourceRequestSpec(priority_0, capability_1_1);

        Task task_1_0 = new Task(application_1, priority_1, new String[] { host_0, host_1 });
        application_1.addTask(task_1_0);

        // Send resource requests to the scheduler
        application_0.schedule(); // allocate
        application_1.schedule(); // allocate

        nodeUpdate(nm_0);

        nodeUpdate(nm_1);

        CapacityScheduler cs = (CapacityScheduler) resourceManager.getResourceScheduler();
        CSQueue origRootQ = cs.getRootQueue();
        CapacitySchedulerInfo oldInfo = new CapacitySchedulerInfo(origRootQ, cs);
        int origNumAppsA = getNumAppsInQueue("a", origRootQ.getChildQueues());
        int origNumAppsRoot = origRootQ.getNumApplications();

        scheduler.moveApplication(application_0.getApplicationId(), "a2");

        CSQueue newRootQ = cs.getRootQueue();
        int newNumAppsA = getNumAppsInQueue("a", newRootQ.getChildQueues());
        int newNumAppsRoot = newRootQ.getNumApplications();
        CapacitySchedulerInfo newInfo = new CapacitySchedulerInfo(newRootQ, cs);
        CapacitySchedulerLeafQueueInfo origOldA1 = (CapacitySchedulerLeafQueueInfo) getQueueInfo("a1",
                oldInfo.getQueues());
        CapacitySchedulerLeafQueueInfo origNewA1 = (CapacitySchedulerLeafQueueInfo) getQueueInfo("a1",
                newInfo.getQueues());
        CapacitySchedulerLeafQueueInfo targetOldA2 = (CapacitySchedulerLeafQueueInfo) getQueueInfo("a2",
                oldInfo.getQueues());
        CapacitySchedulerLeafQueueInfo targetNewA2 = (CapacitySchedulerLeafQueueInfo) getQueueInfo("a2",
                newInfo.getQueues());
        // originally submitted here
        assertEquals(1, origOldA1.getNumApplications());
        assertEquals(1, origNumAppsA);
        assertEquals(2, origNumAppsRoot);
        // after the move
        assertEquals(0, origNewA1.getNumApplications());
        assertEquals(1, newNumAppsA);
        assertEquals(2, newNumAppsRoot);
        // original consumption on a1
        assertEquals(3 * GB, origOldA1.getResourcesUsed().getMemorySize());
        assertEquals(1, origOldA1.getResourcesUsed().getvCores());
        assertEquals(0, origNewA1.getResourcesUsed().getMemorySize()); // after the move
        assertEquals(0, origNewA1.getResourcesUsed().getvCores()); // after the move
        // app moved here with live containers
        assertEquals(3 * GB, targetNewA2.getResourcesUsed().getMemorySize());
        assertEquals(1, targetNewA2.getResourcesUsed().getvCores());
        // it was empty before the move
        assertEquals(0, targetOldA2.getNumApplications());
        assertEquals(0, targetOldA2.getResourcesUsed().getMemorySize());
        assertEquals(0, targetOldA2.getResourcesUsed().getvCores());
        // after the app moved here
        assertEquals(1, targetNewA2.getNumApplications());
        // 1 container on original queue before move
        assertEquals(1, origOldA1.getNumContainers());
        // after the move the resource released
        assertEquals(0, origNewA1.getNumContainers());
        // and moved to the new queue
        assertEquals(1, targetNewA2.getNumContainers());
        // which originally didn't have any
        assertEquals(0, targetOldA2.getNumContainers());
        // 1 user with 3GB
        assertEquals(3 * GB, origOldA1.getUsers().getUsersList().get(0).getResourcesUsed().getMemorySize());
        // 1 user with 1 core
        assertEquals(1, origOldA1.getUsers().getUsersList().get(0).getResourcesUsed().getvCores());
        // user ha no more running app in the orig queue
        assertEquals(0, origNewA1.getUsers().getUsersList().size());
        // 1 user with 3GB
        assertEquals(3 * GB, targetNewA2.getUsers().getUsersList().get(0).getResourcesUsed().getMemorySize());
        // 1 user with 1 core
        assertEquals(1, targetNewA2.getUsers().getUsersList().get(0).getResourcesUsed().getvCores());

        // Get allocations from the scheduler
        application_0.schedule(); // task_0_0
        checkApplicationResourceUsage(3 * GB, application_0);

        application_1.schedule(); // task_1_0
        checkApplicationResourceUsage(1 * GB, application_1);

        // task_1_0 (1G) application_0 moved to b2 with max running app 1 so it is
        // not scheduled
        checkNodeResourceUsage(4 * GB, nm_0);
        checkNodeResourceUsage(0 * GB, nm_1);

    }

    private int getNumAppsInQueue(String name, List<CSQueue> queues) {
        for (CSQueue queue : queues) {
            if (queue.getQueueName().equals(name)) {
                return queue.getNumApplications();
            }
        }
        return -1;
    }

    private CapacitySchedulerQueueInfo getQueueInfo(String name, CapacitySchedulerQueueInfoList info) {
        if (info != null) {
            for (CapacitySchedulerQueueInfo queueInfo : info.getQueueInfoList()) {
                if (queueInfo.getQueueName().equals(name)) {
                    return queueInfo;
                } else {
                    CapacitySchedulerQueueInfo result = getQueueInfo(name, queueInfo.getQueues());
                    if (result == null) {
                        continue;
                    }
                    return result;
                }
            }
        }
        return null;
    }

    @Test
    public void testMoveAllApps() throws Exception {
        MockRM rm = setUpMove();
        AbstractYarnScheduler scheduler = (AbstractYarnScheduler) rm.getResourceScheduler();

        // submit an app
        RMApp app = rm.submitApp(GB, "test-move-1", "user_0", null, "a1");
        ApplicationAttemptId appAttemptId = rm.getApplicationReport(app.getApplicationId())
                .getCurrentApplicationAttemptId();

        // check preconditions
        List<ApplicationAttemptId> appsInA1 = scheduler.getAppsInQueue("a1");
        assertEquals(1, appsInA1.size());

        List<ApplicationAttemptId> appsInA = scheduler.getAppsInQueue("a");
        assertTrue(appsInA.contains(appAttemptId));
        assertEquals(1, appsInA.size());
        String queue = scheduler.getApplicationAttempt(appsInA1.get(0)).getQueue().getQueueName();
        Assert.assertTrue(queue.equals("a1"));

        List<ApplicationAttemptId> appsInRoot = scheduler.getAppsInQueue("root");
        assertTrue(appsInRoot.contains(appAttemptId));
        assertEquals(1, appsInRoot.size());

        List<ApplicationAttemptId> appsInB1 = scheduler.getAppsInQueue("b1");
        assertTrue(appsInB1.isEmpty());

        List<ApplicationAttemptId> appsInB = scheduler.getAppsInQueue("b");
        assertTrue(appsInB.isEmpty());

        // now move the app
        scheduler.moveAllApps("a1", "b1");

        // check postconditions
        Thread.sleep(1000);
        appsInB1 = scheduler.getAppsInQueue("b1");
        assertEquals(1, appsInB1.size());
        queue = scheduler.getApplicationAttempt(appsInB1.get(0)).getQueue().getQueueName();
        Assert.assertTrue(queue.equals("b1"));

        appsInB = scheduler.getAppsInQueue("b");
        assertTrue(appsInB.contains(appAttemptId));
        assertEquals(1, appsInB.size());

        appsInRoot = scheduler.getAppsInQueue("root");
        assertTrue(appsInRoot.contains(appAttemptId));
        assertEquals(1, appsInRoot.size());

        appsInA1 = scheduler.getAppsInQueue("a1");
        assertTrue(appsInA1.isEmpty());

        appsInA = scheduler.getAppsInQueue("a");
        assertTrue(appsInA.isEmpty());

        rm.stop();
    }

    @Test
    public void testMoveAllAppsInvalidDestination() throws Exception {
        MockRM rm = setUpMove();
        YarnScheduler scheduler = rm.getResourceScheduler();

        // submit an app
        RMApp app = rm.submitApp(GB, "test-move-1", "user_0", null, "a1");
        ApplicationAttemptId appAttemptId = rm.getApplicationReport(app.getApplicationId())
                .getCurrentApplicationAttemptId();

        // check preconditions
        List<ApplicationAttemptId> appsInA1 = scheduler.getAppsInQueue("a1");
        assertEquals(1, appsInA1.size());

        List<ApplicationAttemptId> appsInA = scheduler.getAppsInQueue("a");
        assertTrue(appsInA.contains(appAttemptId));
        assertEquals(1, appsInA.size());

        List<ApplicationAttemptId> appsInRoot = scheduler.getAppsInQueue("root");
        assertTrue(appsInRoot.contains(appAttemptId));
        assertEquals(1, appsInRoot.size());

        List<ApplicationAttemptId> appsInB1 = scheduler.getAppsInQueue("b1");
        assertTrue(appsInB1.isEmpty());

        List<ApplicationAttemptId> appsInB = scheduler.getAppsInQueue("b");
        assertTrue(appsInB.isEmpty());

        // now move the app
        try {
            scheduler.moveAllApps("a1", "DOES_NOT_EXIST");
            Assert.fail();
        } catch (YarnException e) {
            // expected
        }

        // check postconditions, app should still be in a1
        appsInA1 = scheduler.getAppsInQueue("a1");
        assertEquals(1, appsInA1.size());

        appsInA = scheduler.getAppsInQueue("a");
        assertTrue(appsInA.contains(appAttemptId));
        assertEquals(1, appsInA.size());

        appsInRoot = scheduler.getAppsInQueue("root");
        assertTrue(appsInRoot.contains(appAttemptId));
        assertEquals(1, appsInRoot.size());

        appsInB1 = scheduler.getAppsInQueue("b1");
        assertTrue(appsInB1.isEmpty());

        appsInB = scheduler.getAppsInQueue("b");
        assertTrue(appsInB.isEmpty());

        rm.stop();
    }

    @Test
    public void testMoveAllAppsInvalidSource() throws Exception {
        MockRM rm = setUpMove();
        YarnScheduler scheduler = rm.getResourceScheduler();

        // submit an app
        RMApp app = rm.submitApp(GB, "test-move-1", "user_0", null, "a1");
        ApplicationAttemptId appAttemptId = rm.getApplicationReport(app.getApplicationId())
                .getCurrentApplicationAttemptId();

        // check preconditions
        List<ApplicationAttemptId> appsInA1 = scheduler.getAppsInQueue("a1");
        assertEquals(1, appsInA1.size());

        List<ApplicationAttemptId> appsInA = scheduler.getAppsInQueue("a");
        assertTrue(appsInA.contains(appAttemptId));
        assertEquals(1, appsInA.size());

        List<ApplicationAttemptId> appsInRoot = scheduler.getAppsInQueue("root");
        assertTrue(appsInRoot.contains(appAttemptId));
        assertEquals(1, appsInRoot.size());

        List<ApplicationAttemptId> appsInB1 = scheduler.getAppsInQueue("b1");
        assertTrue(appsInB1.isEmpty());

        List<ApplicationAttemptId> appsInB = scheduler.getAppsInQueue("b");
        assertTrue(appsInB.isEmpty());

        // now move the app
        try {
            scheduler.moveAllApps("DOES_NOT_EXIST", "b1");
            Assert.fail();
        } catch (YarnException e) {
            // expected
        }

        // check postconditions, app should still be in a1
        appsInA1 = scheduler.getAppsInQueue("a1");
        assertEquals(1, appsInA1.size());

        appsInA = scheduler.getAppsInQueue("a");
        assertTrue(appsInA.contains(appAttemptId));
        assertEquals(1, appsInA.size());

        appsInRoot = scheduler.getAppsInQueue("root");
        assertTrue(appsInRoot.contains(appAttemptId));
        assertEquals(1, appsInRoot.size());

        appsInB1 = scheduler.getAppsInQueue("b1");
        assertTrue(appsInB1.isEmpty());

        appsInB = scheduler.getAppsInQueue("b");
        assertTrue(appsInB.isEmpty());

        rm.stop();
    }

    @Test
    public void testKillAllAppsInQueue() throws Exception {
        MockRM rm = setUpMove();
        AbstractYarnScheduler scheduler = (AbstractYarnScheduler) rm.getResourceScheduler();

        // submit an app
        RMApp app = rm.submitApp(GB, "test-move-1", "user_0", null, "a1");
        ApplicationAttemptId appAttemptId = rm.getApplicationReport(app.getApplicationId())
                .getCurrentApplicationAttemptId();

        // check preconditions
        List<ApplicationAttemptId> appsInA1 = scheduler.getAppsInQueue("a1");
        assertEquals(1, appsInA1.size());

        List<ApplicationAttemptId> appsInA = scheduler.getAppsInQueue("a");
        assertTrue(appsInA.contains(appAttemptId));
        assertEquals(1, appsInA.size());
        String queue = scheduler.getApplicationAttempt(appsInA1.get(0)).getQueue().getQueueName();
        Assert.assertTrue(queue.equals("a1"));

        List<ApplicationAttemptId> appsInRoot = scheduler.getAppsInQueue("root");
        assertTrue(appsInRoot.contains(appAttemptId));
        assertEquals(1, appsInRoot.size());

        // now kill the app
        scheduler.killAllAppsInQueue("a1");

        // check postconditions
        rm.waitForState(app.getApplicationId(), RMAppState.KILLED);
        appsInRoot = scheduler.getAppsInQueue("root");
        assertTrue(appsInRoot.isEmpty());

        appsInA1 = scheduler.getAppsInQueue("a1");
        assertTrue(appsInA1.isEmpty());

        appsInA = scheduler.getAppsInQueue("a");
        assertTrue(appsInA.isEmpty());

        rm.stop();
    }

    @Test
    public void testKillAllAppsInvalidSource() throws Exception {
        MockRM rm = setUpMove();
        YarnScheduler scheduler = rm.getResourceScheduler();

        // submit an app
        RMApp app = rm.submitApp(GB, "test-move-1", "user_0", null, "a1");
        ApplicationAttemptId appAttemptId = rm.getApplicationReport(app.getApplicationId())
                .getCurrentApplicationAttemptId();

        // check preconditions
        List<ApplicationAttemptId> appsInA1 = scheduler.getAppsInQueue("a1");
        assertEquals(1, appsInA1.size());

        List<ApplicationAttemptId> appsInA = scheduler.getAppsInQueue("a");
        assertTrue(appsInA.contains(appAttemptId));
        assertEquals(1, appsInA.size());

        List<ApplicationAttemptId> appsInRoot = scheduler.getAppsInQueue("root");
        assertTrue(appsInRoot.contains(appAttemptId));
        assertEquals(1, appsInRoot.size());

        // now kill the app
        try {
            scheduler.killAllAppsInQueue("DOES_NOT_EXIST");
            Assert.fail();
        } catch (YarnException e) {
            // expected
        }

        // check postconditions, app should still be in a1
        appsInA1 = scheduler.getAppsInQueue("a1");
        assertEquals(1, appsInA1.size());

        appsInA = scheduler.getAppsInQueue("a");
        assertTrue(appsInA.contains(appAttemptId));
        assertEquals(1, appsInA.size());

        appsInRoot = scheduler.getAppsInQueue("root");
        assertTrue(appsInRoot.contains(appAttemptId));
        assertEquals(1, appsInRoot.size());

        rm.stop();
    }

    // Test to ensure that we don't carry out reservation on nodes
    // that have no CPU available when using the DominantResourceCalculator
    @Test(timeout = 30000)
    public void testAppReservationWithDominantResourceCalculator() throws Exception {
        CapacitySchedulerConfiguration csconf = new CapacitySchedulerConfiguration();
        csconf.setResourceComparator(DominantResourceCalculator.class);

        YarnConfiguration conf = new YarnConfiguration(csconf);
        conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);

        MockRM rm = new MockRM(conf);
        rm.start();

        MockNM nm1 = rm.registerNode("127.0.0.1:1234", 10 * GB, 1, 1);

        // register extra nodes to bump up cluster resource
        MockNM nm2 = rm.registerNode("127.0.0.1:1235", 10 * GB, 4, 4);
        rm.registerNode("127.0.0.1:1236", 10 * GB, 4, 4);

        RMApp app1 = rm.submitApp(1024);
        // kick the scheduling
        nm1.nodeHeartbeat(true);
        RMAppAttempt attempt1 = app1.getCurrentAppAttempt();
        MockAM am1 = rm.sendAMLaunched(attempt1.getAppAttemptId());
        am1.registerAppAttempt();
        SchedulerNodeReport report_nm1 = rm.getResourceScheduler().getNodeReport(nm1.getNodeId());

        // check node report
        Assert.assertEquals(1 * GB, report_nm1.getUsedResource().getMemorySize());
        Assert.assertEquals(9 * GB, report_nm1.getAvailableResource().getMemorySize());

        // add request for containers
        am1.addRequests(new String[] { "127.0.0.1", "127.0.0.2" }, 1 * GB, 1, 1);
        am1.schedule(); // send the request

        // kick the scheduler, container reservation should not happen
        nm1.nodeHeartbeat(true);
        Thread.sleep(1000);
        AllocateResponse allocResponse = am1.schedule();
        ApplicationResourceUsageReport report = rm.getResourceScheduler()
                .getAppResourceUsageReport(attempt1.getAppAttemptId());
        Assert.assertEquals(0, allocResponse.getAllocatedContainers().size());
        Assert.assertEquals(0, report.getNumReservedContainers());

        // container should get allocated on this node
        nm2.nodeHeartbeat(true);

        while (allocResponse.getAllocatedContainers().size() == 0) {
            Thread.sleep(100);
            allocResponse = am1.schedule();
        }
        report = rm.getResourceScheduler().getAppResourceUsageReport(attempt1.getAppAttemptId());
        Assert.assertEquals(1, allocResponse.getAllocatedContainers().size());
        Assert.assertEquals(0, report.getNumReservedContainers());
        rm.stop();
    }

    // Test to ensure that we don't carry out reservation on nodes
    // that have no CPU available when using the DominantResourceCalculator
    @Test(timeout = 30000)
    public void testAppReservationWithDominantResourceCalculatorGPU() throws Exception {
        CapacitySchedulerConfiguration csconf = new CapacitySchedulerConfiguration();
        csconf.setResourceComparator(DominantResourceCalculatorGPU.class);

        YarnConfiguration conf = new YarnConfiguration(csconf);
        conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);

        MockRM rm = new MockRM(conf);
        rm.start();

        MockNM nm1 = rm.registerNode("127.0.0.1:1234", 10 * GB, 4, 0);

        // register extra nodes to bump up cluster resource
        MockNM nm2 = rm.registerNode("127.0.0.1:1235", 10 * GB, 4, 4);

        RMApp app1 = rm.submitApp(1024);
        // kick the scheduling
        nm1.nodeHeartbeat(true);
        RMAppAttempt attempt1 = app1.getCurrentAppAttempt();
        MockAM am1 = rm.sendAMLaunched(attempt1.getAppAttemptId());
        am1.registerAppAttempt();
        SchedulerNodeReport report_nm1 = rm.getResourceScheduler().getNodeReport(nm1.getNodeId());

        // check node report
        Assert.assertEquals(1 * GB, report_nm1.getUsedResource().getMemorySize());
        Assert.assertEquals(9 * GB, report_nm1.getAvailableResource().getMemorySize());

        // add request for containers
        am1.addRequests(new String[] { "127.0.0.1", "127.0.0.2" }, 1 * GB, 1, 5);
        am1.getRequests().get(0).getCapability().setGPUs(1);
        am1.schedule(); // send the request

        // kick the scheduler, container reservation should not happen
        nm1.nodeHeartbeat(true);
        Thread.sleep(1000);
        AllocateResponse allocResponse = am1.schedule();
        ApplicationResourceUsageReport report = rm.getResourceScheduler()
                .getAppResourceUsageReport(attempt1.getAppAttemptId());
        Assert.assertEquals(0, allocResponse.getAllocatedContainers().size());
        Assert.assertEquals(0, report.getNumReservedContainers());

        // container should get allocated on this node
        nm1.nodeHeartbeat(true);
        nm2.nodeHeartbeat(true);

        while (allocResponse.getAllocatedContainers().size() == 0) {
            Thread.sleep(100);
            allocResponse = am1.schedule();
        }
        report = rm.getResourceScheduler().getAppResourceUsageReport(attempt1.getAppAttemptId());
        Assert.assertEquals(4, allocResponse.getAllocatedContainers().size());
        Assert.assertEquals(nm2.getNodeId(), allocResponse.getAllocatedContainers().get(0).getNodeId());
        Assert.assertEquals(nm2.getNodeId(), allocResponse.getAllocatedContainers().get(1).getNodeId());
        Assert.assertEquals(nm2.getNodeId(), allocResponse.getAllocatedContainers().get(2).getNodeId());
        Assert.assertEquals(nm2.getNodeId(), allocResponse.getAllocatedContainers().get(3).getNodeId());
        Assert.assertEquals(1, report.getNumReservedContainers());
        Assert.assertEquals(1, report.getReservedResources().getGPUs());

        report_nm1 = rm.getResourceScheduler().getNodeReport(nm1.getNodeId());
        SchedulerNodeReport report_nm2 = rm.getResourceScheduler().getNodeReport(nm2.getNodeId());

        Assert.assertEquals(1 * GB, report_nm1.getUsedResource().getMemorySize());
        Assert.assertEquals(0, report_nm1.getUsedResource().getGPUs());

        Assert.assertEquals(4, report_nm2.getUsedResource().getGPUs());
        Assert.assertEquals(4 * GB, report_nm2.getUsedResource().getMemorySize());

        rm.stop();
    }

    @Test
    public void testPreemptionDisabled() throws Exception {
        CapacityScheduler cs = new CapacityScheduler();
        CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
        conf.setBoolean(YarnConfiguration.RM_SCHEDULER_ENABLE_MONITORS, true);
        RMContextImpl rmContext = new RMContextImpl(null, null, null, null, null, null,
                new RMContainerTokenSecretManager(conf), new NMTokenSecretManagerInRM(conf),
                new ClientToAMTokenSecretManagerInRM(), null);
        setupQueueConfiguration(conf);
        cs.setConf(new YarnConfiguration());
        cs.setRMContext(resourceManager.getRMContext());
        cs.init(conf);
        cs.start();
        cs.reinitialize(conf, rmContext);

        CSQueue rootQueue = cs.getRootQueue();
        CSQueue queueB = findQueue(rootQueue, B);
        CSQueue queueB2 = findQueue(queueB, B2);

        // When preemption turned on for the whole system
        // (yarn.resourcemanager.scheduler.monitor.enable=true), and with no other 
        // preemption properties set, queue root.b.b2 should be preemptable.
        assertFalse("queue " + B2 + " should default to preemptable", queueB2.getPreemptionDisabled());

        // Disable preemption at the root queue level.
        // The preemption property should be inherited from root all the
        // way down so that root.b.b2 should NOT be preemptable.
        conf.setPreemptionDisabled(rootQueue.getQueuePath(), true);
        cs.reinitialize(conf, rmContext);
        assertTrue("queue " + B2 + " should have inherited non-preemptability from root",
                queueB2.getPreemptionDisabled());

        // Enable preemption for root (grandparent) but disable for root.b (parent).
        // root.b.b2 should inherit property from parent and NOT be preemptable
        conf.setPreemptionDisabled(rootQueue.getQueuePath(), false);
        conf.setPreemptionDisabled(queueB.getQueuePath(), true);
        cs.reinitialize(conf, rmContext);
        assertTrue("queue " + B2 + " should have inherited non-preemptability from parent",
                queueB2.getPreemptionDisabled());

        // When preemption is turned on for root.b.b2, it should be preemptable
        // even though preemption is disabled on root.b (parent).
        conf.setPreemptionDisabled(queueB2.getQueuePath(), false);
        cs.reinitialize(conf, rmContext);
        assertFalse("queue " + B2 + " should have been preemptable", queueB2.getPreemptionDisabled());
    }

    @Test
    public void testRefreshQueuesMaxAllocationRefresh() throws Exception {
        // queue refresh should not allow changing the maximum allocation setting
        // per queue to be smaller than previous setting
        CapacityScheduler cs = new CapacityScheduler();
        CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
        setupQueueConfiguration(conf);
        cs.setConf(new YarnConfiguration());
        cs.setRMContext(resourceManager.getRMContext());
        cs.init(conf);
        cs.start();
        cs.reinitialize(conf, mockContext);
        checkQueueCapacities(cs, A_CAPACITY, B_CAPACITY);

        assertEquals("max allocation in CS", YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_MB,
                cs.getMaximumResourceCapability().getMemorySize());
        assertEquals("max gpu allocation in CS", YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_GPUS,
                cs.getMaximumResourceCapability().getGPUs());
        assertEquals("max allocation for A1", YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_MB,
                conf.getMaximumAllocationPerQueue(A1).getMemorySize());
        assertEquals("max allocation", YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_MB,
                conf.getMaximumAllocation().getMemorySize());

        CSQueue rootQueue = cs.getRootQueue();
        CSQueue queueA = findQueue(rootQueue, A);
        CSQueue queueA1 = findQueue(queueA, A1);
        assertEquals("queue max allocation", ((LeafQueue) queueA1).getMaximumAllocation().getMemorySize(), 8192);

        setMaxAllocMb(conf, A1, 4096);

        try {
            cs.reinitialize(conf, mockContext);
            fail("should have thrown exception");
        } catch (IOException e) {
            assertTrue("max allocation exception", e.getCause().toString().contains("not be decreased"));
        }

        setMaxAllocMb(conf, A1, 8192);
        cs.reinitialize(conf, mockContext);

        setMaxAllocVcores(conf, A1, YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_VCORES - 1);
        try {
            cs.reinitialize(conf, mockContext);
            fail("should have thrown exception");
        } catch (IOException e) {
            assertTrue("max allocation exception", e.getCause().toString().contains("not be decreased"));
        }
    }

    @Test
    public void testRefreshQueuesMaxAllocationPerQueueLarge() throws Exception {
        // verify we can't set the allocation per queue larger then cluster setting
        CapacityScheduler cs = new CapacityScheduler();
        cs.setConf(new YarnConfiguration());
        cs.setRMContext(resourceManager.getRMContext());
        CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
        setupQueueConfiguration(conf);
        cs.init(conf);
        cs.start();
        // change max allocation for B3 queue to be larger then cluster max
        setMaxAllocMb(conf, B3, YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_MB + 2048);
        try {
            cs.reinitialize(conf, mockContext);
            fail("should have thrown exception");
        } catch (IOException e) {
            assertTrue("maximum allocation exception", e.getCause().getMessage().contains("maximum allocation"));
        }

        setMaxAllocMb(conf, B3, YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_MB);
        cs.reinitialize(conf, mockContext);

        setMaxAllocVcores(conf, B3, YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_VCORES + 1);
        try {
            cs.reinitialize(conf, mockContext);
            fail("should have thrown exception");
        } catch (IOException e) {
            assertTrue("maximum allocation exception", e.getCause().getMessage().contains("maximum allocation"));
        }
    }

    @Test
    public void testRefreshQueuesMaxAllocationRefreshLarger() throws Exception {
        // queue refresh should allow max allocation per queue to go larger
        CapacityScheduler cs = new CapacityScheduler();
        cs.setConf(new YarnConfiguration());
        cs.setRMContext(resourceManager.getRMContext());
        CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
        setupQueueConfiguration(conf);
        setMaxAllocMb(conf, YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_MB);
        setMaxAllocVcores(conf, YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_VCORES);
        setMaxAllocMb(conf, A1, 4096);
        setMaxAllocVcores(conf, A1, 2);
        cs.init(conf);
        cs.start();
        cs.reinitialize(conf, mockContext);
        checkQueueCapacities(cs, A_CAPACITY, B_CAPACITY);

        assertEquals("max capability MB in CS", YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_MB,
                cs.getMaximumResourceCapability().getMemorySize());
        assertEquals("max capability vcores in CS",
                YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_VCORES,
                cs.getMaximumResourceCapability().getVirtualCores());
        assertEquals("max allocation MB A1", 4096, conf.getMaximumAllocationPerQueue(A1).getMemorySize());
        assertEquals("max allocation vcores A1", 2, conf.getMaximumAllocationPerQueue(A1).getVirtualCores());
        assertEquals("cluster max allocation MB", YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_MB,
                conf.getMaximumAllocation().getMemorySize());
        assertEquals("cluster max allocation vcores",
                YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_VCORES,
                conf.getMaximumAllocation().getVirtualCores());

        CSQueue rootQueue = cs.getRootQueue();
        CSQueue queueA = findQueue(rootQueue, A);
        CSQueue queueA1 = findQueue(queueA, A1);
        assertEquals("queue max allocation", ((LeafQueue) queueA1).getMaximumAllocation().getMemorySize(), 4096);

        setMaxAllocMb(conf, A1, 6144);
        setMaxAllocVcores(conf, A1, 3);
        cs.reinitialize(conf, null);
        // conf will have changed but we shouldn't be able to change max allocation
        // for the actual queue
        assertEquals("max allocation MB A1", 6144, conf.getMaximumAllocationPerQueue(A1).getMemorySize());
        assertEquals("max allocation vcores A1", 3, conf.getMaximumAllocationPerQueue(A1).getVirtualCores());
        assertEquals("max allocation MB cluster", YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_MB,
                conf.getMaximumAllocation().getMemorySize());
        assertEquals("max allocation vcores cluster",
                YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_VCORES,
                conf.getMaximumAllocation().getVirtualCores());
        assertEquals("queue max allocation MB", 6144, ((LeafQueue) queueA1).getMaximumAllocation().getMemorySize());
        assertEquals("queue max allocation vcores", 3,
                ((LeafQueue) queueA1).getMaximumAllocation().getVirtualCores());
        assertEquals("max capability MB cluster", YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_MB,
                cs.getMaximumResourceCapability().getMemorySize());
        assertEquals("cluster max capability vcores",
                YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_VCORES,
                cs.getMaximumResourceCapability().getVirtualCores());
    }

    @Test
    public void testRefreshQueuesMaxAllocationCSError() throws Exception {
        // Try to refresh the cluster level max allocation size to be smaller
        // and it should error out
        CapacityScheduler cs = new CapacityScheduler();
        cs.setConf(new YarnConfiguration());
        cs.setRMContext(resourceManager.getRMContext());
        CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
        setupQueueConfiguration(conf);
        setMaxAllocMb(conf, 10240);
        setMaxAllocVcores(conf, 10);
        setMaxAllocMb(conf, A1, 4096);
        setMaxAllocVcores(conf, A1, 4);
        cs.init(conf);
        cs.start();
        cs.reinitialize(conf, mockContext);
        checkQueueCapacities(cs, A_CAPACITY, B_CAPACITY);

        assertEquals("max allocation MB in CS", 10240, cs.getMaximumResourceCapability().getMemorySize());
        assertEquals("max allocation vcores in CS", 10, cs.getMaximumResourceCapability().getVirtualCores());

        setMaxAllocMb(conf, 6144);
        try {
            cs.reinitialize(conf, mockContext);
            fail("should have thrown exception");
        } catch (IOException e) {
            assertTrue("max allocation exception", e.getCause().toString().contains("not be decreased"));
        }

        setMaxAllocMb(conf, 10240);
        cs.reinitialize(conf, mockContext);

        setMaxAllocVcores(conf, 8);
        try {
            cs.reinitialize(conf, mockContext);
            fail("should have thrown exception");
        } catch (IOException e) {
            assertTrue("max allocation exception", e.getCause().toString().contains("not be decreased"));
        }
    }

    @Test
    public void testRefreshQueuesMaxAllocationCSLarger() throws Exception {
        // Try to refresh the cluster level max allocation size to be larger
        // and verify that if there is no setting per queue it uses the
        // cluster level setting.
        CapacityScheduler cs = new CapacityScheduler();
        cs.setConf(new YarnConfiguration());
        cs.setRMContext(resourceManager.getRMContext());
        CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
        setupQueueConfiguration(conf);
        setMaxAllocMb(conf, 10240);
        setMaxAllocVcores(conf, 10);
        setMaxAllocMb(conf, A1, 4096);
        setMaxAllocVcores(conf, A1, 4);
        cs.init(conf);
        cs.start();
        cs.reinitialize(conf, mockContext);
        checkQueueCapacities(cs, A_CAPACITY, B_CAPACITY);

        assertEquals("max allocation MB in CS", 10240, cs.getMaximumResourceCapability().getMemorySize());
        assertEquals("max allocation vcores in CS", 10, cs.getMaximumResourceCapability().getVirtualCores());

        CSQueue rootQueue = cs.getRootQueue();
        CSQueue queueA = findQueue(rootQueue, A);
        CSQueue queueB = findQueue(rootQueue, B);
        CSQueue queueA1 = findQueue(queueA, A1);
        CSQueue queueA2 = findQueue(queueA, A2);
        CSQueue queueB2 = findQueue(queueB, B2);

        assertEquals("queue A1 max allocation MB", 4096,
                ((LeafQueue) queueA1).getMaximumAllocation().getMemorySize());
        assertEquals("queue A1 max allocation vcores", 4,
                ((LeafQueue) queueA1).getMaximumAllocation().getVirtualCores());
        assertEquals("queue A2 max allocation MB", 10240,
                ((LeafQueue) queueA2).getMaximumAllocation().getMemorySize());
        assertEquals("queue A2 max allocation vcores", 10,
                ((LeafQueue) queueA2).getMaximumAllocation().getVirtualCores());
        assertEquals("queue B2 max allocation MB", 10240,
                ((LeafQueue) queueB2).getMaximumAllocation().getMemorySize());
        assertEquals("queue B2 max allocation vcores", 10,
                ((LeafQueue) queueB2).getMaximumAllocation().getVirtualCores());

        setMaxAllocMb(conf, 12288);
        setMaxAllocVcores(conf, 12);
        cs.reinitialize(conf, null);
        // cluster level setting should change and any queues without
        // per queue setting
        assertEquals("max allocation MB in CS", 12288, cs.getMaximumResourceCapability().getMemorySize());
        assertEquals("max allocation vcores in CS", 12, cs.getMaximumResourceCapability().getVirtualCores());
        assertEquals("queue A1 max MB allocation", 4096,
                ((LeafQueue) queueA1).getMaximumAllocation().getMemorySize());
        assertEquals("queue A1 max vcores allocation", 4,
                ((LeafQueue) queueA1).getMaximumAllocation().getVirtualCores());
        assertEquals("queue A2 max MB allocation", 12288,
                ((LeafQueue) queueA2).getMaximumAllocation().getMemorySize());
        assertEquals("queue A2 max vcores allocation", 12,
                ((LeafQueue) queueA2).getMaximumAllocation().getVirtualCores());
        assertEquals("queue B2 max MB allocation", 12288,
                ((LeafQueue) queueB2).getMaximumAllocation().getMemorySize());
        assertEquals("queue B2 max vcores allocation", 12,
                ((LeafQueue) queueB2).getMaximumAllocation().getVirtualCores());
    }

    private void waitContainerAllocated(MockAM am, int mem, int nContainer, int startContainerId, MockRM rm,
            MockNM nm) throws Exception {
        for (int cId = startContainerId; cId < startContainerId + nContainer; cId++) {
            am.allocate("*", mem, 1, new ArrayList<ContainerId>());
            ContainerId containerId = ContainerId.newContainerId(am.getApplicationAttemptId(), cId);
            Assert.assertTrue(rm.waitForState(nm, containerId, RMContainerState.ALLOCATED, 10 * 1000));
        }
    }

    @Test
    public void testHierarchyQueuesCurrentLimits() throws Exception {
        /*
         * Queue tree:
         *          Root
         *        /     \
         *       A       B
         *      / \    / | \
         *     A1 A2  B1 B2 B3
         */
        YarnConfiguration conf = new YarnConfiguration(
                setupQueueConfiguration(new CapacitySchedulerConfiguration()));
        conf.setBoolean(CapacitySchedulerConfiguration.ENABLE_USER_METRICS, true);

        MemoryRMStateStore memStore = new MemoryRMStateStore();
        memStore.init(conf);
        MockRM rm1 = new MockRM(conf, memStore);
        rm1.start();
        MockNM nm1 = new MockNM("127.0.0.1:1234", 100 * GB, rm1.getResourceTrackerService());
        nm1.registerNode();

        RMApp app1 = rm1.submitApp(1 * GB, "app", "user", null, "b1");
        MockAM am1 = MockRM.launchAndRegisterAM(app1, rm1, nm1);

        waitContainerAllocated(am1, 1 * GB, 1, 2, rm1, nm1);

        // Maximum resoure of b1 is 100 * 0.895 * 0.792 = 71 GB
        // 2 GBs used by am, so it's 71 - 2 = 69G.
        Assert.assertEquals(69 * GB, am1.doHeartbeat().getAvailableResources().getMemorySize());

        RMApp app2 = rm1.submitApp(1 * GB, "app", "user", null, "b2");
        MockAM am2 = MockRM.launchAndRegisterAM(app2, rm1, nm1);

        // Allocate 5 containers, each one is 8 GB in am2 (40 GB in total)
        waitContainerAllocated(am2, 8 * GB, 5, 2, rm1, nm1);

        // Allocated one more container with 1 GB resource in b1
        waitContainerAllocated(am1, 1 * GB, 1, 3, rm1, nm1);

        // Total is 100 GB, 
        // B2 uses 41 GB (5 * 8GB containers and 1 AM container)
        // B1 uses 3 GB (2 * 1GB containers and 1 AM container)
        // Available is 100 - 41 - 3 = 56 GB
        Assert.assertEquals(56 * GB, am1.doHeartbeat().getAvailableResources().getMemorySize());

        // Now we submit app3 to a1 (in higher level hierarchy), to see if headroom
        // of app1 (in queue b1) updated correctly
        RMApp app3 = rm1.submitApp(1 * GB, "app", "user", null, "a1");
        MockAM am3 = MockRM.launchAndRegisterAM(app3, rm1, nm1);

        // Allocate 3 containers, each one is 8 GB in am3 (24 GB in total)
        waitContainerAllocated(am3, 8 * GB, 3, 2, rm1, nm1);

        // Allocated one more container with 4 GB resource in b1
        waitContainerAllocated(am1, 1 * GB, 1, 4, rm1, nm1);

        // Total is 100 GB, 
        // B2 uses 41 GB (5 * 8GB containers and 1 AM container)
        // B1 uses 4 GB (3 * 1GB containers and 1 AM container)
        // A1 uses 25 GB (3 * 8GB containers and 1 AM container)
        // Available is 100 - 41 - 4 - 25 = 30 GB
        Assert.assertEquals(30 * GB, am1.doHeartbeat().getAvailableResources().getMemorySize());
    }

    @Test
    public void testParentQueueMaxCapsAreRespected() throws Exception {
        /*
         * Queue tree:
         *          Root
         *        /     \
         *       A       B
         *      / \
         *     A1 A2 
         */
        CapacitySchedulerConfiguration csConf = new CapacitySchedulerConfiguration();
        csConf.setQueues(CapacitySchedulerConfiguration.ROOT, new String[] { "a", "b" });
        csConf.setCapacity(A, 50);
        csConf.setMaximumCapacity(A, 50);
        csConf.setCapacity(B, 50);

        // Define 2nd-level queues
        csConf.setQueues(A, new String[] { "a1", "a2" });
        csConf.setCapacity(A1, 50);
        csConf.setUserLimitFactor(A1, 100.0f);
        csConf.setCapacity(A2, 50);
        csConf.setUserLimitFactor(A2, 100.0f);
        csConf.setCapacity(B1, B1_CAPACITY);
        csConf.setUserLimitFactor(B1, 100.0f);

        YarnConfiguration conf = new YarnConfiguration(csConf);
        conf.setBoolean(CapacitySchedulerConfiguration.ENABLE_USER_METRICS, true);

        MemoryRMStateStore memStore = new MemoryRMStateStore();
        memStore.init(conf);
        MockRM rm1 = new MockRM(conf, memStore);
        rm1.start();
        MockNM nm1 = new MockNM("127.0.0.1:1234", 24 * GB, rm1.getResourceTrackerService());
        nm1.registerNode();

        // Launch app1 in a1, resource usage is 1GB (am) + 4GB * 2 = 9GB 
        RMApp app1 = rm1.submitApp(1 * GB, "app", "user", null, "a1");
        MockAM am1 = MockRM.launchAndRegisterAM(app1, rm1, nm1);
        waitContainerAllocated(am1, 4 * GB, 2, 2, rm1, nm1);

        // Try to launch app2 in a2, asked 2GB, should success 
        RMApp app2 = rm1.submitApp(2 * GB, "app", "user", null, "a2");
        MockAM am2 = MockRM.launchAndRegisterAM(app2, rm1, nm1);
        try {
            // Try to allocate a container, a's usage=11G/max=12
            // a1's usage=9G/max=12
            // a2's usage=2G/max=12
            // In this case, if a2 asked 2G, should fail.
            waitContainerAllocated(am2, 2 * GB, 1, 2, rm1, nm1);
        } catch (AssertionError failure) {
            // Expected, return;
            return;
        }
        Assert.fail("Shouldn't successfully allocate containers for am2, "
                + "queue-a's max capacity will be violated if container allocated");
    }

    @SuppressWarnings("unchecked")
    private <E> Set<E> toSet(E... elements) {
        Set<E> set = Sets.newHashSet(elements);
        return set;
    }

    @Test
    public void testQueueHierarchyPendingResourceUpdate() throws Exception {
        Configuration conf = TestUtils.getConfigurationWithQueueLabels(new Configuration(false));
        conf.setBoolean(YarnConfiguration.NODE_LABELS_ENABLED, true);

        final RMNodeLabelsManager mgr = new NullRMNodeLabelsManager();
        mgr.init(conf);
        mgr.addToCluserNodeLabelsWithDefaultExclusivity(ImmutableSet.of("x", "y"));
        mgr.addLabelsToNode(ImmutableMap.of(NodeId.newInstance("h1", 0), toSet("x")));

        MemoryRMStateStore memStore = new MemoryRMStateStore();
        memStore.init(conf);
        MockRM rm = new MockRM(conf, memStore) {
            protected RMNodeLabelsManager createNodeLabelManager() {
                return mgr;
            }
        };

        rm.start();
        MockNM nm1 = // label = x
                new MockNM("h1:1234", 200 * GB, rm.getResourceTrackerService());
        nm1.registerNode();

        MockNM nm2 = // label = ""
                new MockNM("h2:1234", 200 * GB, rm.getResourceTrackerService());
        nm2.registerNode();

        // Launch app1 in queue=a1
        RMApp app1 = rm.submitApp(1 * GB, "app", "user", null, "a1");
        MockAM am1 = MockRM.launchAndRegisterAM(app1, rm, nm2);

        // Launch app2 in queue=b1  
        RMApp app2 = rm.submitApp(8 * GB, "app", "user", null, "b1");
        MockAM am2 = MockRM.launchAndRegisterAM(app2, rm, nm2);

        // am1 asks for 8 * 1GB container for no label
        am1.allocate(Arrays.asList(
                ResourceRequest.newInstance(Priority.newInstance(1), "*", Resources.createResource(1 * GB), 8)),
                null);

        checkPendingResource(rm, "a1", 8 * GB, null);
        checkPendingResource(rm, "a", 8 * GB, null);
        checkPendingResource(rm, "root", 8 * GB, null);

        // am2 asks for 8 * 1GB container for no label
        am2.allocate(Arrays.asList(
                ResourceRequest.newInstance(Priority.newInstance(1), "*", Resources.createResource(1 * GB), 8)),
                null);

        checkPendingResource(rm, "a1", 8 * GB, null);
        checkPendingResource(rm, "a", 8 * GB, null);
        checkPendingResource(rm, "b1", 8 * GB, null);
        checkPendingResource(rm, "b", 8 * GB, null);
        // root = a + b
        checkPendingResource(rm, "root", 16 * GB, null);

        // am2 asks for 8 * 1GB container in another priority for no label
        am2.allocate(Arrays.asList(
                ResourceRequest.newInstance(Priority.newInstance(2), "*", Resources.createResource(1 * GB), 8)),
                null);

        checkPendingResource(rm, "a1", 8 * GB, null);
        checkPendingResource(rm, "a", 8 * GB, null);
        checkPendingResource(rm, "b1", 16 * GB, null);
        checkPendingResource(rm, "b", 16 * GB, null);
        // root = a + b
        checkPendingResource(rm, "root", 24 * GB, null);

        // am1 asks 4 GB resource instead of 8 * GB for priority=1
        am1.allocate(Arrays.asList(
                ResourceRequest.newInstance(Priority.newInstance(1), "*", Resources.createResource(4 * GB), 1)),
                null);

        checkPendingResource(rm, "a1", 4 * GB, null);
        checkPendingResource(rm, "a", 4 * GB, null);
        checkPendingResource(rm, "b1", 16 * GB, null);
        checkPendingResource(rm, "b", 16 * GB, null);
        // root = a + b
        checkPendingResource(rm, "root", 20 * GB, null);

        // am1 asks 8 * GB resource which label=x
        am1.allocate(Arrays.asList(ResourceRequest.newInstance(Priority.newInstance(2), "*",
                Resources.createResource(8 * GB), 1, true, "x")), null);

        checkPendingResource(rm, "a1", 4 * GB, null);
        checkPendingResource(rm, "a", 4 * GB, null);
        checkPendingResource(rm, "a1", 8 * GB, "x");
        checkPendingResource(rm, "a", 8 * GB, "x");
        checkPendingResource(rm, "b1", 16 * GB, null);
        checkPendingResource(rm, "b", 16 * GB, null);
        // root = a + b
        checkPendingResource(rm, "root", 20 * GB, null);
        checkPendingResource(rm, "root", 8 * GB, "x");

        // some containers allocated for am1, pending resource should decrease
        ContainerId containerId = ContainerId.newContainerId(am1.getApplicationAttemptId(), 2);
        Assert.assertTrue(rm.waitForState(nm1, containerId, RMContainerState.ALLOCATED, 10 * 1000));
        containerId = ContainerId.newContainerId(am1.getApplicationAttemptId(), 3);
        Assert.assertTrue(rm.waitForState(nm2, containerId, RMContainerState.ALLOCATED, 10 * 1000));

        checkPendingResource(rm, "a1", 0 * GB, null);
        checkPendingResource(rm, "a", 0 * GB, null);
        checkPendingResource(rm, "a1", 0 * GB, "x");
        checkPendingResource(rm, "a", 0 * GB, "x");
        // some containers could be allocated for am2 when we allocating containers
        // for am1, just check if pending resource of b1/b/root > 0 
        checkPendingResourceGreaterThanZero(rm, "b1", null);
        checkPendingResourceGreaterThanZero(rm, "b", null);
        // root = a + b
        checkPendingResourceGreaterThanZero(rm, "root", null);
        checkPendingResource(rm, "root", 0 * GB, "x");

        // complete am2, pending resource should be 0 now
        AppAttemptRemovedSchedulerEvent appRemovedEvent = new AppAttemptRemovedSchedulerEvent(
                am2.getApplicationAttemptId(), RMAppAttemptState.FINISHED, false);
        rm.getResourceScheduler().handle(appRemovedEvent);

        checkPendingResource(rm, "a1", 0 * GB, null);
        checkPendingResource(rm, "a", 0 * GB, null);
        checkPendingResource(rm, "a1", 0 * GB, "x");
        checkPendingResource(rm, "a", 0 * GB, "x");
        checkPendingResource(rm, "b1", 0 * GB, null);
        checkPendingResource(rm, "b", 0 * GB, null);
        checkPendingResource(rm, "root", 0 * GB, null);
        checkPendingResource(rm, "root", 0 * GB, "x");
    }

    private void checkPendingResource(MockRM rm, String queueName, int memory, String label) {
        CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();
        CSQueue queue = cs.getQueue(queueName);
        Assert.assertEquals(memory, queue.getQueueResourceUsage()
                .getPending(label == null ? RMNodeLabelsManager.NO_LABEL : label).getMemorySize());
    }

    private void checkPendingResourceGreaterThanZero(MockRM rm, String queueName, String label) {
        CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();
        CSQueue queue = cs.getQueue(queueName);
        Assert.assertTrue(queue.getQueueResourceUsage()
                .getPending(label == null ? RMNodeLabelsManager.NO_LABEL : label).getMemorySize() > 0);
    }

    // Test verifies AM Used resource for LeafQueue when AM ResourceRequest is
    // lesser than minimumAllocation
    @Test(timeout = 30000)
    public void testAMUsedResource() throws Exception {
        MockRM rm = setUpMove();
        rm.registerNode("127.0.0.1:1234", 4 * GB);

        Configuration conf = rm.getConfig();
        int minAllocMb = conf.getInt(YarnConfiguration.RM_SCHEDULER_MINIMUM_ALLOCATION_MB,
                YarnConfiguration.DEFAULT_RM_SCHEDULER_MINIMUM_ALLOCATION_MB);
        int amMemory = 50;
        assertTrue("AM memory is greater than or equal to minAllocation", amMemory < minAllocMb);
        Resource minAllocResource = Resource.newInstance(minAllocMb, 1);
        String queueName = "a1";
        RMApp rmApp = rm.submitApp(amMemory, "app-1", "user_0", null, queueName);

        assertEquals("RMApp does not containes minimum allocation", minAllocResource,
                rmApp.getAMResourceRequest().getCapability());

        ResourceScheduler scheduler = rm.getRMContext().getScheduler();
        LeafQueue queueA = (LeafQueue) ((CapacityScheduler) scheduler).getQueue(queueName);
        assertEquals("Minimum Resource for AM is incorrect", minAllocResource,
                queueA.getUser("user_0").getResourceUsage().getAMUsed());
        rm.stop();
    }

    // Verifies headroom passed to ApplicationMaster has been updated in
    // RMAppAttemptMetrics
    @Test
    public void testApplicationHeadRoom() throws Exception {
        Configuration conf = new Configuration();
        conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);
        MockRM rm = new MockRM(conf);
        rm.start();
        CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();

        ApplicationId appId = BuilderUtils.newApplicationId(100, 1);
        ApplicationAttemptId appAttemptId = BuilderUtils.newApplicationAttemptId(appId, 1);

        RMAppAttemptMetrics attemptMetric = new RMAppAttemptMetrics(appAttemptId, rm.getRMContext());
        RMAppImpl app = mock(RMAppImpl.class);
        when(app.getApplicationId()).thenReturn(appId);
        RMAppAttemptImpl attempt = mock(RMAppAttemptImpl.class);
        Container container = mock(Container.class);
        when(attempt.getMasterContainer()).thenReturn(container);
        ApplicationSubmissionContext submissionContext = mock(ApplicationSubmissionContext.class);
        when(attempt.getSubmissionContext()).thenReturn(submissionContext);
        when(attempt.getAppAttemptId()).thenReturn(appAttemptId);
        when(attempt.getRMAppAttemptMetrics()).thenReturn(attemptMetric);
        when(app.getCurrentAppAttempt()).thenReturn(attempt);

        rm.getRMContext().getRMApps().put(appId, app);

        SchedulerEvent addAppEvent = new AppAddedSchedulerEvent(appId, "default", "user");
        cs.handle(addAppEvent);
        SchedulerEvent addAttemptEvent = new AppAttemptAddedSchedulerEvent(appAttemptId, false);
        cs.handle(addAttemptEvent);

        Allocation allocate = cs.allocate(appAttemptId, Collections.<ResourceRequest>emptyList(),
                Collections.<ContainerId>emptyList(), null, null, null, null);

        Assert.assertNotNull(attempt);

        Assert.assertEquals(Resource.newInstance(0, 0), allocate.getResourceLimit());
        Assert.assertEquals(Resource.newInstance(0, 0), attemptMetric.getApplicationAttemptHeadroom());

        // Add a node to cluster
        Resource newResource = Resource.newInstance(4 * GB, 1);
        RMNode node = MockNodes.newNodeInfo(0, newResource, 1, "127.0.0.1");
        cs.handle(new NodeAddedSchedulerEvent(node));

        allocate = cs.allocate(appAttemptId, Collections.<ResourceRequest>emptyList(),
                Collections.<ContainerId>emptyList(), null, null, null, null);

        // All resources should be sent as headroom
        Assert.assertEquals(newResource, allocate.getResourceLimit());
        Assert.assertEquals(newResource, attemptMetric.getApplicationAttemptHeadroom());

        rm.stop();
    }

    @Test
    public void testHeadRoomCalculationWithDRC() throws Exception {
        // test with total cluster resource of 20GB memory and 20 vcores.
        // the queue where two apps running has user limit 0.8
        // allocate 10GB memory and 1 vcore to app 1.
        // app 1 should have headroom
        // 20GB*0.8 - 10GB = 6GB memory available and 15 vcores.
        // allocate 1GB memory and 1 vcore to app2.
        // app 2 should have headroom 20GB - 10 - 1 = 1GB memory,
        // and 20*0.8 - 1 = 15 vcores.

        CapacitySchedulerConfiguration csconf = new CapacitySchedulerConfiguration();
        csconf.setResourceComparator(DominantResourceCalculator.class);

        YarnConfiguration conf = new YarnConfiguration(csconf);
        conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);

        MockRM rm = new MockRM(conf);
        rm.start();

        CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();
        LeafQueue qb = (LeafQueue) cs.getQueue("default");
        qb.setUserLimitFactor((float) 0.8);

        // add app 1
        ApplicationId appId = BuilderUtils.newApplicationId(100, 1);
        ApplicationAttemptId appAttemptId = BuilderUtils.newApplicationAttemptId(appId, 1);

        RMAppAttemptMetrics attemptMetric = new RMAppAttemptMetrics(appAttemptId, rm.getRMContext());
        RMAppImpl app = mock(RMAppImpl.class);
        when(app.getApplicationId()).thenReturn(appId);
        RMAppAttemptImpl attempt = mock(RMAppAttemptImpl.class);
        Container container = mock(Container.class);
        when(attempt.getMasterContainer()).thenReturn(container);
        ApplicationSubmissionContext submissionContext = mock(ApplicationSubmissionContext.class);
        when(attempt.getSubmissionContext()).thenReturn(submissionContext);
        when(attempt.getAppAttemptId()).thenReturn(appAttemptId);
        when(attempt.getRMAppAttemptMetrics()).thenReturn(attemptMetric);
        when(app.getCurrentAppAttempt()).thenReturn(attempt);

        rm.getRMContext().getRMApps().put(appId, app);

        SchedulerEvent addAppEvent = new AppAddedSchedulerEvent(appId, "default", "user1");
        cs.handle(addAppEvent);
        SchedulerEvent addAttemptEvent = new AppAttemptAddedSchedulerEvent(appAttemptId, false);
        cs.handle(addAttemptEvent);

        // add app 2
        ApplicationId appId2 = BuilderUtils.newApplicationId(100, 2);
        ApplicationAttemptId appAttemptId2 = BuilderUtils.newApplicationAttemptId(appId2, 1);

        RMAppAttemptMetrics attemptMetric2 = new RMAppAttemptMetrics(appAttemptId2, rm.getRMContext());
        RMAppImpl app2 = mock(RMAppImpl.class);
        when(app2.getApplicationId()).thenReturn(appId2);
        RMAppAttemptImpl attempt2 = mock(RMAppAttemptImpl.class);
        when(attempt2.getMasterContainer()).thenReturn(container);
        when(attempt2.getSubmissionContext()).thenReturn(submissionContext);
        when(attempt2.getAppAttemptId()).thenReturn(appAttemptId2);
        when(attempt2.getRMAppAttemptMetrics()).thenReturn(attemptMetric2);
        when(app2.getCurrentAppAttempt()).thenReturn(attempt2);

        rm.getRMContext().getRMApps().put(appId2, app2);
        addAppEvent = new AppAddedSchedulerEvent(appId2, "default", "user2");
        cs.handle(addAppEvent);
        addAttemptEvent = new AppAttemptAddedSchedulerEvent(appAttemptId2, false);
        cs.handle(addAttemptEvent);

        // add nodes  to cluster, so cluster have 20GB and 20 vcores
        Resource newResource = Resource.newInstance(10 * GB, 10);
        RMNode node = MockNodes.newNodeInfo(0, newResource, 1, "127.0.0.1");
        cs.handle(new NodeAddedSchedulerEvent(node));

        Resource newResource2 = Resource.newInstance(10 * GB, 10);
        RMNode node2 = MockNodes.newNodeInfo(0, newResource2, 1, "127.0.0.2");
        cs.handle(new NodeAddedSchedulerEvent(node2));

        FiCaSchedulerApp fiCaApp1 = cs.getSchedulerApplications().get(app.getApplicationId())
                .getCurrentAppAttempt();

        FiCaSchedulerApp fiCaApp2 = cs.getSchedulerApplications().get(app2.getApplicationId())
                .getCurrentAppAttempt();
        Priority u0Priority = TestUtils.createMockPriority(1);
        RecordFactory recordFactory = RecordFactoryProvider.getRecordFactory(null);

        // allocate container for app1 with 10GB memory and 1 vcore
        fiCaApp1.updateResourceRequests(Collections.singletonList(
                TestUtils.createResourceRequest(ResourceRequest.ANY, 10 * GB, 1, true, u0Priority, recordFactory)));
        cs.handle(new NodeUpdateSchedulerEvent(node));
        cs.handle(new NodeUpdateSchedulerEvent(node2));
        assertEquals(6 * GB, fiCaApp1.getHeadroom().getMemorySize());
        assertEquals(15, fiCaApp1.getHeadroom().getVirtualCores());

        // allocate container for app2 with 1GB memory and 1 vcore
        fiCaApp2.updateResourceRequests(Collections.singletonList(
                TestUtils.createResourceRequest(ResourceRequest.ANY, 1 * GB, 1, true, u0Priority, recordFactory)));
        cs.handle(new NodeUpdateSchedulerEvent(node));
        cs.handle(new NodeUpdateSchedulerEvent(node2));
        assertEquals(9 * GB, fiCaApp2.getHeadroom().getMemorySize());
        assertEquals(15, fiCaApp2.getHeadroom().getVirtualCores());
    }

    @Test
    public void testDefaultNodeLabelExpressionQueueConfig() throws Exception {
        CapacityScheduler cs = new CapacityScheduler();
        CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
        setupQueueConfiguration(conf);
        conf.setDefaultNodeLabelExpression("root.a", " x");
        conf.setDefaultNodeLabelExpression("root.b", " y ");
        cs.setConf(new YarnConfiguration());
        cs.setRMContext(resourceManager.getRMContext());
        cs.init(conf);
        cs.start();

        QueueInfo queueInfoA = cs.getQueueInfo("a", true, false);
        Assert.assertEquals(queueInfoA.getQueueName(), "a");
        Assert.assertEquals(queueInfoA.getDefaultNodeLabelExpression(), "x");

        QueueInfo queueInfoB = cs.getQueueInfo("b", true, false);
        Assert.assertEquals(queueInfoB.getQueueName(), "b");
        Assert.assertEquals(queueInfoB.getDefaultNodeLabelExpression(), "y");
    }

    @Test(timeout = 60000)
    public void testAMLimitUsage() throws Exception {

        CapacitySchedulerConfiguration config = new CapacitySchedulerConfiguration();

        config.set(CapacitySchedulerConfiguration.RESOURCE_CALCULATOR_CLASS,
                DefaultResourceCalculator.class.getName());
        verifyAMLimitForLeafQueue(config);

        config.set(CapacitySchedulerConfiguration.RESOURCE_CALCULATOR_CLASS,
                DominantResourceCalculator.class.getName());
        verifyAMLimitForLeafQueue(config);
    }

    private FiCaSchedulerApp getFiCaSchedulerApp(MockRM rm, ApplicationId appId) {
        CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();
        return cs.getSchedulerApplications().get(appId).getCurrentAppAttempt();
    }

    @Test
    public void testPendingResourceUpdatedAccordingToIncreaseRequestChanges() throws Exception {
        Configuration conf = TestUtils.getConfigurationWithQueueLabels(new Configuration(false));
        conf.setBoolean(YarnConfiguration.NODE_LABELS_ENABLED, true);

        final RMNodeLabelsManager mgr = new NullRMNodeLabelsManager();
        mgr.init(conf);

        MemoryRMStateStore memStore = new MemoryRMStateStore();
        memStore.init(conf);
        MockRM rm = new MockRM(conf, memStore) {
            protected RMNodeLabelsManager createNodeLabelManager() {
                return mgr;
            }
        };

        rm.start();

        MockNM nm1 = // label = ""
                new MockNM("h1:1234", 200 * GB, rm.getResourceTrackerService());
        nm1.registerNode();

        // Launch app1 in queue=a1
        RMApp app1 = rm.submitApp(1 * GB, "app", "user", null, "a1");
        MockAM am1 = MockRM.launchAndRegisterAM(app1, rm, nm1);

        // Allocate two more containers
        am1.allocate(Arrays.asList(
                ResourceRequest.newInstance(Priority.newInstance(1), "*", Resources.createResource(2 * GB), 2)),
                null);
        ContainerId containerId1 = ContainerId.newContainerId(am1.getApplicationAttemptId(), 1);
        ContainerId containerId2 = ContainerId.newContainerId(am1.getApplicationAttemptId(), 2);
        ContainerId containerId3 = ContainerId.newContainerId(am1.getApplicationAttemptId(), 3);
        Assert.assertTrue(rm.waitForState(nm1, containerId3, RMContainerState.ALLOCATED, 10 * 1000));
        // Acquire them
        am1.allocate(null, null);
        sentRMContainerLaunched(rm, ContainerId.newContainerId(am1.getApplicationAttemptId(), 1L));
        sentRMContainerLaunched(rm, ContainerId.newContainerId(am1.getApplicationAttemptId(), 2L));
        sentRMContainerLaunched(rm, ContainerId.newContainerId(am1.getApplicationAttemptId(), 3L));

        // am1 asks to change its AM container from 1GB to 3GB
        am1.sendContainerResizingRequest(Arrays.asList(UpdateContainerRequest.newInstance(0, containerId1,
                ContainerUpdateType.INCREASE_RESOURCE, Resources.createResource(3 * GB))));

        FiCaSchedulerApp app = getFiCaSchedulerApp(rm, app1.getApplicationId());

        Assert.assertEquals(2 * GB, app.getAppAttemptResourceUsage().getPending().getMemorySize());
        checkPendingResource(rm, "a1", 2 * GB, null);
        checkPendingResource(rm, "a", 2 * GB, null);
        checkPendingResource(rm, "root", 2 * GB, null);

        // am1 asks to change containerId2 (2G -> 3G) and containerId3 (2G -> 5G)
        am1.sendContainerResizingRequest(Arrays.asList(
                UpdateContainerRequest.newInstance(0, containerId2, ContainerUpdateType.INCREASE_RESOURCE,
                        Resources.createResource(3 * GB)),
                UpdateContainerRequest.newInstance(0, containerId3, ContainerUpdateType.INCREASE_RESOURCE,
                        Resources.createResource(5 * GB))));

        Assert.assertEquals(6 * GB, app.getAppAttemptResourceUsage().getPending().getMemorySize());
        checkPendingResource(rm, "a1", 6 * GB, null);
        checkPendingResource(rm, "a", 6 * GB, null);
        checkPendingResource(rm, "root", 6 * GB, null);

        // am1 asks to change containerId1 (1G->3G), containerId2 (2G -> 4G) and
        // containerId3 (2G -> 2G)
        am1.sendContainerResizingRequest(Arrays.asList(
                UpdateContainerRequest.newInstance(0, containerId1, ContainerUpdateType.INCREASE_RESOURCE,
                        Resources.createResource(3 * GB)),
                UpdateContainerRequest.newInstance(0, containerId2, ContainerUpdateType.INCREASE_RESOURCE,
                        Resources.createResource(4 * GB)),
                UpdateContainerRequest.newInstance(0, containerId3, ContainerUpdateType.INCREASE_RESOURCE,
                        Resources.createResource(2 * GB))));
        Assert.assertEquals(4 * GB, app.getAppAttemptResourceUsage().getPending().getMemorySize());
        checkPendingResource(rm, "a1", 4 * GB, null);
        checkPendingResource(rm, "a", 4 * GB, null);
        checkPendingResource(rm, "root", 4 * GB, null);
    }

    private void verifyAMLimitForLeafQueue(CapacitySchedulerConfiguration config) throws Exception {
        MockRM rm = setUpMove(config);
        final int nodeMemory = 4 * GB;
        rm.registerNode("127.0.0.1:1234", nodeMemory);

        String queueName = "a1";
        String userName = "user_0";
        final ResourceScheduler scheduler = rm.getRMContext().getScheduler();
        LeafQueue queueA = (LeafQueue) ((CapacityScheduler) scheduler).getQueue(queueName);
        Resource amResourceLimit = queueA.getAMResourceLimit();

        Resource amResource1 = Resource.newInstance(amResourceLimit.getMemorySize() + 1024,
                amResourceLimit.getVirtualCores() + 1);
        Resource amResource2 = Resource.newInstance(amResourceLimit.getMemorySize() + 2048,
                amResourceLimit.getVirtualCores() + 1);

        // Wait for the scheduler to be updated with new node capacity
        GenericTestUtils.waitFor(new Supplier<Boolean>() {
            @Override
            public Boolean get() {
                return scheduler.getMaximumResourceCapability().getMemorySize() == nodeMemory;
            }
        }, 100, 60 * 1000);

        rm.submitApp(amResource1, "app-1", userName, null, queueName);

        rm.submitApp(amResource2, "app-2", userName, null, queueName);

        // When AM limit is exceeded, 1 applications will be activated.Rest all
        // applications will be in pending
        Assert.assertEquals("PendingApplications should be 1", 1, queueA.getNumPendingApplications());
        Assert.assertEquals("Active applications should be 1", 1, queueA.getNumActiveApplications());

        Assert.assertEquals("User PendingApplications should be 1", 1,
                queueA.getUser(userName).getPendingApplications());
        Assert.assertEquals("User Active applications should be 1", 1,
                queueA.getUser(userName).getActiveApplications());
        rm.stop();
    }

    private void setMaxAllocMb(Configuration conf, int maxAllocMb) {
        conf.setInt(YarnConfiguration.RM_SCHEDULER_MAXIMUM_ALLOCATION_MB, maxAllocMb);
    }

    private void setMaxAllocMb(CapacitySchedulerConfiguration conf, String queueName, int maxAllocMb) {
        String propName = CapacitySchedulerConfiguration.getQueuePrefix(queueName)
                + CapacitySchedulerConfiguration.MAXIMUM_ALLOCATION_MB;
        conf.setInt(propName, maxAllocMb);
    }

    private void setMaxAllocVcores(Configuration conf, int maxAllocVcores) {
        conf.setInt(YarnConfiguration.RM_SCHEDULER_MAXIMUM_ALLOCATION_VCORES, maxAllocVcores);
    }

    private void setMaxAllocVcores(CapacitySchedulerConfiguration conf, String queueName, int maxAllocVcores) {
        String propName = CapacitySchedulerConfiguration.getQueuePrefix(queueName)
                + CapacitySchedulerConfiguration.MAXIMUM_ALLOCATION_VCORES;
        conf.setInt(propName, maxAllocVcores);
    }

    private void sentRMContainerLaunched(MockRM rm, ContainerId containerId) {
        CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();
        RMContainer rmContainer = cs.getRMContainer(containerId);
        if (rmContainer != null) {
            rmContainer.handle(new RMContainerEvent(containerId, RMContainerEventType.LAUNCHED));
        } else {
            Assert.fail("Cannot find RMContainer");
        }
    }

    @Test
    public void testResourceUpdateDecommissioningNode() throws Exception {
        // Mock the RMNodeResourceUpdate event handler to update SchedulerNode
        // to have 0 available resource
        RMContext spyContext = Mockito.spy(resourceManager.getRMContext());
        Dispatcher mockDispatcher = mock(AsyncDispatcher.class);
        when(mockDispatcher.getEventHandler()).thenReturn(new EventHandler() {
            @Override
            public void handle(Event event) {
                if (event instanceof RMNodeResourceUpdateEvent) {
                    RMNodeResourceUpdateEvent resourceEvent = (RMNodeResourceUpdateEvent) event;
                    resourceManager.getResourceScheduler().getSchedulerNode(resourceEvent.getNodeId())
                            .updateTotalResource(resourceEvent.getResourceOption().getResource());
                }
            }
        });
        Mockito.doReturn(mockDispatcher).when(spyContext).getDispatcher();
        ((CapacityScheduler) resourceManager.getResourceScheduler()).setRMContext(spyContext);
        ((AsyncDispatcher) mockDispatcher).start();
        // Register node
        String host_0 = "host_0";
        org.apache.hadoop.yarn.server.resourcemanager.NodeManager nm_0 = registerNode(host_0, 1234, 2345,
                NetworkTopology.DEFAULT_RACK, Resources.createResource(8 * GB, 4));
        // ResourceRequest priorities
        Priority priority_0 = org.apache.hadoop.yarn.server.resourcemanager.resource.Priority.create(0);

        // Submit an application
        Application application_0 = new Application("user_0", "a1", resourceManager);
        application_0.submit();

        application_0.addNodeManager(host_0, 1234, nm_0);

        Resource capability_0_0 = Resources.createResource(1 * GB, 1);
        application_0.addResourceRequestSpec(priority_0, capability_0_0);

        Task task_0_0 = new Task(application_0, priority_0, new String[] { host_0 });
        application_0.addTask(task_0_0);

        // Send resource requests to the scheduler
        application_0.schedule();

        nodeUpdate(nm_0);
        // Kick off another heartbeat with the node state mocked to decommissioning
        // This should update the schedulernodes to have 0 available resource
        RMNode spyNode = Mockito.spy(resourceManager.getRMContext().getRMNodes().get(nm_0.getNodeId()));
        when(spyNode.getState()).thenReturn(NodeState.DECOMMISSIONING);
        resourceManager.getResourceScheduler().handle(new NodeUpdateSchedulerEvent(spyNode));

        // Get allocations from the scheduler
        application_0.schedule();

        // Check the used resource is 1 GB 1 core
        Assert.assertEquals(1 * GB, nm_0.getUsed().getMemorySize());
        Resource usedResource = resourceManager.getResourceScheduler()

                .getSchedulerNode(nm_0.getNodeId()).getUsedResource();
        Assert.assertEquals(usedResource.getMemorySize(), 1 * GB);

        Assert.assertEquals(usedResource.getVirtualCores(), 1);
        // Check total resource of scheduler node is also changed to 1 GB 1 core
        Resource totalResource = resourceManager.getResourceScheduler().getSchedulerNode(nm_0.getNodeId())
                .getTotalResource();
        Assert.assertEquals(totalResource.getMemorySize(), 1 * GB);
        Assert.assertEquals(totalResource.getVirtualCores(), 1);
        // Check the available resource is 0/0
        Resource availableResource = resourceManager.getResourceScheduler().getSchedulerNode(nm_0.getNodeId())
                .getAvailableResource();
        Assert.assertEquals(availableResource.getMemorySize(), 0);

        Assert.assertEquals(availableResource.getVirtualCores(), 0);
    }

    @Test
    public void testCSReservationWithRootUnblocked() throws Exception {
        CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
        conf.setResourceComparator(DominantResourceCalculator.class);
        setupOtherBlockedQueueConfiguration(conf);
        conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);
        MockRM rm = new MockRM(conf);
        rm.start();
        CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();
        ParentQueue q = (ParentQueue) cs.getQueue("p1");

        Assert.assertNotNull(q);
        String host = "127.0.0.1";
        String host1 = "test";
        RMNode node = MockNodes.newNodeInfo(0, Resource.newInstance(8 * GB, 8), 1, host);
        RMNode node1 = MockNodes.newNodeInfo(0, Resource.newInstance(8 * GB, 8), 2, host1);
        cs.handle(new NodeAddedSchedulerEvent(node));
        cs.handle(new NodeAddedSchedulerEvent(node1));
        ApplicationAttemptId appAttemptId1 = appHelper(rm, cs, 100, 1, "x1", "userX1");
        ApplicationAttemptId appAttemptId2 = appHelper(rm, cs, 100, 2, "x2", "userX2");
        ApplicationAttemptId appAttemptId3 = appHelper(rm, cs, 100, 3, "y1", "userY1");
        RecordFactory recordFactory = RecordFactoryProvider.getRecordFactory(null);

        Priority priority = TestUtils.createMockPriority(1);
        ResourceRequest y1Req = null;
        ResourceRequest x1Req = null;
        ResourceRequest x2Req = null;
        for (int i = 0; i < 4; i++) {
            y1Req = TestUtils.createResourceRequest(ResourceRequest.ANY, 1 * GB, 1, true, priority, recordFactory);
            cs.allocate(appAttemptId3, Collections.<ResourceRequest>singletonList(y1Req),
                    Collections.<ContainerId>emptyList(), null, null, null, null);
            CapacityScheduler.schedule(cs);
        }
        assertEquals("Y1 Used Resource should be 4 GB", 4 * GB,
                cs.getQueue("y1").getUsedResources().getMemorySize());
        assertEquals("P2 Used Resource should be 4 GB", 4 * GB,
                cs.getQueue("p2").getUsedResources().getMemorySize());

        for (int i = 0; i < 7; i++) {
            x1Req = TestUtils.createResourceRequest(ResourceRequest.ANY, 1 * GB, 1, true, priority, recordFactory);
            cs.allocate(appAttemptId1, Collections.<ResourceRequest>singletonList(x1Req),
                    Collections.<ContainerId>emptyList(), null, null, null, null);
            CapacityScheduler.schedule(cs);
        }
        assertEquals("X1 Used Resource should be 7 GB", 7 * GB,
                cs.getQueue("x1").getUsedResources().getMemorySize());
        assertEquals("P1 Used Resource should be 7 GB", 7 * GB,
                cs.getQueue("p1").getUsedResources().getMemorySize());

        x2Req = TestUtils.createResourceRequest(ResourceRequest.ANY, 2 * GB, 1, true, priority, recordFactory);
        cs.allocate(appAttemptId2, Collections.<ResourceRequest>singletonList(x2Req),
                Collections.<ContainerId>emptyList(), null, null, null, null);
        CapacityScheduler.schedule(cs);
        assertEquals("X2 Used Resource should be 0", 0, cs.getQueue("x2").getUsedResources().getMemorySize());
        assertEquals("P1 Used Resource should be 7 GB", 7 * GB,
                cs.getQueue("p1").getUsedResources().getMemorySize());
        //this assign should fail
        x1Req = TestUtils.createResourceRequest(ResourceRequest.ANY, 1 * GB, 1, true, priority, recordFactory);
        cs.allocate(appAttemptId1, Collections.<ResourceRequest>singletonList(x1Req),
                Collections.<ContainerId>emptyList(), null, null, null, null);
        CapacityScheduler.schedule(cs);
        assertEquals("X1 Used Resource should be 7 GB", 7 * GB,
                cs.getQueue("x1").getUsedResources().getMemorySize());
        assertEquals("P1 Used Resource should be 7 GB", 7 * GB,
                cs.getQueue("p1").getUsedResources().getMemorySize());

        //this should get thru
        for (int i = 0; i < 4; i++) {
            y1Req = TestUtils.createResourceRequest(ResourceRequest.ANY, 1 * GB, 1, true, priority, recordFactory);
            cs.allocate(appAttemptId3, Collections.<ResourceRequest>singletonList(y1Req),
                    Collections.<ContainerId>emptyList(), null, null, null, null);
            CapacityScheduler.schedule(cs);
        }
        assertEquals("P2 Used Resource should be 8 GB", 8 * GB,
                cs.getQueue("p2").getUsedResources().getMemorySize());

        //Free a container from X1
        ContainerId containerId = ContainerId.newContainerId(appAttemptId1, 2);
        cs.handle(new ContainerExpiredSchedulerEvent(containerId));

        //Schedule pending request
        CapacityScheduler.schedule(cs);
        assertEquals("X2 Used Resource should be 2 GB", 2 * GB,
                cs.getQueue("x2").getUsedResources().getMemorySize());
        assertEquals("P1 Used Resource should be 8 GB", 8 * GB,
                cs.getQueue("p1").getUsedResources().getMemorySize());
        assertEquals("P2 Used Resource should be 8 GB", 8 * GB,
                cs.getQueue("p2").getUsedResources().getMemorySize());
        assertEquals("Root Used Resource should be 16 GB", 16 * GB,
                cs.getRootQueue().getUsedResources().getMemorySize());
        rm.stop();
    }

    @Test
    public void testCSQueueBlocked() throws Exception {
        CapacitySchedulerConfiguration conf = new CapacitySchedulerConfiguration();
        setupBlockedQueueConfiguration(conf);
        conf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class);
        MockRM rm = new MockRM(conf);
        rm.start();
        CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();
        LeafQueue q = (LeafQueue) cs.getQueue("a");

        Assert.assertNotNull(q);
        String host = "127.0.0.1";
        String host1 = "test";
        RMNode node = MockNodes.newNodeInfo(0, Resource.newInstance(8 * GB, 8), 1, host);
        RMNode node1 = MockNodes.newNodeInfo(0, Resource.newInstance(8 * GB, 8), 2, host1);
        cs.handle(new NodeAddedSchedulerEvent(node));
        cs.handle(new NodeAddedSchedulerEvent(node1));
        //add app begin
        ApplicationAttemptId appAttemptId1 = appHelper(rm, cs, 100, 1, "a", "user1");
        ApplicationAttemptId appAttemptId2 = appHelper(rm, cs, 100, 2, "b", "user2");
        //add app end

        RecordFactory recordFactory = RecordFactoryProvider.getRecordFactory(null);

        Priority priority = TestUtils.createMockPriority(1);
        ResourceRequest r1 = TestUtils.createResourceRequest(ResourceRequest.ANY, 2 * GB, 1, true, priority,
                recordFactory);
        //This will allocate for app1
        cs.allocate(appAttemptId1, Collections.<ResourceRequest>singletonList(r1),
                Collections.<ContainerId>emptyList(), null, null, null, null).getContainers().size();
        CapacityScheduler.schedule(cs);
        ResourceRequest r2 = null;
        for (int i = 0; i < 13; i++) {
            r2 = TestUtils.createResourceRequest(ResourceRequest.ANY, 1 * GB, 1, true, priority, recordFactory);
            cs.allocate(appAttemptId2, Collections.<ResourceRequest>singletonList(r2),
                    Collections.<ContainerId>emptyList(), null, null, null, null);
            CapacityScheduler.schedule(cs);
        }
        assertEquals("A Used Resource should be 2 GB", 2 * GB, cs.getQueue("a").getUsedResources().getMemorySize());
        assertEquals("B Used Resource should be 13 GB", 13 * GB,
                cs.getQueue("b").getUsedResources().getMemorySize());
        r1 = TestUtils.createResourceRequest(ResourceRequest.ANY, 2 * GB, 1, true, priority, recordFactory);
        r2 = TestUtils.createResourceRequest(ResourceRequest.ANY, 1 * GB, 1, true, priority, recordFactory);
        cs.allocate(appAttemptId1, Collections.<ResourceRequest>singletonList(r1),
                Collections.<ContainerId>emptyList(), null, null, null, null).getContainers().size();
        CapacityScheduler.schedule(cs);

        cs.allocate(appAttemptId2, Collections.<ResourceRequest>singletonList(r2),
                Collections.<ContainerId>emptyList(), null, null, null, null);
        CapacityScheduler.schedule(cs);
        //Check blocked Resource
        assertEquals("A Used Resource should be 2 GB", 2 * GB, cs.getQueue("a").getUsedResources().getMemorySize());
        assertEquals("B Used Resource should be 13 GB", 13 * GB,
                cs.getQueue("b").getUsedResources().getMemorySize());

        ContainerId containerId1 = ContainerId.newContainerId(appAttemptId2, 10);
        ContainerId containerId2 = ContainerId.newContainerId(appAttemptId2, 11);

        cs.handle(new ContainerExpiredSchedulerEvent(containerId1));
        rm.drainEvents();
        CapacityScheduler.schedule(cs);

        cs.handle(new ContainerExpiredSchedulerEvent(containerId2));
        CapacityScheduler.schedule(cs);
        rm.drainEvents();

        assertEquals("A Used Resource should be 4 GB", 4 * GB, cs.getQueue("a").getUsedResources().getMemorySize());
        assertEquals("B Used Resource should be 12 GB", 12 * GB,
                cs.getQueue("b").getUsedResources().getMemorySize());
        assertEquals("Used Resource on Root should be 16 GB", 16 * GB,
                cs.getRootQueue().getUsedResources().getMemorySize());
        rm.stop();
    }

    private ApplicationAttemptId appHelper(MockRM rm, CapacityScheduler cs, int clusterTs, int appId, String queue,
            String user) {
        ApplicationId appId1 = BuilderUtils.newApplicationId(clusterTs, appId);
        ApplicationAttemptId appAttemptId1 = BuilderUtils.newApplicationAttemptId(appId1, appId);

        RMAppAttemptMetrics attemptMetric1 = new RMAppAttemptMetrics(appAttemptId1, rm.getRMContext());
        RMAppImpl app1 = mock(RMAppImpl.class);
        when(app1.getApplicationId()).thenReturn(appId1);
        RMAppAttemptImpl attempt1 = mock(RMAppAttemptImpl.class);
        Container container = mock(Container.class);
        when(attempt1.getMasterContainer()).thenReturn(container);
        ApplicationSubmissionContext submissionContext = mock(ApplicationSubmissionContext.class);
        when(attempt1.getSubmissionContext()).thenReturn(submissionContext);
        when(attempt1.getAppAttemptId()).thenReturn(appAttemptId1);
        when(attempt1.getRMAppAttemptMetrics()).thenReturn(attemptMetric1);
        when(app1.getCurrentAppAttempt()).thenReturn(attempt1);
        rm.getRMContext().getRMApps().put(appId1, app1);

        SchedulerEvent addAppEvent1 = new AppAddedSchedulerEvent(appId1, queue, user);
        cs.handle(addAppEvent1);
        SchedulerEvent addAttemptEvent1 = new AppAttemptAddedSchedulerEvent(appAttemptId1, false);
        cs.handle(addAttemptEvent1);
        return appAttemptId1;
    }

    @Test
    public void testCapacitySchedulerTh() throws Exception {
        for (int j = 20; j < 100; j += 10) {
            float total = 0;
            for (int i = 0; i < 5; i++) {
                total += testCapacitySchedulerThInt(j);
                tearDown();
                setUp();
            }
            float result = total / 5;
            LOG.error("avg nb heartbeats handled per seconds " + j + "\t" + result);
        }
    }

    public float testCapacitySchedulerThInt(int nbTasks) throws Exception {

        LOG.info("--- START: testCapacityScheduler ---");

        // Register node1
        String host_0 = "host_0";
        org.apache.hadoop.yarn.server.resourcemanager.NodeManager nm_0 = registerNode(host_0, 1234, 2345,
                NetworkTopology.DEFAULT_RACK, Resources.createResource((nbTasks + 4) * GB, 1));

        // ResourceRequest priorities
        Priority priority_1 = org.apache.hadoop.yarn.server.resourcemanager.resource.Priority.create(1);

        // Submit an application
        Application application_1 = new Application("user_1", "b2", resourceManager);
        application_1.submit();

        application_1.addNodeManager(host_0, 1234, nm_0);

        Resource capability_1_0 = Resources.createResource(1 * GB, 1);
        application_1.addResourceRequestSpec(priority_1, capability_1_0);

        Task task_1_0 = new Task(application_1, priority_1, new String[] { host_0 });
        application_1.addTask(task_1_0);

        // Send resource requests to the scheduler
        application_1.schedule();

        // Send a heartbeat to kick the tires on the Scheduler
        LOG.info("Kick!");

        // task_0_0 and task_1_0 allocated, used=4G
        nodeUpdate(nm_0);

        // Get allocations from the scheduler

        application_1.schedule(); // task_1_0
        checkApplicationResourceUsage(1 * GB, application_1);

        checkNodeResourceUsage(1 * GB, nm_0);

        long start = System.currentTimeMillis();
        int nbHb = 0;
        long duration = 0;
        long totalTimeInHB = 0;
        do {
            checkApplicationResourceUsage(1 * GB, application_1);
            List<Task> tasks = new ArrayList<>();
            for (int i = 0; i < nbTasks; i++) {
                Task task = new Task(application_1, priority_1, new String[] { ResourceRequest.ANY });
                application_1.addTask(task);
                tasks.add(task);
            }
            application_1.schedule();
            // Send a heartbeat to kick the tires on the Scheduler
            // nothing new, used=4G
            long s = System.nanoTime();
            nodeUpdate(nm_0);
            totalTimeInHB += System.nanoTime() - s;
            nbHb++;

            // task_0_1 is prefer as locality, used=2G

            // Get allocations from the scheduler
            while (!allTasksRunning(tasks)) {
                application_1.schedule();
                //      checkApplicationResourceUsage(5 * GB, application_1);

                s = System.nanoTime();
                nodeUpdate(nm_0);
                totalTimeInHB += System.nanoTime() - s;
                nbHb++;
            }

            for (Task task : tasks) {
                application_1.finishTask(task);
            }
            application_1.schedule();

            s = System.nanoTime();
            nodeUpdate(nm_0);
            totalTimeInHB += System.nanoTime() - s;
            nbHb++;

            //      checkNodeResourceUsage(4 * GB, nm_0);
            //      checkNodeResourceUsage(0 * GB, nm_1);
            duration = System.currentTimeMillis() - start;
        } while (duration < 10000);
        float avgHBDur = (float) totalTimeInHB / nbHb;
        float nbHbps = 1000000000 / avgHBDur;
        LOG.error("nb heartbeats handled per seconds: " + nbHbps + " (" + avgHBDur + ")");
        return nbHbps;
        //
        //    LOG.info("--- END: testCapacityScheduler ---");
    }

    private boolean allTasksRunning(List<Task> tasks) {
        for (Task task : tasks) {
            if (!task.getState().equals(Task.State.RUNNING)) {
                return false;
            }
        }
        return true;
    }
}