public List<Container> allocateAndWaitForContainers(int nContainer,
    int memory, MockNM nm) throws Exception {
  // AM request for containers
  allocate("ANY", memory, nContainer, null);
  // kick the scheduler
  nm.nodeHeartbeat(true);
  List<Container> conts =
      allocate(new ArrayList<ResourceRequest>(), null)
          .getAllocatedContainers();
  while (conts.size() < nContainer) {
    nm.nodeHeartbeat(true);
    conts.addAll(allocate(new ArrayList<ResourceRequest>(),
        new ArrayList<ContainerId>()).getAllocatedContainers());
    Thread.sleep(500);
  }
  return conts;
}
 

private static void normalizeNodeLabelExpressionInRequest(
    ResourceRequest resReq, QueueInfo queueInfo) {

  String labelExp = resReq.getNodeLabelExpression();

  // if queue has default label expression, and RR doesn't have, use the
  // default label expression of queue
  if (labelExp == null && queueInfo != null && ResourceRequest.ANY
      .equals(resReq.getResourceName())) {
    labelExp = queueInfo.getDefaultNodeLabelExpression();
  }

  // If labelExp still equals to null, set it to be NO_LABEL
  if (labelExp == null) {
    labelExp = RMNodeLabelsManager.NO_LABEL;
  }
  resReq.setNodeLabelExpression(labelExp);
}
 

private Resource getTotalResource(List<ResourceRequest> requests) {
  Resource totalResource = Resource.newInstance(0, 0, 0);
  if (requests == null) {
    return totalResource;
  }
  for (ResourceRequest request : requests) {
    if (request.getNumContainers() == 0) {
      continue;
    }
    if (request.getResourceName().equals(ResourceRequest.ANY)) {
      Resources.addTo(
        totalResource,
        Resources.multiply(request.getCapability(),
          request.getNumContainers()));
    }
  }
  return totalResource;
}
 

@Public
@Stable
public static AllocateRequest newInstance(int responseID, float appProgress,
    List<ResourceRequest> resourceAsk,
    List<ContainerId> containersToBeReleased,
    ResourceBlacklistRequest resourceBlacklistRequest,
    List<ContainerResourceIncreaseRequest> increaseRequests) {
  AllocateRequest allocateRequest = Records.newRecord(AllocateRequest.class);
  allocateRequest.setResponseId(responseID);
  allocateRequest.setProgress(appProgress);
  allocateRequest.setAskList(resourceAsk);
  allocateRequest.setReleaseList(containersToBeReleased);
  allocateRequest.setResourceBlacklistRequest(resourceBlacklistRequest);
  allocateRequest.setIncreaseRequests(increaseRequests);
  return allocateRequest;
}
 

private LogAggregationContext getLogAggregationContextFromContainerToken(
    MockRM rm1, MockNM nm1, LogAggregationContext logAggregationContext)
    throws Exception {
  RMApp app2 = rm1.submitApp(200, logAggregationContext);
  MockAM am2 = MockRM.launchAndRegisterAM(app2, rm1, nm1);
  nm1.nodeHeartbeat(true);
  // request a container.
  am2.allocate("127.0.0.1", 512, 1, new ArrayList<ContainerId>());
  ContainerId containerId =
      ContainerId.newContainerId(am2.getApplicationAttemptId(), 2);
  rm1.waitForState(nm1, containerId, RMContainerState.ALLOCATED);

  // acquire the container.
  List<Container> containers =
      am2.allocate(new ArrayList<ResourceRequest>(),
        new ArrayList<ContainerId>()).getAllocatedContainers();
  Assert.assertEquals(containerId, containers.get(0).getId());
  // container token is generated.
  Assert.assertNotNull(containers.get(0).getContainerToken());
  ContainerTokenIdentifier token =
      BuilderUtils.newContainerTokenIdentifier(containers.get(0)
        .getContainerToken());
  return token.getLogAggregationContext();
}
 

@Override
public void updateDemand() {
  demand = Resources.createResource(0);
  // Demand is current consumption plus outstanding requests
  Resources.addTo(demand, getCurrentConsumption());

  // Add up outstanding resource requests
  synchronized (this) {
    for (Priority p : getPriorities()) {
      for (ResourceRequest r : getResourceRequests(p).values()) {
        Resource total = Resources.multiply(r.getCapability(), r.getNumContainers());
        Resources.addTo(demand, total);
      }
    }
  }
}
 

