List of usage examples for org.apache.hadoop.yarn.api.records Priority newInstance
@Public @Stable public static Priority newInstance(int p)
From source file:org.apache.tez.dag.app.rm.TestDagAwareYarnTaskScheduler.java
License:Apache License
@Test(timeout = 50000) public void testIdleContainerAssignmentReuseNewContainers() throws Exception { AMRMClientAsyncWrapperForTest mockRMClient = spy(new AMRMClientAsyncWrapperForTest()); String appHost = "host"; int appPort = 0; String appUrl = "url"; Configuration conf = new Configuration(); conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_ENABLED, true); conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_REUSE_LOCALITY_DELAY_ALLOCATION_MILLIS, 100); conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_RACK_FALLBACK_ENABLED, false); conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NON_LOCAL_FALLBACK_ENABLED, false); conf.setBoolean(TezConfiguration.TEZ_AM_CONTAINER_REUSE_NEW_CONTAINERS_ENABLED, true); conf.setInt(TezConfiguration.TEZ_AM_RM_HEARTBEAT_INTERVAL_MS_MAX, 100); conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MIN_MILLIS, 4000); conf.setInt(TezConfiguration.TEZ_AM_CONTAINER_IDLE_RELEASE_TIMEOUT_MAX_MILLIS, 5000); conf.setInt(TezConfiguration.TEZ_AM_SESSION_MIN_HELD_CONTAINERS, 5); DagInfo mockDagInfo = mock(DagInfo.class); when(mockDagInfo.getTotalVertices()).thenReturn(10); when(mockDagInfo.getVertexDescendants(anyInt())).thenReturn(new BitSet()); TaskSchedulerContext mockApp = setupMockTaskSchedulerContext(appHost, appPort, appUrl, conf); when(mockApp.getCurrentDagInfo()).thenReturn(mockDagInfo); when(mockApp.isSession()).thenReturn(true); TaskSchedulerContextDrainable drainableAppCallback = createDrainableContext(mockApp); MockClock clock = new MockClock(1000); NewTaskSchedulerForTest scheduler = new NewTaskSchedulerForTest(drainableAppCallback, mockRMClient, clock); scheduler.initialize();/* ww w.j a va 2 s . com*/ drainableAppCallback.drain(); scheduler.start(); drainableAppCallback.drain(); verify(mockRMClient).start(); verify(mockRMClient).registerApplicationMaster(appHost, appPort, appUrl); RegisterApplicationMasterResponse regResponse = mockRMClient.getRegistrationResponse(); verify(mockApp).setApplicationRegistrationData(regResponse.getMaximumResourceCapability(), regResponse.getApplicationACLs(), regResponse.getClientToAMTokenMasterKey(), regResponse.getQueue()); assertEquals(scheduler.getClusterNodeCount(), mockRMClient.getClusterNodeCount()); final String rack1 = "/r1"; final String rack2 = "/r2"; final String node1Rack1 = "n1r1"; final String node2Rack1 = "n2r1"; final String node1Rack2 = "n1r2"; MockDNSToSwitchMapping.addRackMapping(node1Rack1, rack1); MockDNSToSwitchMapping.addRackMapping(node2Rack1, rack1); MockDNSToSwitchMapping.addRackMapping(node1Rack2, rack2); Priority priorityv0 = Priority.newInstance(1); MockTaskInfo taskv0t0 = new MockTaskInfo("taskv0t0", priorityv0, node1Rack1, rack1); ApplicationAttemptId attemptId = ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 1), 1); ContainerId cid1 = ContainerId.newContainerId(attemptId, 1); NodeId n2r1 = NodeId.newInstance(node2Rack1, 1); Container container1 = Container.newInstance(cid1, n2r1, null, taskv0t0.capability, priorityv0, null); // verify idle container is kept for now scheduler.onContainersAllocated(Collections.singletonList(container1)); clock.incrementTime(2000); drainableAppCallback.drain(); verify(mockApp, never()).containerBeingReleased(cid1); verify(mockRMClient, never()).releaseAssignedContainer(cid1); // verify idle container is released without being assigned to a task because rack-local reuse is // disabled TaskRequestCaptor taskRequestCaptor = new TaskRequestCaptor(mockRMClient, scheduler, drainableAppCallback); TaskRequest reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0); clock.incrementTime(10000); drainableAppCallback.drain(); verify(mockApp, never()).taskAllocated(taskv0t0.task, taskv0t0.cookie, container1); verify(mockRMClient, never()).removeContainerRequest(reqv0t0); verify(mockApp, never()).containerBeingReleased(cid1); verify(mockRMClient).releaseAssignedContainer(cid1); // cancel the task request assertFalse(scheduler.deallocateTask(taskv0t0.task, false, null, null)); // allocate another container that's node-local ContainerId cid2 = ContainerId.newContainerId(attemptId, 2); NodeId n1r1 = NodeId.newInstance(node1Rack1, 1); Container container2 = Container.newInstance(cid2, n1r1, null, taskv0t0.capability, priorityv0, null); scheduler.onContainersAllocated(Collections.singletonList(container2)); clock.incrementTime(2000); drainableAppCallback.drain(); verify(mockApp, never()).containerBeingReleased(cid2); verify(mockRMClient, never()).releaseAssignedContainer(cid2); // reschedule the task, verify it's now scheduled without a container request // since node-local idle container is available reqv0t0 = taskRequestCaptor.scheduleTask(taskv0t0, false); verify(mockApp).taskAllocated(taskv0t0.task, taskv0t0.cookie, container2); verify(mockRMClient).removeContainerRequest(reqv0t0); }
From source file:org.apache.tez.dag.app.rm.TestLlapTaskSchedulerService.java
License:Apache License
@Test(timeout = 5000) public void testSimpleLocalAllocation() throws IOException, InterruptedException { TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(); try {/*from w w w . j a v a2 s . c o m*/ Priority priority1 = Priority.newInstance(1); String[] hosts1 = new String[] { HOST1 }; Object task1 = new Object(); Object clientCookie1 = new Object(); tsWrapper.controlScheduler(true); tsWrapper.allocateTask(task1, hosts1, priority1, clientCookie1); tsWrapper.signalSchedulerRun(); tsWrapper.awaitSchedulerRun(); verify(tsWrapper.mockAppCallback).taskAllocated(eq(task1), eq(clientCookie1), any(Container.class)); // TODO Verify this is on host1. assertEquals(1, tsWrapper.ts.dagStats.numLocalAllocations); } finally { tsWrapper.shutdown(); } }
From source file:org.apache.tez.dag.app.rm.TestLlapTaskSchedulerService.java
License:Apache License
@Test(timeout = 5000) public void testSimpleNoLocalityAllocation() throws IOException, InterruptedException { TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(); try {/*from www . j av a 2 s . c o m*/ Priority priority1 = Priority.newInstance(1); Object task1 = new Object(); Object clientCookie1 = new Object(); tsWrapper.controlScheduler(true); tsWrapper.allocateTask(task1, null, priority1, clientCookie1); tsWrapper.signalSchedulerRun(); tsWrapper.awaitSchedulerRun(); verify(tsWrapper.mockAppCallback).taskAllocated(eq(task1), eq(clientCookie1), any(Container.class)); assertEquals(1, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest); } finally { tsWrapper.shutdown(); } }
From source file:org.apache.tez.dag.app.rm.TestLlapTaskSchedulerService.java
License:Apache License
@Test(timeout = 5000) public void testPreemption() throws InterruptedException, IOException { Priority priority1 = Priority.newInstance(1); Priority priority2 = Priority.newInstance(2); String[] hosts = new String[] { HOST1 }; TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(2000, hosts, 1, 1); try {/*from ww w. j a va 2 s .c o m*/ Object task1 = "task1"; Object clientCookie1 = "cookie1"; Object task2 = "task2"; Object clientCookie2 = "cookie2"; Object task3 = "task3"; Object clientCookie3 = "cookie3"; Object task4 = "task4"; Object clientCookie4 = "cookie4"; tsWrapper.controlScheduler(true); tsWrapper.allocateTask(task1, hosts, priority2, clientCookie1); tsWrapper.allocateTask(task2, hosts, priority2, clientCookie2); tsWrapper.allocateTask(task3, hosts, priority2, clientCookie3); while (true) { tsWrapper.signalSchedulerRun(); tsWrapper.awaitSchedulerRun(); if (tsWrapper.ts.dagStats.numLocalAllocations == 2) { break; } } verify(tsWrapper.mockAppCallback, times(2)).taskAllocated(any(Object.class), any(Object.class), any(Container.class)); assertEquals(2, tsWrapper.ts.dagStats.numLocalAllocations); assertEquals(0, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest); reset(tsWrapper.mockAppCallback); tsWrapper.allocateTask(task4, hosts, priority1, clientCookie4); while (true) { tsWrapper.signalSchedulerRun(); tsWrapper.awaitSchedulerRun(); if (tsWrapper.ts.dagStats.numPreemptedTasks == 1) { break; } } verify(tsWrapper.mockAppCallback).preemptContainer(any(ContainerId.class)); tsWrapper.deallocateTask(task2, false, TaskAttemptEndReason.INTERNAL_PREEMPTION); while (true) { tsWrapper.signalSchedulerRun(); tsWrapper.awaitSchedulerRun(); if (tsWrapper.ts.dagStats.numTotalAllocations == 3) { break; } } verify(tsWrapper.mockAppCallback, times(1)).taskAllocated(eq(task4), eq(clientCookie4), any(Container.class)); } finally { tsWrapper.shutdown(); } }
From source file:org.apache.tez.dag.app.rm.TestLlapTaskSchedulerService.java
License:Apache License
@Test(timeout = 5000) public void testNodeDisabled() throws IOException, InterruptedException { TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(10000l); try {//from w w w . j av a 2 s . com Priority priority1 = Priority.newInstance(1); String[] hosts1 = new String[] { HOST1 }; Object task1 = new Object(); Object clientCookie1 = new Object(); tsWrapper.controlScheduler(true); tsWrapper.allocateTask(task1, hosts1, priority1, clientCookie1); while (true) { tsWrapper.signalSchedulerRun(); tsWrapper.awaitSchedulerRun(); if (tsWrapper.ts.dagStats.numTotalAllocations == 1) { break; } } verify(tsWrapper.mockAppCallback).taskAllocated(eq(task1), eq(clientCookie1), any(Container.class)); assertEquals(1, tsWrapper.ts.dagStats.numLocalAllocations); assertEquals(0, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest); assertEquals(0, tsWrapper.ts.dagStats.numNonLocalAllocations); assertEquals(1, tsWrapper.ts.dagStats.numTotalAllocations); tsWrapper.resetAppCallback(); tsWrapper.clock.setTime(10000l); tsWrapper.rejectExecution(task1); // Verify that the node is blacklisted assertEquals(1, tsWrapper.ts.dagStats.numRejectedTasks); assertEquals(3, tsWrapper.ts.instanceToNodeMap.size()); LlapTaskSchedulerService.NodeInfo disabledNodeInfo = tsWrapper.ts.disabledNodesQueue.peek(); assertNotNull(disabledNodeInfo); assertEquals(HOST1, disabledNodeInfo.serviceInstance.getHost()); assertEquals((10000l), disabledNodeInfo.getDelay(TimeUnit.MILLISECONDS)); assertEquals((10000l + 10000l), disabledNodeInfo.expireTimeMillis); Object task2 = new Object(); Object clientCookie2 = new Object(); tsWrapper.allocateTask(task2, hosts1, priority1, clientCookie2); while (true) { tsWrapper.signalSchedulerRun(); tsWrapper.awaitSchedulerRun(); if (tsWrapper.ts.dagStats.numTotalAllocations == 2) { break; } } verify(tsWrapper.mockAppCallback).taskAllocated(eq(task2), eq(clientCookie2), any(Container.class)); assertEquals(1, tsWrapper.ts.dagStats.numLocalAllocations); assertEquals(0, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest); assertEquals(1, tsWrapper.ts.dagStats.numNonLocalAllocations); assertEquals(2, tsWrapper.ts.dagStats.numTotalAllocations); } finally { tsWrapper.shutdown(); } }
From source file:org.apache.tez.dag.app.rm.TestLlapTaskSchedulerService.java
License:Apache License
@Test(timeout = 5000) public void testNodeReEnabled() throws InterruptedException, IOException { // Based on actual timing. TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(1000l); try {/* w w w . ja v a 2 s. c o m*/ Priority priority1 = Priority.newInstance(1); String[] hosts1 = new String[] { HOST1 }; String[] hosts2 = new String[] { HOST2 }; String[] hosts3 = new String[] { HOST3 }; Object task1 = new Object(); Object clientCookie1 = new Object(); Object task2 = new Object(); Object clientCookie2 = new Object(); Object task3 = new Object(); Object clientCookie3 = new Object(); tsWrapper.controlScheduler(true); tsWrapper.allocateTask(task1, hosts1, priority1, clientCookie1); tsWrapper.allocateTask(task2, hosts2, priority1, clientCookie2); tsWrapper.allocateTask(task3, hosts3, priority1, clientCookie3); while (true) { tsWrapper.signalSchedulerRun(); tsWrapper.awaitSchedulerRun(); if (tsWrapper.ts.dagStats.numTotalAllocations == 3) { break; } } verify(tsWrapper.mockAppCallback, times(3)).taskAllocated(any(Object.class), any(Object.class), any(Container.class)); assertEquals(3, tsWrapper.ts.dagStats.numLocalAllocations); assertEquals(0, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest); assertEquals(3, tsWrapper.ts.dagStats.numTotalAllocations); tsWrapper.resetAppCallback(); tsWrapper.rejectExecution(task1); tsWrapper.rejectExecution(task2); tsWrapper.rejectExecution(task3); // Verify that the node is blacklisted assertEquals(3, tsWrapper.ts.dagStats.numRejectedTasks); assertEquals(3, tsWrapper.ts.instanceToNodeMap.size()); assertEquals(3, tsWrapper.ts.disabledNodesQueue.size()); Object task4 = new Object(); Object clientCookie4 = new Object(); Object task5 = new Object(); Object clientCookie5 = new Object(); Object task6 = new Object(); Object clientCookie6 = new Object(); tsWrapper.allocateTask(task4, hosts1, priority1, clientCookie4); tsWrapper.allocateTask(task5, hosts2, priority1, clientCookie5); tsWrapper.allocateTask(task6, hosts3, priority1, clientCookie6); while (true) { tsWrapper.signalSchedulerRun(); tsWrapper.awaitSchedulerRun(); if (tsWrapper.ts.dagStats.numTotalAllocations == 6) { break; } } ArgumentCaptor<Container> argumentCaptor = ArgumentCaptor.forClass(Container.class); verify(tsWrapper.mockAppCallback, times(3)).taskAllocated(any(Object.class), any(Object.class), argumentCaptor.capture()); // which affects the locality matching assertEquals(0, tsWrapper.ts.dagStats.numAllocationsNoLocalityRequest); assertEquals(6, tsWrapper.ts.dagStats.numTotalAllocations); } finally { tsWrapper.shutdown(); } }
From source file:org.apache.tez.dag.app.rm.TestLlapTaskSchedulerService.java
License:Apache License
private void forceLocalityTest1(boolean forceLocality) throws IOException, InterruptedException { Priority priority1 = Priority.newInstance(1); String[] hosts = new String[] { HOST1, HOST2 }; String[] hostsH1 = new String[] { HOST1 }; String[] hostsH2 = new String[] { HOST2 }; TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(2000, hosts, 1, 1, (forceLocality ? -1l : 0l)); try {//from w w w . j a va 2 s. c om Object task1 = "task1"; Object clientCookie1 = "cookie1"; Object task2 = "task2"; Object clientCookie2 = "cookie2"; Object task3 = "task3"; Object clientCookie3 = "cookie3"; Object task4 = "task4"; Object clientCookie4 = "cookie4"; Object task5 = "task5"; Object clientCookie5 = "cookie5"; tsWrapper.controlScheduler(true); //H1 - should allocate tsWrapper.allocateTask(task1, hostsH1, priority1, clientCookie1); //H1 - should allocate tsWrapper.allocateTask(task2, hostsH1, priority1, clientCookie2); //H1 - no capacity if force, should allocate otherwise tsWrapper.allocateTask(task3, hostsH1, priority1, clientCookie3); //H2 - should allocate tsWrapper.allocateTask(task4, hostsH2, priority1, clientCookie4); //No location - should allocate if force, no capacity otherwise tsWrapper.allocateTask(task5, null, priority1, clientCookie5); while (true) { tsWrapper.signalSchedulerRun(); tsWrapper.awaitSchedulerRun(); if (tsWrapper.ts.dagStats.numTotalAllocations == 4) { break; } } // Verify no preemption requests - since everything is at the same priority verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class)); ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class); verify(tsWrapper.mockAppCallback, times(4)).taskAllocated(argumentCaptor.capture(), any(Object.class), any(Container.class)); assertEquals(4, argumentCaptor.getAllValues().size()); assertEquals(task1, argumentCaptor.getAllValues().get(0)); assertEquals(task2, argumentCaptor.getAllValues().get(1)); if (forceLocality) { // task3 not allocated assertEquals(task4, argumentCaptor.getAllValues().get(2)); assertEquals(task5, argumentCaptor.getAllValues().get(3)); } else { assertEquals(task3, argumentCaptor.getAllValues().get(2)); assertEquals(task4, argumentCaptor.getAllValues().get(3)); } //Complete one task on host1. tsWrapper.deallocateTask(task1, true, null); reset(tsWrapper.mockAppCallback); // Try scheduling again. while (true) { tsWrapper.signalSchedulerRun(); tsWrapper.awaitSchedulerRun(); if (tsWrapper.ts.dagStats.numTotalAllocations == 5) { break; } } verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class)); argumentCaptor = ArgumentCaptor.forClass(Object.class); verify(tsWrapper.mockAppCallback, times(1)).taskAllocated(argumentCaptor.capture(), any(Object.class), any(Container.class)); assertEquals(1, argumentCaptor.getAllValues().size()); if (forceLocality) { assertEquals(task3, argumentCaptor.getAllValues().get(0)); } else { assertEquals(task5, argumentCaptor.getAllValues().get(0)); } } finally { tsWrapper.shutdown(); } }
From source file:org.apache.tez.dag.app.rm.TestLlapTaskSchedulerService.java
License:Apache License
@Test(timeout = 5000) public void testForcedLocalityUnknownHost() throws IOException, InterruptedException { Priority priority1 = Priority.newInstance(1); String[] hostsKnown = new String[] { HOST1 }; String[] hostsUnknown = new String[] { HOST2 }; TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(2000, hostsKnown, 1, 1, -1l);/*from www .j a v a 2s .c om*/ try { Object task1 = "task1"; Object clientCookie1 = "cookie1"; Object task2 = "task2"; Object clientCookie2 = "cookie2"; tsWrapper.controlScheduler(true); // Should allocate since H2 is not known. tsWrapper.allocateTask(task1, hostsUnknown, priority1, clientCookie1); tsWrapper.allocateTask(task2, hostsKnown, priority1, clientCookie2); while (true) { tsWrapper.signalSchedulerRun(); tsWrapper.awaitSchedulerRun(); if (tsWrapper.ts.dagStats.numTotalAllocations == 2) { break; } } ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class); verify(tsWrapper.mockAppCallback, times(2)).taskAllocated(argumentCaptor.capture(), any(Object.class), any(Container.class)); assertEquals(2, argumentCaptor.getAllValues().size()); assertEquals(task1, argumentCaptor.getAllValues().get(0)); assertEquals(task2, argumentCaptor.getAllValues().get(1)); } finally { tsWrapper.shutdown(); } }
From source file:org.apache.tez.dag.app.rm.TestLlapTaskSchedulerService.java
License:Apache License
@Test(timeout = 5000) public void testForcedLocalityPreemption() throws IOException, InterruptedException { Priority priority1 = Priority.newInstance(1); Priority priority2 = Priority.newInstance(2); String[] hosts = new String[] { HOST1, HOST2 }; String[] hostsH1 = new String[] { HOST1 }; String[] hostsH2 = new String[] { HOST2 }; TestTaskSchedulerServiceWrapper tsWrapper = new TestTaskSchedulerServiceWrapper(2000, hosts, 1, 1, -1l); // Fill up host1 with p2 tasks. // Leave host2 empty // Try running p1 task on host1 - should preempt try {/*from ww w .j av a 2 s . c om*/ Object task1 = "task1"; Object clientCookie1 = "cookie1"; Object task2 = "task2"; Object clientCookie2 = "cookie2"; Object task3 = "task3"; Object clientCookie3 = "cookie3"; Object task4 = "task4"; Object clientCookie4 = "cookie4"; tsWrapper.controlScheduler(true); tsWrapper.allocateTask(task1, hostsH1, priority2, clientCookie1); tsWrapper.allocateTask(task2, hostsH1, priority2, clientCookie2); // This request at a lower priority should not affect anything. tsWrapper.allocateTask(task3, hostsH1, priority2, clientCookie3); while (true) { tsWrapper.signalSchedulerRun(); tsWrapper.awaitSchedulerRun(); if (tsWrapper.ts.dagStats.numLocalAllocations == 2) { break; } } verify(tsWrapper.mockAppCallback, never()).preemptContainer(any(ContainerId.class)); ArgumentCaptor<Object> argumentCaptor = ArgumentCaptor.forClass(Object.class); verify(tsWrapper.mockAppCallback, times(2)).taskAllocated(argumentCaptor.capture(), any(Object.class), any(Container.class)); assertEquals(2, argumentCaptor.getAllValues().size()); assertEquals(task1, argumentCaptor.getAllValues().get(0)); assertEquals(task2, argumentCaptor.getAllValues().get(1)); reset(tsWrapper.mockAppCallback); // Allocate t4 at higher priority. t3 should not be allocated, // and a preemption should be attempted on host1, despite host2 having available capacity tsWrapper.allocateTask(task4, hostsH1, priority1, clientCookie4); while (true) { tsWrapper.signalSchedulerRun(); tsWrapper.awaitSchedulerRun(); if (tsWrapper.ts.dagStats.numPreemptedTasks == 1) { break; } } verify(tsWrapper.mockAppCallback).preemptContainer(any(ContainerId.class)); tsWrapper.deallocateTask(task1, false, TaskAttemptEndReason.INTERNAL_PREEMPTION); while (true) { tsWrapper.signalSchedulerRun(); tsWrapper.awaitSchedulerRun(); if (tsWrapper.ts.dagStats.numTotalAllocations == 3) { break; } } verify(tsWrapper.mockAppCallback, times(1)).taskAllocated(eq(task4), eq(clientCookie4), any(Container.class)); } finally { tsWrapper.shutdown(); } }
From source file:org.apache.tez.dag.app.rm.TestLocalTaskScheduler.java
License:Apache License
@Test(timeout = 5000) public void maxTasksAllocationsCannotBeExceeded() { final int MAX_TASKS = 4; TezConfiguration tezConf = new TezConfiguration(); tezConf.setInt(TezConfiguration.TEZ_AM_INLINE_TASK_EXECUTION_MAX_TASKS, MAX_TASKS); LocalContainerFactory containerFactory = new LocalContainerFactory(createMockAppContext()); HashMap<Object, Container> taskAllocations = new LinkedHashMap<Object, Container>(); PriorityBlockingQueue<TaskRequest> taskRequestQueue = new PriorityBlockingQueue<TaskRequest>(); TaskSchedulerAppCallback appClientDelegate = mock(TaskSchedulerAppCallback.class); // Object under test AsyncDelegateRequestHandler requestHandler = new AsyncDelegateRequestHandler(taskRequestQueue, containerFactory, taskAllocations, appClientDelegate, tezConf); // Allocate up to max tasks for (int i = 0; i < MAX_TASKS; i++) { Priority priority = Priority.newInstance(20); requestHandler.addAllocateTaskRequest(new Long(i), null, priority, null); requestHandler.processRequest(); }//from w ww. jav a 2 s .com // Only MAX_TASKS number of tasks should have been allocated Assert.assertEquals("Wrong number of allocate tasks", MAX_TASKS, taskAllocations.size()); Assert.assertTrue("Another allocation should not fit", requestHandler.shouldWait()); // Deallocate down to zero for (int i = 0; i < MAX_TASKS; i++) { requestHandler.addDeallocateTaskRequest(new Long(i)); requestHandler.processRequest(); } // All allocated tasks should have been removed Assert.assertEquals("Wrong number of allocate tasks", 0, taskAllocations.size()); }