TieredQueue implements InternalTaskQueue { @Override public void setSla(TaskQueueSla sla) throws IllegalArgumentException { if (sla != null && !(sla instanceof TieredQueueSlas)) { throw new IllegalArgumentException("Queue SLA must be an instance of " + TieredQueueSlas.class.getName() + ", can't accept " + sla.getClass().getName()); } slasQueue.offer(sla == null? new TieredQueueSlas(Collections.emptyMap(), Collections.emptyMap()) : (TieredQueueSlas)sla); } TieredQueue(int numTiers); int getNumTiers(); @Override void queueTask(QueuableTask task); @Override void setSla(TaskQueueSla sla); @Override Assignable<QueuableTask> next(); @Override boolean reset(); @Override UsageTrackedQueue getUsageTracker(); @Override Map<TaskState, Collection<QueuableTask>> getAllTasks();  }
@Test public void testTierSlas() throws Exception { TaskQueue queue = new TieredQueue(2); Map<String, Double> bucketWeights = new HashMap<>(); bucketWeights.put("A", 1.0); bucketWeights.put("B", 2.0); bucketWeights.put("C", 1.0); queue.setSla(new TieredQueueSlas(Collections.emptyMap(), getTierAllocsForBuckets(bucketWeights))); final AtomicInteger countA = new AtomicInteger(); final AtomicInteger countB = new AtomicInteger(); final AtomicInteger countC = new AtomicInteger(); final CountDownLatch latch = new CountDownLatch(40); final TaskSchedulingService schedulingService = setupSchedSvcWithAssignmentCounters(queue, new AtomicReference<>(latch), countA, countB, countC); createTasksForBuckets(queue, 25, "A", "B", "C"); schedulingService.addLeases(LeaseProvider.getLeases(10, 4.0, 4000.0, 4000.0, 1, 100)); schedulingService.start(); Assert.assertTrue("Timeout waiting for assignments", latch.await(2, TimeUnit.SECONDS)); Thread.sleep(500); Assert.assertEquals(10, countA.get()); Assert.assertEquals(20, countB.get()); Assert.assertEquals(10, countC.get()); }
@Test public void testTierSlas2() throws Exception { TaskQueue queue = new TieredQueue(2); Map<String, Double> bucketWeights = new HashMap<>(); bucketWeights.put("A", 1.0); bucketWeights.put("B", 2.0); bucketWeights.put("C", 1.0); ResAllocs tier1Capacity = new ResAllocsBuilder("tier#0") .withCores(10 * 4) .withMemory(10 * 4000) .withNetworkMbps(10 * 4000) .build(); queue.setSla(new TieredQueueSlas(singletonMap(1, tier1Capacity), getTierAllocsForBuckets(bucketWeights))); final AtomicInteger countA = new AtomicInteger(); final AtomicInteger countB = new AtomicInteger(); final AtomicInteger countC = new AtomicInteger(); final CountDownLatch latch = new CountDownLatch(40); final AtomicReference<CountDownLatch> latchRef = new AtomicReference<>(latch); final TaskSchedulingService schedulingService = setupSchedSvcWithAssignmentCounters(queue, latchRef, countA, countB, countC); createTasksForBuckets(queue, 25, "A", "B", "C"); schedulingService.addLeases(LeaseProvider.getLeases(10, 4.0, 4000.0, 4000.0, 1, 100)); schedulingService.start(); Assert.assertTrue("Timeout waiting for assignments", latch.await(2, TimeUnit.SECONDS)); Thread.sleep(500); Assert.assertEquals(10, countA.get()); Assert.assertEquals(20, countB.get()); Assert.assertEquals(10, countC.get()); bucketWeights.clear(); bucketWeights.put("A", 1.0); bucketWeights.put("B", 1.0); bucketWeights.put("C", 2.0); ResAllocs doubledTier1Capacity = new ResAllocsBuilder("tier#0") .withCores(20 * 4) .withMemory(20 * 4000) .withNetworkMbps(20 * 4000) .build(); queue.setSla(new TieredQueueSlas(singletonMap(0, doubledTier1Capacity), getTierAllocsForBuckets(bucketWeights))); latchRef.set(new CountDownLatch(40)); countA.set(0); countB.set(0); countC.set(0); createTasksForBuckets(queue, 25, "A", "B", "C"); schedulingService.addLeases(LeaseProvider.getLeases(10, 4.0, 4000.0, 4000.0, 1, 100)); Assert.assertTrue("Timeout waiting for assignments", latchRef.get().await(2, TimeUnit.SECONDS)); Thread.sleep(500); Assert.assertEquals(10, countA.get()); Assert.assertEquals(0, countB.get()); Assert.assertEquals(30, countC.get()); }
@Test public void testTierSlasWithPreviouslyRunningTasks() throws Exception { TaskQueue queue = new TieredQueue(2); Map<String, Double> bucketWeights = new HashMap<>(); bucketWeights.put("A", 1.0); bucketWeights.put("B", 2.0); bucketWeights.put("C", 1.0); queue.setSla(new TieredQueueSlas(Collections.emptyMap(), getTierAllocsForBuckets(bucketWeights))); final AtomicInteger countA = new AtomicInteger(); final AtomicInteger countB = new AtomicInteger(); final AtomicInteger countC = new AtomicInteger(); final CountDownLatch latch = new CountDownLatch(32); final AtomicReference<CountDownLatch> latchRef = new AtomicReference<>(latch); final TaskSchedulingService schedulingService = setupSchedSvcWithAssignmentCounters(queue, latchRef, countA, countB, countC); initializeRunningTasksForBuckets(schedulingService, 5, "A"); initializeRunningTasksForBuckets(schedulingService, 3, "C"); createTasksForBuckets(queue, 25, "A", "B", "C"); schedulingService.addLeases(LeaseProvider.getLeases(8, 4.0, 4000.0, 4000.0, 1, 100)); schedulingService.start(); Assert.assertTrue("Timeout waiting for assignments", latchRef.get().await(2, TimeUnit.SECONDS)); Thread.sleep(500); Assert.assertEquals(10-5, countA.get()); Assert.assertEquals(20, countB.get()); Assert.assertEquals(10-3, countC.get()); }
@Test public void testTierSlasWithOneBucketWithNoTasks() throws Exception { TaskQueue queue = new TieredQueue(2); Map<String, Double> bucketWeights = new HashMap<>(); bucketWeights.put("A", 1.0); bucketWeights.put("B", 2.0); bucketWeights.put("C", 1.0); ResAllocs tier1Capacity = new ResAllocsBuilder("tier#0") .withCores(70 * 4) .withMemory(70 * 4000) .withNetworkMbps(70 * 4000) .build(); queue.setSla(new TieredQueueSlas(singletonMap(0, tier1Capacity), getTierAllocsForBuckets(bucketWeights))); final AtomicInteger countA = new AtomicInteger(); final AtomicInteger countB = new AtomicInteger(); final AtomicInteger countC = new AtomicInteger(); final CountDownLatch latch = new CountDownLatch(40); final AtomicReference<CountDownLatch> latchRef = new AtomicReference<>(latch); final TaskSchedulingService schedulingService = setupSchedSvcWithAssignmentCounters(queue, latchRef, countA, countB, countC); createTasksForBuckets(queue, 25, "A", "C"); schedulingService.addLeases(LeaseProvider.getLeases(10, 4.0, 4000.0, 4000.0, 1, 100)); schedulingService.start(); Assert.assertTrue("Timeout waiting for assignments", latch.await(2, TimeUnit.SECONDS)); Thread.sleep(500); Assert.assertEquals(20, countA.get()); Assert.assertEquals(0, countB.get()); Assert.assertEquals(20, countC.get()); }