YarnNodeCapacityManager extends BaseInterceptor { @VisibleForTesting protected void handleContainerAllocation(RMNode rmNode) { String host = rmNode.getNodeID().getHost(); ConsumedOffer consumedOffer = offerLifecycleMgr.drainConsumedOffer(host); if (consumedOffer == null) { LOGGER.debug("No offer consumed for {}", host); return; } Node node = nodeStore.getNode(host); Set<RMContainer> containersBeforeSched = node.getContainerSnapshot(); Set<RMContainer> containersAfterSched = new HashSet<>(node.getNode().getRunningContainers()); Set<RMContainer> containersAllocatedByMesosOffer = (containersBeforeSched == null) ? containersAfterSched : Sets.difference( containersAfterSched, containersBeforeSched); if (containersAllocatedByMesosOffer.isEmpty()) { LOGGER.debug("No containers allocated using Mesos offers for host: {}", host); for (Protos.Offer offer : consumedOffer.getOffers()) { offerLifecycleMgr.declineOffer(offer); } decrementNodeCapacity(rmNode, OfferUtils.getYarnResourcesFromMesosOffers(consumedOffer.getOffers())); } else { LOGGER.debug("Containers allocated using Mesos offers for host: {} count: {}", host, containersAllocatedByMesosOffer.size()); List<Protos.TaskInfo> tasks = Lists.newArrayList(); Resource resUsed = Resource.newInstance(0, 0); for (RMContainer newContainer : containersAllocatedByMesosOffer) { tasks.add(getTaskInfoForContainer(newContainer, consumedOffer, node)); resUsed = Resources.add(resUsed, newContainer.getAllocatedResource()); } Resource resOffered = OfferUtils.getYarnResourcesFromMesosOffers(consumedOffer.getOffers()); Resource resUnused = Resources.subtract(resOffered, resUsed); decrementNodeCapacity(rmNode, resUnused); myriadDriver.getDriver().launchTasks(consumedOffer.getOfferIds(), tasks); } node.removeContainerSnapshot(); } @Inject  YarnNodeCapacityManager(InterceptorRegistry registry, AbstractYarnScheduler yarnScheduler, RMContext rmContext,
                                 MyriadDriver myriadDriver, OfferLifecycleManager offerLifecycleMgr, NodeStore nodeStore,
                                 SchedulerState state, TaskUtils taskUtils); @Override CallBackFilter getCallBackFilter(); @Override void beforeReleaseContainers(List<ContainerId> containerIds, SchedulerApplicationAttempt attempt); @Override void beforeCompletedContainer(RMContainer rmContainer, ContainerStatus containerStatus, RMContainerEventType type); @Override void afterSchedulerEventHandled(SchedulerEvent event); void incrementNodeCapacity(RMNode rmNode, Resource addedCapacity); void decrementNodeCapacity(RMNode rmNode, Resource removedCapacity); @SuppressWarnings("unchecked") void setNodeCapacity(RMNode rmNode, Resource newCapacity);  }
@Test public void testHandleContainerAllocation() throws Exception { Offer offer = TestObjectFactory.getOffer("zero-localhost-one", "slave-one", "mock-framework", "offer-one", 0.1, 512.0); sNodeOne.allocateContainer(containerOne); NodeTask task = TestObjectFactory.getNodeTask("small", "localhost-one", Double.valueOf(0.1), Double.valueOf(512.0), Long.parseLong("1"), Long.parseLong("256")); state.addNodes(Lists.newArrayList(task)); olManager.addOffers(offer); olManager.markAsConsumed(offer); manager.handleContainerAllocation(nodeOne); store.getNode("localhost-one").snapshotRunningContainers(); assertEquals(1, store.getNode("localhost-one").getNode().getRunningContainers().size()); assertEquals(1, store.getNode("localhost-one").getContainerSnapshot().size()); }