public Task(Application application, Priority priority, String[] hosts) {
  this.applicationId = application.getApplicationId();
  this.priority = priority;
  
  taskId = application.getNextTaskId();
  state = State.PENDING;
  
  // Special case: Don't care about locality
  if (!(hosts.length == 1 && 
      hosts[0].equals(ResourceRequest.ANY))) {
    for (String host : hosts) {
      this.hosts.add(host);
      this.racks.add(Application.resolve(host));
    }
  }
  LOG.info("Task " + taskId + " added to application " + this.applicationId + 
      " with " + this.hosts.size() + " hosts, " + racks.size() + " racks");
}
 

public static void normalizeAndValidateRequest(ResourceRequest resReq,
    Resource maximumResource, String queueName, YarnScheduler scheduler,
    boolean isRecovery, RMContext rmContext, QueueInfo queueInfo)
    throws InvalidResourceRequestException {

  if (queueInfo == null) {
    try {
      queueInfo = scheduler.getQueueInfo(queueName, false, false);
    } catch (IOException e) {
      // it is possible queue cannot get when queue mapping is set, just ignore
      // the queueInfo here, and move forward
    }
  }
  SchedulerUtils.normalizeNodeLabelExpressionInRequest(resReq, queueInfo);
  if (!isRecovery) {
    validateResourceRequest(resReq, maximumResource, queueInfo, rmContext);
  }
}
 

public List<ResourceRequest> createReq(String[] hosts, int memory, int priority,
    int containers, String labelExpression) throws Exception {
  List<ResourceRequest> reqs = new ArrayList<ResourceRequest>();
  if (hosts != null) {
    for (String host : hosts) {
      // only add host/rack request when asked host isn't ANY
      if (!host.equals(ResourceRequest.ANY)) {
        ResourceRequest hostReq =
            createResourceReq(host, memory, priority, containers,
                labelExpression);
        reqs.add(hostReq);
        ResourceRequest rackReq =
            createResourceReq("/default-rack", memory, priority, containers,
                labelExpression);
        reqs.add(rackReq);
      }
    }
  }

  ResourceRequest offRackReq = createResourceReq(ResourceRequest.ANY, memory,
      priority, containers, labelExpression);
  reqs.add(offRackReq);
  return reqs;
}
 

/**
 * Whether this app has containers requests that could be satisfied on the
 * given node, if the node had full space.
 */
public boolean hasContainerForNode(Priority prio, FSSchedulerNode node) {
  ResourceRequest anyRequest = getResourceRequest(prio, ResourceRequest.ANY);
  ResourceRequest rackRequest = getResourceRequest(prio, node.getRackName());
  ResourceRequest nodeRequest = getResourceRequest(prio, node.getNodeName());

  return
      // There must be outstanding requests at the given priority:
      anyRequest != null && anyRequest.getNumContainers() > 0 &&
          // If locality relaxation is turned off at *-level, there must be a
          // non-zero request for the node's rack:
          (anyRequest.getRelaxLocality() ||
              (rackRequest != null && rackRequest.getNumContainers() > 0)) &&
          // If locality relaxation is turned off at rack-level, there must be a
          // non-zero request at the node:
          (rackRequest == null || rackRequest.getRelaxLocality() ||
              (nodeRequest != null && nodeRequest.getNumContainers() > 0)) &&
          // The requested container must be able to fit on the node:
          Resources.lessThanOrEqual(RESOURCE_CALCULATOR, null,
              anyRequest.getCapability(), node.getRMNode().getTotalCapability());
}
 

public static ResourceRequest newResourceRequest(Priority priority,
    String hostName, Resource capability, int numContainers) {
  ResourceRequest request = recordFactory
      .newRecordInstance(ResourceRequest.class);
  request.setPriority(priority);
  request.setResourceName(hostName);
  request.setCapability(capability);
  request.setNumContainers(numContainers);
  return request;
}
 

private void addAsksToProto() {
  maybeInitBuilder();
  builder.clearAsk();
  if (ask == null)
    return;
  Iterable<ResourceRequestProto> iterable = new Iterable<ResourceRequestProto>() {
    @Override
    public Iterator<ResourceRequestProto> iterator() {
      return new Iterator<ResourceRequestProto>() {

        Iterator<ResourceRequest> iter = ask.iterator();

        @Override
        public boolean hasNext() {
          return iter.hasNext();
        }

        @Override
        public ResourceRequestProto next() {
          return convertToProtoFormat(iter.next());
        }

        @Override
        public void remove() {
          throw new UnsupportedOperationException();

        }
      };

    }
  };
  builder.addAllAsk(iterable);
}
 

synchronized private void decrementOutstanding(
    ResourceRequest offSwitchRequest) {
  int numOffSwitchContainers = offSwitchRequest.getNumContainers() - 1;

  // Do not remove ANY
  offSwitchRequest.setNumContainers(numOffSwitchContainers);
  
  // Do we have any outstanding requests?
  // If there is nothing, we need to deactivate this application
  if (numOffSwitchContainers == 0) {
    checkForDeactivation();
  }
}
 

