Java tutorial
/* * Copyright 2013 The Regents of The University California * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package edu.berkeley.sparrow.daemon.nodemonitor; import static org.junit.Assert.assertEquals; import java.net.InetSocketAddress; import java.util.List; import org.apache.commons.configuration.PropertiesConfiguration; import org.apache.log4j.BasicConfigurator; import org.junit.Before; import org.junit.Test; import com.google.common.collect.Lists; import edu.berkeley.sparrow.daemon.nodemonitor.TaskScheduler.TaskSpec; import edu.berkeley.sparrow.thrift.TEnqueueTaskReservationsRequest; import edu.berkeley.sparrow.thrift.TFullTaskId; import edu.berkeley.sparrow.thrift.THostPort; import edu.berkeley.sparrow.thrift.TUserGroupInfo; public class TestTaskScheduler { int requestId; @Before public void setUp() { // Set up a simple configuration that logs on the console. BasicConfigurator.configure(); requestId = 1; } private TEnqueueTaskReservationsRequest createTaskReservationRequest(int numTasks, TaskScheduler scheduler, String userId) { return createTaskReservationRequest(numTasks, scheduler, userId, 0); } private TEnqueueTaskReservationsRequest createTaskReservationRequest(int numTasks, TaskScheduler scheduler, String userId, int priority) { String idStr = Integer.toString(requestId++); TUserGroupInfo user = new TUserGroupInfo(userId, "group", priority); THostPort schedulerAddress = new THostPort("1.2.3.4", 52); return new TEnqueueTaskReservationsRequest("appId", user, idStr, schedulerAddress, numTasks); } /** * Tests the fifo task scheduler. */ @Test public void testFifo() { TaskScheduler scheduler = new FifoTaskScheduler(4); scheduler.initialize(new PropertiesConfiguration(), 12345); final String testApp = "test app"; final InetSocketAddress backendAddress = new InetSocketAddress("123.4.5.6", 2); // Make sure that tasks are launched right away, if resources are available. scheduler.submitTaskReservations(createTaskReservationRequest(1, scheduler, testApp), backendAddress); assertEquals(1, scheduler.runnableTasks()); TaskSpec task = scheduler.getNextTask(); assertEquals("1", task.requestId); assertEquals(0, scheduler.runnableTasks()); scheduler.submitTaskReservations(createTaskReservationRequest(2, scheduler, testApp), backendAddress); assertEquals(2, scheduler.runnableTasks()); // Make sure the request to schedule 3 tasks is appropriately split, with one task running // now and others started later. scheduler.submitTaskReservations(createTaskReservationRequest(3, scheduler, testApp), backendAddress); /* 4 tasks have been launched but one was already removed from the runnable queue using * getTask(), leaving 3 runnable tasks. */ assertEquals(3, scheduler.runnableTasks()); task = scheduler.getNextTask(); assertEquals("2", task.requestId); task = scheduler.getNextTask(); assertEquals("2", task.requestId); /* Make a list of task ids to use in every call to tasksFinished, and just update the request * id for each call. */ TFullTaskId fullTaskId = new TFullTaskId(); fullTaskId.taskId = ""; List<TFullTaskId> completedTasks = Lists.newArrayList(); completedTasks.add(fullTaskId); // Have a few tasks complete before the last runnable task is removed from the queue. fullTaskId.requestId = "2"; scheduler.tasksFinished(completedTasks); scheduler.tasksFinished(completedTasks); fullTaskId.requestId = "1"; scheduler.tasksFinished(completedTasks); task = scheduler.getNextTask(); assertEquals("3", task.requestId); task = scheduler.getNextTask(); assertEquals("3", task.requestId); task = scheduler.getNextTask(); assertEquals("3", task.requestId); assertEquals(0, scheduler.runnableTasks()); } /** * Tests the round robin task scheduler. */ @Test public void testBasicRoundRobin() { TaskScheduler scheduler = new RoundRobinTaskScheduler(4); scheduler.initialize(new PropertiesConfiguration(), 12345); final String user1 = "user1"; final InetSocketAddress address1 = new InetSocketAddress("localhost", 1); final String user2 = "user2"; final InetSocketAddress address2 = new InetSocketAddress("localhost", 1); final String user3 = "user3"; final InetSocketAddress address3 = new InetSocketAddress("localhost", 1); final String user4 = "user4"; final InetSocketAddress address4 = new InetSocketAddress("localhost", 1); // Submit enough tasks to saturate the existing capacity. scheduler.submitTaskReservations(createTaskReservationRequest(1, scheduler, user1), address1); assertEquals(1, scheduler.runnableTasks()); scheduler.getNextTask(); assertEquals(0, scheduler.runnableTasks()); scheduler.submitTaskReservations(createTaskReservationRequest(1, scheduler, user2), address2); assertEquals(1, scheduler.runnableTasks()); scheduler.getNextTask(); assertEquals(0, scheduler.runnableTasks()); scheduler.submitTaskReservations(createTaskReservationRequest(1, scheduler, user3), address3); assertEquals(1, scheduler.runnableTasks()); scheduler.getNextTask(); assertEquals(0, scheduler.runnableTasks()); scheduler.submitTaskReservations(createTaskReservationRequest(1, scheduler, user4), address4); assertEquals(1, scheduler.runnableTasks()); scheduler.getNextTask(); assertEquals(0, scheduler.runnableTasks()); /* Create the following backlogs. * user1: 2 tasks * user2: 3 tasks * user3: 4 tasks */ scheduler.submitTaskReservations(createTaskReservationRequest(2, scheduler, user1), address1); scheduler.submitTaskReservations(createTaskReservationRequest(1, scheduler, user2), address2); scheduler.submitTaskReservations(createTaskReservationRequest(1, scheduler, user2), address2); scheduler.submitTaskReservations(createTaskReservationRequest(1, scheduler, user2), address2); scheduler.submitTaskReservations(createTaskReservationRequest(4, scheduler, user3), address3); assertEquals(0, scheduler.runnableTasks()); /* Make sure that as tasks finish (and space is freed up) new tasks are added to the runqueue * in round-robin order. * Make a list of task ids to use in every call to tasksFinished, and just update the request * id. */ TFullTaskId fullTaskId = new TFullTaskId(); fullTaskId.taskId = ""; List<TFullTaskId> completedTasks = Lists.newArrayList(); completedTasks.add(fullTaskId); fullTaskId.requestId = "1"; scheduler.tasksFinished(completedTasks); assertEquals(1, scheduler.runnableTasks()); TaskSpec task = scheduler.getNextTask(); assertEquals("5", task.requestId); assertEquals(0, scheduler.runnableTasks()); fullTaskId.requestId = task.requestId; scheduler.tasksFinished(completedTasks); assertEquals(1, scheduler.runnableTasks()); task = scheduler.getNextTask(); assertEquals("6", task.requestId); assertEquals(0, scheduler.runnableTasks()); fullTaskId.requestId = task.requestId; scheduler.tasksFinished(completedTasks); assertEquals(1, scheduler.runnableTasks()); task = scheduler.getNextTask(); assertEquals("9", task.requestId); assertEquals(0, scheduler.runnableTasks()); fullTaskId.requestId = task.requestId; scheduler.tasksFinished(completedTasks); assertEquals(1, scheduler.runnableTasks()); task = scheduler.getNextTask(); assertEquals("5", task.requestId); assertEquals(0, scheduler.runnableTasks()); fullTaskId.requestId = task.requestId; scheduler.tasksFinished(completedTasks); assertEquals(1, scheduler.runnableTasks()); task = scheduler.getNextTask(); assertEquals("7", task.requestId); assertEquals(0, scheduler.runnableTasks()); fullTaskId.requestId = task.requestId; scheduler.tasksFinished(completedTasks); assertEquals(1, scheduler.runnableTasks()); task = scheduler.getNextTask(); assertEquals("9", task.requestId); assertEquals(0, scheduler.runnableTasks()); fullTaskId.requestId = task.requestId; scheduler.tasksFinished(completedTasks); assertEquals(1, scheduler.runnableTasks()); task = scheduler.getNextTask(); assertEquals("8", task.requestId); assertEquals(0, scheduler.runnableTasks()); fullTaskId.requestId = task.requestId; scheduler.tasksFinished(completedTasks); assertEquals(1, scheduler.runnableTasks()); task = scheduler.getNextTask(); assertEquals("9", task.requestId); assertEquals(0, scheduler.runnableTasks()); fullTaskId.requestId = task.requestId; scheduler.tasksFinished(completedTasks); assertEquals(1, scheduler.runnableTasks()); task = scheduler.getNextTask(); assertEquals("9", task.requestId); assertEquals(0, scheduler.runnableTasks()); } @Test public void testRoundRobinDoesNotGiveUpUserSpotWhenGetTaskFails() { TaskScheduler scheduler = new RoundRobinTaskScheduler(4); scheduler.initialize(new PropertiesConfiguration(), 12345); final String user1 = "user1"; final InetSocketAddress address1 = new InetSocketAddress("localhost", 1); final String user2 = "user2"; final InetSocketAddress address2 = new InetSocketAddress("localhost", 1); final String user3 = "user3"; final InetSocketAddress address3 = new InetSocketAddress("localhost", 1); final String user4 = "user4"; final InetSocketAddress address4 = new InetSocketAddress("localhost", 1); // Submit enough tasks to saturate the existing capacity. scheduler.submitTaskReservations(createTaskReservationRequest(1, scheduler, user1), address1); assertEquals(1, scheduler.runnableTasks()); scheduler.getNextTask(); assertEquals(0, scheduler.runnableTasks()); scheduler.submitTaskReservations(createTaskReservationRequest(1, scheduler, user2), address2); assertEquals(1, scheduler.runnableTasks()); TaskSpec user2Task = scheduler.getNextTask(); assertEquals(0, scheduler.runnableTasks()); scheduler.submitTaskReservations(createTaskReservationRequest(1, scheduler, user3), address3); assertEquals(1, scheduler.runnableTasks()); scheduler.getNextTask(); assertEquals(0, scheduler.runnableTasks()); scheduler.submitTaskReservations(createTaskReservationRequest(1, scheduler, user4), address4); assertEquals(1, scheduler.runnableTasks()); scheduler.getNextTask(); assertEquals(0, scheduler.runnableTasks()); /* Create the following backlogs. * user1: 2 tasks * user2: 3 tasks * user3: 4 tasks */ scheduler.submitTaskReservations(createTaskReservationRequest(2, scheduler, user1), address1); scheduler.submitTaskReservations(createTaskReservationRequest(1, scheduler, user2), address2); scheduler.submitTaskReservations(createTaskReservationRequest(1, scheduler, user2), address2); scheduler.submitTaskReservations(createTaskReservationRequest(1, scheduler, user2), address2); scheduler.submitTaskReservations(createTaskReservationRequest(4, scheduler, user3), address3); assertEquals(0, scheduler.runnableTasks()); // If the getTask() for user 2 fails, the scheduler should try to get another task // for that user rather than using the usual round-robin ordering. scheduler.noTaskForReservation(user2Task); assertEquals(1, scheduler.runnableTasks()); user2Task = scheduler.getNextTask(); assertEquals(user2, user2Task.user.user); assertEquals("6", user2Task.requestId); assertEquals(0, scheduler.runnableTasks()); // The scheduler should resume round-robin from user 1 (and not start the round-robin // again after user 2). TFullTaskId fullTaskId = new TFullTaskId(); fullTaskId.taskId = ""; List<TFullTaskId> completedTasks = Lists.newArrayList(); completedTasks.add(fullTaskId); fullTaskId.requestId = "1"; scheduler.tasksFinished(completedTasks); assertEquals(1, scheduler.runnableTasks()); TaskSpec task = scheduler.getNextTask(); assertEquals(user1, task.user.user); assertEquals("5", task.requestId); assertEquals(0, scheduler.runnableTasks()); // If the scheduler eventually exhausts the queue of tasks for user 2, and cannot launch any, // it sound continue in round-robin order. scheduler.noTaskForReservation(user2Task); assertEquals(1, scheduler.runnableTasks()); user2Task = scheduler.getNextTask(); assertEquals(user2, user2Task.user.user); assertEquals("7", user2Task.requestId); assertEquals(0, scheduler.runnableTasks()); scheduler.noTaskForReservation(user2Task); assertEquals(1, scheduler.runnableTasks()); user2Task = scheduler.getNextTask(); assertEquals(user2, user2Task.user.user); assertEquals("8", user2Task.requestId); assertEquals(0, scheduler.runnableTasks()); scheduler.noTaskForReservation(user2Task); assertEquals(1, scheduler.runnableTasks()); task = scheduler.getNextTask(); assertEquals(user3, task.user.user); assertEquals("9", task.requestId); assertEquals(0, scheduler.runnableTasks()); } @Test public void testPriority() { /* Submit tasks at priority 0, 1, and 2 and ensure that the highest priority tasks are * run soonest. */ TaskScheduler scheduler = new PriorityTaskScheduler(4); scheduler.initialize(new PropertiesConfiguration(), 12345); final InetSocketAddress appBackendAddress = new InetSocketAddress("localhost", 1); final String user = "user"; // Submit enough tasks to saturate the existing capacity (with one task queued). scheduler.submitTaskReservations(createTaskReservationRequest(5, scheduler, user, 2), appBackendAddress); assertEquals(4, scheduler.runnableTasks()); // Submit 3 tasks for higher priority users (2 at priority 1 and 1 at priority 0). scheduler.submitTaskReservations(createTaskReservationRequest(2, scheduler, user, 1), appBackendAddress); scheduler.submitTaskReservations(createTaskReservationRequest(1, scheduler, user, 0), appBackendAddress); TaskSpec task = scheduler.getNextTask(); assertEquals("1", task.requestId); assertEquals(3, scheduler.runnableTasks()); task = scheduler.getNextTask(); assertEquals("1", task.requestId); assertEquals(2, scheduler.runnableTasks()); task = scheduler.getNextTask(); assertEquals("1", task.requestId); assertEquals(1, scheduler.runnableTasks()); task = scheduler.getNextTask(); assertEquals("1", task.requestId); assertEquals(0, scheduler.runnableTasks()); /* Make sure that as tasks finish, new tasks are added to the runqueue in strictly priority * order. */ TFullTaskId fullTaskId = new TFullTaskId(); fullTaskId.taskId = ""; List<TFullTaskId> completedTasks = Lists.newArrayList(); completedTasks.add(fullTaskId); fullTaskId.requestId = "1"; scheduler.tasksFinished(completedTasks); assertEquals(1, scheduler.runnableTasks()); task = scheduler.getNextTask(); assertEquals("3", task.requestId); assertEquals(0, scheduler.runnableTasks()); scheduler.tasksFinished(completedTasks); assertEquals(1, scheduler.runnableTasks()); task = scheduler.getNextTask(); assertEquals("2", task.requestId); assertEquals(0, scheduler.runnableTasks()); scheduler.tasksFinished(completedTasks); assertEquals(1, scheduler.runnableTasks()); task = scheduler.getNextTask(); assertEquals("2", task.requestId); assertEquals(0, scheduler.runnableTasks()); scheduler.tasksFinished(completedTasks); assertEquals(1, scheduler.runnableTasks()); task = scheduler.getNextTask(); assertEquals("1", task.requestId); assertEquals(0, scheduler.runnableTasks()); } }