CapacityScheduler extends
    AbstractYarnScheduler<FiCaSchedulerApp, FiCaSchedulerNode> implements
    PreemptableResourceScheduler, CapacitySchedulerContext, Configurable { @Override @Lock(Lock.NoLock.class) public Allocation allocate(ApplicationAttemptId applicationAttemptId, List<ResourceRequest> ask, List<ContainerId> release, List<String> blacklistAdditions, List<String> blacklistRemovals, List<ContainerResourceChangeRequest> increaseRequests, List<ContainerResourceChangeRequest> decreaseRequests) { FiCaSchedulerApp application = getApplicationAttempt(applicationAttemptId); if (application == null) { return EMPTY_ALLOCATION; } SchedulerUtils.normalizeRequests( ask, getResourceCalculator(), getClusterResource(), getMinimumResourceCapability(), getMaximumResourceCapability()); List<SchedContainerChangeRequest> normalizedIncreaseRequests = checkAndNormalizeContainerChangeRequests(increaseRequests, true); List<SchedContainerChangeRequest> normalizedDecreaseRequests = checkAndNormalizeContainerChangeRequests(decreaseRequests, false); releaseContainers(release, application); Allocation allocation; LeafQueue updateDemandForQueue = null; synchronized (application) { if (application.isStopped()) { return EMPTY_ALLOCATION; } if (!ask.isEmpty()) { if(LOG.isDebugEnabled()) { LOG.debug("allocate: pre-update " + applicationAttemptId + " ask size =" + ask.size()); application.showRequests(); } if (application.updateResourceRequests(ask)) { updateDemandForQueue = (LeafQueue) application.getQueue(); } if(LOG.isDebugEnabled()) { LOG.debug("allocate: post-update"); application.showRequests(); } } if (application.updateIncreaseRequests(normalizedIncreaseRequests) && (updateDemandForQueue == null)) { updateDemandForQueue = (LeafQueue) application.getQueue(); } if (application.isWaitingForAMContainer()) { application.updateAMBlacklist( blacklistAdditions, blacklistRemovals); } else { application.updateBlacklist(blacklistAdditions, blacklistRemovals); } decreaseContainers(normalizedDecreaseRequests, application); allocation = application.getAllocation(getResourceCalculator(), clusterResource, getMinimumResourceCapability()); } if (updateDemandForQueue != null) { updateDemandForQueue.getOrderingPolicy().demandUpdated(application); } return allocation; } CapacityScheduler(); @Override void setConf(Configuration conf); @Override Configuration getConf(); @Override QueueMetrics getRootQueueMetrics(); CSQueue getRootQueue(); @Override CapacitySchedulerConfiguration getConfiguration(); @Override synchronized RMContainerTokenSecretManager getContainerTokenSecretManager(); @Override ResourceCalculator getResourceCalculator(); @Override Comparator<CSQueue> getNonPartitionedQueueComparator(); @Override PartitionedQueueComparator getPartitionedQueueComparator(); @Override int getNumClusterNodes(); @Override synchronized RMContext getRMContext(); @Override synchronized void setRMContext(RMContext rmContext); @Override void serviceInit(Configuration conf); @Override void serviceStart(); @Override void serviceStop(); @Override synchronized void reinitialize(Configuration conf, RMContext rmContext); @VisibleForTesting synchronized UserGroupMappingPlacementRule getUserGroupMappingPlacementRule(); @VisibleForTesting static void setQueueAcls(YarnAuthorizationProvider authorizer,
      Map<String, CSQueue> queues); CSQueue getQueue(String queueName); @Override // Note: when AM asks to decrease container or release container, we will // acquire scheduler lock @Lock(Lock.NoLock.class) Allocation allocate(ApplicationAttemptId applicationAttemptId,
      List<ResourceRequest> ask, List<ContainerId> release,
      List<String> blacklistAdditions, List<String> blacklistRemovals,
      List<ContainerResourceChangeRequest> increaseRequests,
      List<ContainerResourceChangeRequest> decreaseRequests); @Override @Lock(Lock.NoLock.class) QueueInfo getQueueInfo(String queueName, 
      boolean includeChildQueues, boolean recursive); @Override @Lock(Lock.NoLock.class) List<QueueUserACLInfo> getQueueUserAclInfo(); @Override void handle(SchedulerEvent event); @Lock(Lock.NoLock.class) @VisibleForTesting @Override FiCaSchedulerApp getApplicationAttempt(
      ApplicationAttemptId applicationAttemptId); @Lock(Lock.NoLock.class) FiCaSchedulerNode getNode(NodeId nodeId); @Override @Lock(Lock.NoLock.class) void recover(RMState state); @Override void dropContainerReservation(RMContainer container); @Override void preemptContainer(ApplicationAttemptId aid, RMContainer cont); @Override void killContainer(RMContainer cont); @Override synchronized boolean checkAccess(UserGroupInformation callerUGI,
      QueueACL acl, String queueName); @Override List<ApplicationAttemptId> getAppsInQueue(String queueName); @Override synchronized void removeQueue(String queueName); @Override synchronized void addQueue(Queue queue); @Override synchronized void setEntitlement(String inQueue,
      QueueEntitlement entitlement); @Override synchronized String moveApplication(ApplicationId appId,
      String targetQueueName); @Override EnumSet<SchedulerResourceTypes> getSchedulingResourceTypes(); @Override Resource getMaximumResourceCapability(String queueName); @Override Set<String> getPlanQueues(); SchedulerHealth getSchedulerHealth(); @Override Priority checkAndGetApplicationPriority(Priority priorityFromContext,
      String user, String queueName, ApplicationId applicationId); @Override void updateApplicationPriority(Priority newPriority,
      ApplicationId applicationId); @Private
  static final String ROOT_QUEUE; }
