AMSimulator extends TaskRunner.Task { protected List<ResourceRequest> packageRequests( List<ContainerSimulator> csList, int priority) { Map<Long, Map<String, ResourceRequest>> rackLocalRequests = new HashMap<>(); Map<Long, Map<String, ResourceRequest>> nodeLocalRequests = new HashMap<>(); Map<Long, ResourceRequest> anyRequests = new HashMap<>(); for (ContainerSimulator cs : csList) { long allocationId = cs.getAllocationId(); ResourceRequest anyRequest = anyRequests.get(allocationId); if (cs.getHostname() != null) { Map<String, ResourceRequest> rackLocalRequestMap; if (rackLocalRequests.containsKey(allocationId)) { rackLocalRequestMap = rackLocalRequests.get(allocationId); } else { rackLocalRequestMap = new HashMap<>(); rackLocalRequests.put(allocationId, rackLocalRequestMap); } String[] rackHostNames = SLSUtils.getRackHostName(cs.getHostname()); String rackname = "/" + rackHostNames[0]; if (rackLocalRequestMap.containsKey(rackname)) { rackLocalRequestMap.get(rackname).setNumContainers( rackLocalRequestMap.get(rackname).getNumContainers() + 1); } else { ResourceRequest request = createResourceRequest(cs.getResource(), cs.getExecutionType(), rackname, priority, cs.getAllocationId(), 1); rackLocalRequestMap.put(rackname, request); } Map<String, ResourceRequest> nodeLocalRequestMap; if (nodeLocalRequests.containsKey(allocationId)) { nodeLocalRequestMap = nodeLocalRequests.get(allocationId); } else { nodeLocalRequestMap = new HashMap<>(); nodeLocalRequests.put(allocationId, nodeLocalRequestMap); } String hostname = rackHostNames[1]; if (nodeLocalRequestMap.containsKey(hostname)) { nodeLocalRequestMap.get(hostname).setNumContainers( nodeLocalRequestMap.get(hostname).getNumContainers() + 1); } else { ResourceRequest request = createResourceRequest(cs.getResource(), cs.getExecutionType(), hostname, priority, cs.getAllocationId(), 1); nodeLocalRequestMap.put(hostname, request); } } if (anyRequest == null) { anyRequest = createResourceRequest(cs.getResource(), cs.getExecutionType(), ResourceRequest.ANY, priority, cs.getAllocationId(), 1); anyRequests.put(allocationId, anyRequest); } else { anyRequest.setNumContainers(anyRequest.getNumContainers() + 1); } } List<ResourceRequest> ask = new ArrayList<ResourceRequest>(); for (Map<String, ResourceRequest> nodeLocalRequestMap : nodeLocalRequests.values()) { ask.addAll(nodeLocalRequestMap.values()); } for (Map<String, ResourceRequest> rackLocalRequestMap : rackLocalRequests.values()) { ask.addAll(rackLocalRequestMap.values()); } ask.addAll(anyRequests.values()); return ask; } AMSimulator(); @SuppressWarnings("checkstyle:parameternumber") void init(int heartbeatInterval,
      List<ContainerSimulator> containerList, ResourceManager resourceManager,
      SLSRunner slsRunnner, long startTime, long finishTime, String simUser,
      String simQueue, boolean tracked, String oldApp, long baseTimeMS,
      Resource amResource, String nodeLabelExpr, Map<String, String> params,
      Map<ApplicationId, AMSimulator> appIdAMSim); @Override void firstStep(); synchronized void notifyAMContainerLaunched(Container masterContainer); @Override void middleStep(); @Override void lastStep(); abstract void initReservation(
      ReservationId reservationId, long deadline, long now); void untrackApp(); String getQueue(); String getAMType(); long getDuration(); int getNumTasks(); ApplicationId getApplicationId(); ApplicationAttemptId getApplicationAttemptId();  }
@Test public void testPackageRequests() { MockAMSimulator app = new MockAMSimulator(); List<ContainerSimulator> containerSimulators = new ArrayList<>(); Resource resource = Resources.createResource(1024); int priority = 1; ExecutionType execType = ExecutionType.GUARANTEED; String type = "map"; ContainerSimulator s1 = new ContainerSimulator(resource, 100, "/default-rack/h1", priority, type, execType); ContainerSimulator s2 = new ContainerSimulator(resource, 100, "/default-rack/h1", priority, type, execType); ContainerSimulator s3 = new ContainerSimulator(resource, 100, "/default-rack/h2", priority, type, execType); containerSimulators.add(s1); containerSimulators.add(s2); containerSimulators.add(s3); List<ResourceRequest> res = app.packageRequests(containerSimulators, priority); Assert.assertEquals(4, res.size()); int anyRequestCount = 0; int rackRequestCount = 0; int nodeRequestCount = 0; for (ResourceRequest request : res) { String resourceName = request.getResourceName(); if (resourceName.equals("*")) { anyRequestCount++; } else if (resourceName.equals("/default-rack")) { rackRequestCount++; } else { nodeRequestCount++; } } Assert.assertEquals(1, anyRequestCount); Assert.assertEquals(1, rackRequestCount); Assert.assertEquals(2, nodeRequestCount); containerSimulators.clear(); s1 = new ContainerSimulator(resource, 100, "/default-rack/h1", priority, type, execType, 1, 0); s2 = new ContainerSimulator(resource, 100, "/default-rack/h1", priority, type, execType, 2, 0); s3 = new ContainerSimulator(resource, 100, "/default-rack/h2", priority, type, execType, 1, 0); containerSimulators.add(s1); containerSimulators.add(s2); containerSimulators.add(s3); res = app.packageRequests(containerSimulators, priority); Assert.assertEquals(7, res.size()); anyRequestCount = 0; rackRequestCount = 0; nodeRequestCount = 0; for (ResourceRequest request : res) { String resourceName = request.getResourceName(); long allocationId = request.getAllocationRequestId(); Assert.assertTrue(allocationId == 1 || allocationId == 2); if (resourceName.equals("*")) { anyRequestCount++; } else if (resourceName.equals("/default-rack")) { rackRequestCount++; } else { nodeRequestCount++; } } Assert.assertEquals(2, anyRequestCount); Assert.assertEquals(2, rackRequestCount); Assert.assertEquals(3, nodeRequestCount); }