private Resource assignNodeLocalContainers(Resource clusterResource,
    ResourceRequest nodeLocalResourceRequest, FiCaSchedulerNode node,
    FiCaSchedulerApp application, Priority priority,
    RMContainer reservedContainer, MutableObject allocatedContainer,
    ResourceLimits currentResoureLimits) {
  if (canAssign(application, priority, node, NodeType.NODE_LOCAL, 
      reservedContainer)) {
    return assignContainer(clusterResource, node, application, priority,
        nodeLocalResourceRequest, NodeType.NODE_LOCAL, reservedContainer,
        allocatedContainer, currentResoureLimits);
  }
  
  return Resources.none();
}
 

@Override
public void setAMContainerResourceRequest(ResourceRequest request) {
  maybeInitBuilder();
  if (request == null) {
    builder.clearAmContainerResourceRequest();
  }
  this.amResourceRequest = request;
}
 

public static ResourceRequest newResourceRequest(Priority priority,
    String hostName, Resource capability, int numContainers) {
  ResourceRequest request = recordFactory
      .newRecordInstance(ResourceRequest.class);
  request.setPriority(priority);
  request.setResourceName(hostName);
  request.setCapability(capability);
  request.setNumContainers(numContainers);
  return request;
}
 

String getMatchingLocation() {
  switch (state) {
  case MATCHING_LOCAL:
    return getHost();
  case MATCHING_RACK:
    return getRack();
  case MATCHING_ANY:
    return ResourceRequest.ANY;
  default:
    throw new IllegalStateException("Container " + getId() + " trying to match in state " + state);
  }
}
 

@Override
public List<ResourceRequest> getResourceRequests() {
  try {
    readLock.lock();
    return resourceRequests;
  } finally {
    readLock.unlock();
  }
}
 

@Public
@Stable
public static AllocateRequest newInstance(int responseID, float appProgress,
    List<ResourceRequest> resourceAsk,
    List<ContainerId> containersToBeReleased,
    ResourceBlacklistRequest resourceBlacklistRequest) {
  return newInstance(responseID, appProgress, resourceAsk,
      containersToBeReleased, resourceBlacklistRequest, null);
}
 

@GuardedBy("DagAwareYarnTaskScheduler.this")
void add(HeldContainer hc) {
  add(hc, hc.getHost());
  add(hc, hc.getRack());
  add(hc, ResourceRequest.ANY);
  ++numContainers;
}
 

private ResourceRequest createResourceRequest(int memory, String host,
    int priority, int numContainers) {
  ResourceRequest request = recordFactory
      .newRecordInstance(ResourceRequest.class);
  request.setCapability(Resources.createResource(memory));
  request.setResourceName(host);
  request.setNumContainers(numContainers);
  Priority prio = recordFactory.newRecordInstance(Priority.class);
  prio.setPriority(priority);
  request.setPriority(prio);
  return request;
}
 

protected AllocateRequest createAllocateRequest(List<ResourceRequest> ask,
    List<ContainerId> toRelease) {
  AllocateRequest allocateRequest =
          recordFactory.newRecordInstance(AllocateRequest.class);
  allocateRequest.setResponseId(RESPONSE_ID ++);
  allocateRequest.setAskList(ask);
  allocateRequest.setReleaseList(toRelease);
  return allocateRequest;
}
 

private int getMaxAllocatableContainers(FiCaSchedulerApp application,
    Priority priority, FiCaSchedulerNode node, NodeType type) {
  int maxContainers = 0;
  
  ResourceRequest offSwitchRequest = 
    application.getResourceRequest(priority, ResourceRequest.ANY);
  if (offSwitchRequest != null) {
    maxContainers = offSwitchRequest.getNumContainers();
  }

  if (type == NodeType.OFF_SWITCH) {
    return maxContainers;
  }

  if (type == NodeType.RACK_LOCAL) {
    ResourceRequest rackLocalRequest = 
      application.getResourceRequest(priority, node.getRMNode().getRackName());
    if (rackLocalRequest == null) {
      return maxContainers;
    }

    maxContainers = Math.min(maxContainers, rackLocalRequest.getNumContainers());
  }

  if (type == NodeType.NODE_LOCAL) {
    ResourceRequest nodeLocalRequest = 
      application.getResourceRequest(priority, node.getRMNode().getNodeAddress());
    if (nodeLocalRequest != null) {
      maxContainers = Math.min(maxContainers, nodeLocalRequest.getNumContainers());
    }
  }

  return maxContainers;
}
 