@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); 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); final CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler(); 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 e) { e.printStackTrace(); } catch (BrokenBarrierException e) { e.printStackTrace(); } } } }); otherThread.start(); barrier.await(); AllocateRequest allocateRequest = AllocateRequest.newInstance(0, 0.0f, null, null, null); client.allocate(allocateRequest); barrier.await(); otherThread.join(); rm.stop(); }
@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 = new MockNM("h1:1234", 200 * GB, rm.getResourceTrackerService()); nm1.registerNode(); RMApp app1 = rm.submitApp(1 * GB, "app", "user", null, "a1"); MockAM am1 = MockRM.launchAndRegisterAM(app1, rm, nm1); 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)); 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.sendContainerResizingRequest(Arrays.asList( ContainerResourceChangeRequest .newInstance(containerId1, Resources.createResource(3 * GB))), null); FiCaSchedulerApp app = getFiCaSchedulerApp(rm, app1.getApplicationId()); Assert.assertEquals(2 * GB, app.getAppAttemptResourceUsage().getPending().getMemory()); checkPendingResource(rm, "a1", 2 * GB, null); checkPendingResource(rm, "a", 2 * GB, null); checkPendingResource(rm, "root", 2 * GB, null); am1.sendContainerResizingRequest(Arrays.asList( ContainerResourceChangeRequest .newInstance(containerId2, Resources.createResource(3 * GB)), ContainerResourceChangeRequest .newInstance(containerId3, Resources.createResource(5 * GB))), null); Assert.assertEquals(6 * GB, app.getAppAttemptResourceUsage().getPending().getMemory()); checkPendingResource(rm, "a1", 6 * GB, null); checkPendingResource(rm, "a", 6 * GB, null); checkPendingResource(rm, "root", 6 * GB, null); am1.sendContainerResizingRequest(Arrays.asList( ContainerResourceChangeRequest .newInstance(containerId1, Resources.createResource(3 * GB)), ContainerResourceChangeRequest .newInstance(containerId2, Resources.createResource(4 * GB)), ContainerResourceChangeRequest .newInstance(containerId3, Resources.createResource(2 * GB))), null); Assert.assertEquals(4 * GB, app.getAppAttemptResourceUsage().getPending().getMemory()); checkPendingResource(rm, "a1", 4 * GB, null); checkPendingResource(rm, "a", 4 * GB, null); checkPendingResource(rm, "root", 4 * GB, null); }