@GuardedBy("this")
@Nullable
private TaskRequest tryAssignReuseContainerAppRunning(HeldContainer hc) {
  if (!hc.isAssignable()) {
    LOG.debug("Skipping scheduling of container {} because it state is {}", hc.getId(), hc.getState());
    return null;
  }

  TaskRequest assignedRequest = tryAssignReuseContainerForAffinity(hc);
  if (assignedRequest != null) {
    return assignedRequest;
  }

  for (Entry<Priority,RequestPriorityStats> entry : requestTracker.getStatsEntries()) {
    Priority priority = entry.getKey();
    RequestPriorityStats stats = entry.getValue();
    if (!stats.allowedVertices.intersects(stats.vertices)) {
      LOG.debug("Skipping requests at priority {} because all requesting vertices are blocked by higher priority requests",
          priority);
      continue;
    }

    String matchLocation = hc.getMatchingLocation();
    if (stats.localityCount <= 0) {
      LOG.debug("Overriding locality match of container {} to ANY since there are no locality requests at priority {}",
          hc.getId(), priority);
      matchLocation = ResourceRequest.ANY;
    }
    assignedRequest = tryAssignReuseContainerForPriority(hc, matchLocation,
        priority, stats.allowedVertices);
    if (assignedRequest != null) {
      break;
    }
  }
  return assignedRequest;
}
 

private boolean checkResourceRequestMatchingNodeLabel(ResourceRequest offswitchResourceRequest,
    FiCaSchedulerNode node) {
  String askedNodeLabel = offswitchResourceRequest.getNodeLabelExpression();
  if (null == askedNodeLabel) {
    askedNodeLabel = RMNodeLabelsManager.NO_LABEL;
  }
  return askedNodeLabel.equals(node.getPartition());
}
 

@Override
public ResourceRequest getAMContainerResourceRequest() {
  ApplicationSubmissionContextProtoOrBuilder p = viaProto ? proto : builder;
  if (this.amResourceRequest != null) {
    return amResourceRequest;
  } // Else via proto
  if (!p.hasAmContainerResourceRequest()) {
    return null;
  }
  amResourceRequest = convertFromProtoFormat(p.getAmContainerResourceRequest());
  return amResourceRequest;
}
 

ResourceRequestInfo(Priority priority, String resourceName,
    Resource capability, boolean relaxLocality) {
  remoteRequest = ResourceRequest.newInstance(priority, resourceName,
      capability, 0);
  remoteRequest.setRelaxLocality(relaxLocality);
  containerRequests = new LinkedHashSet<T>();
}
 

public static ResourceRequest newResourceRequest(ResourceRequest r) {
  ResourceRequest request = recordFactory
      .newRecordInstance(ResourceRequest.class);
  request.setPriority(r.getPriority());
  request.setResourceName(r.getResourceName());
  request.setCapability(r.getCapability());
  request.setNumContainers(r.getNumContainers());
  return request;
}
 

private Resource assignNodeLocalContainers(Resource clusterResource,
    ResourceRequest nodeLocalResourceRequest, FiCaSchedulerNode node,
    FiCaSchedulerApp application, Priority priority,
    RMContainer reservedContainer, MutableObject allocatedContainer,
    ResourceLimits currentResoureLimits) {
  if (canAssign(application, priority, node, NodeType.NODE_LOCAL, 
      reservedContainer)) {
    return assignContainer(clusterResource, node, application, priority,
        nodeLocalResourceRequest, NodeType.NODE_LOCAL, reservedContainer,
        allocatedContainer, currentResoureLimits);
  }
  
  return Resources.none();
}
 

@Test
public void testContainerTokenGeneratedOnPullRequest() throws Exception {
  MockRM rm1 = new MockRM(conf);
  rm1.start();
  MockNM nm1 = rm1.registerNode("127.0.0.1:1234", 8000);
  RMApp app1 = rm1.submitApp(200);
  MockAM am1 = MockRM.launchAndRegisterAM(app1, rm1, nm1);
  // request a container.
  am1.allocate("127.0.0.1", 1024, 1, new ArrayList<ContainerId>());
  ContainerId containerId2 =
      ContainerId.newContainerId(am1.getApplicationAttemptId(), 2);
  rm1.waitForState(nm1, containerId2, RMContainerState.ALLOCATED);

  RMContainer container =
      rm1.getResourceScheduler().getRMContainer(containerId2);
  // no container token is generated.
  Assert.assertEquals(containerId2, container.getContainerId());
  Assert.assertNull(container.getContainer().getContainerToken());

  // acquire the container.
  List<Container> containers =
      am1.allocate(new ArrayList<ResourceRequest>(),
        new ArrayList<ContainerId>()).getAllocatedContainers();
  Assert.assertEquals(containerId2, containers.get(0).getId());
  // container token is generated.
  Assert.assertNotNull(containers.get(0).getContainerToken());
  rm1.stop();
}