Merge pull request #99 from spodila/chgRmTaskSignature

change signature of TaskSchedulingService.removeTask

Author: spodila

fenzo-core/src/main/java/com/netflix/fenzo/TaskSchedulingService.java

@@ -19,6 +19,7 @@
 import com.netflix.fenzo.functions.Action0;
 import com.netflix.fenzo.functions.Action1;
 import com.netflix.fenzo.queues.*;
+import com.netflix.fenzo.queues.TaskQueue;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -49,11 +50,24 @@
  *         assigned from within this scheduling service. This service assigns the tasks before making the result
  *         available to you via the callback. To mark tasks as running for those tasks that were running from
  *         before this service was created, use {@link #initializeRunningTask(QueuableTask, String)}. Later, call
- *         {@link #removeTask(QueuableTask, String)} when tasks complete or they no longer need resource assignments.
+ *         {@link #removeTask(String, QAttributes, String)} when tasks complete or they no longer need resource assignments.
  *     </LI>
  * </UL>
  */
 public class TaskSchedulingService {
+
+    private static class RemoveTaskRequest {
+        private final String taskId;
+        private final QAttributes qAttributes;
+        private final String hostname;
+
+        public RemoveTaskRequest(String taskId, QAttributes qAttributes, String hostname) {
+            this.taskId = taskId;
+            this.qAttributes = qAttributes;
+            this.hostname = hostname;
+        }
+    }
+
     private static final Logger logger = LoggerFactory.getLogger(TaskSchedulingService.class);
     private final TaskScheduler taskScheduler;
     private final Action1<SchedulingResult> schedulingResultCallback;
@@ -63,7 +77,7 @@ public class TaskSchedulingService {
     private final Action0 preHook;
     private final BlockingQueue<VirtualMachineLease> leaseBlockingQueue = new LinkedBlockingQueue<>();
     private final BlockingQueue<Map<String, QueuableTask>> addRunningTasksQueue = new LinkedBlockingQueue<>();
-    private final BlockingQueue<Map<QueuableTask, String>> removeTasksQueue = new LinkedBlockingQueue<>();
+    private final BlockingQueue<RemoveTaskRequest> removeTasksQueue = new LinkedBlockingQueue<>();
     private final BlockingQueue<Action1<Map<TaskQueue.TaskState, Collection<QueuableTask>>>> taskMapRequest = new LinkedBlockingQueue<>(10);
     private final BlockingQueue<Action1<Map<String, Map<VMResource, Double[]>>>> resStatusRequest = new LinkedBlockingQueue<>(10);
     private final BlockingQueue<Action1<List<VirtualMachineCurrentState>>> vmCurrStateRequest = new LinkedBlockingQueue<>(10);
@@ -161,20 +175,18 @@ private void addPendingRunningTasks() {
 
     private void removeTasks() {
         if (removeTasksQueue.peek() != null) {
-            List<Map<QueuableTask, String>> r = new LinkedList<>();
-            removeTasksQueue.drainTo(r);
-            for (Map<QueuableTask, String> m: r) {
-                for (Map.Entry<QueuableTask, String> e: m.entrySet()) {
-                    // remove it from the queue and call taskScheduler to unassign it if hostname is not null
-                    try {
-                        taskQueue.getUsageTracker().removeTask(e.getKey().getId(), e.getKey().getQAttributes());
-                    } catch (TaskQueueException e1) {
-                        // shouldn't happen since we're calling outside of scheduling iteration
-                        logger.warn("Unexpected to get exception outside of scheduling iteration: " + e1.getMessage(), e1);
-                    }
-                    if (e.getValue() != null)
-                        taskScheduler.getTaskUnAssigner().call(e.getKey().getId(), e.getValue());
+            List<RemoveTaskRequest> l = new LinkedList<>();
+            removeTasksQueue.drainTo(l);
+            for (RemoveTaskRequest r: l) {
+                // remove it from the queue and call taskScheduler to unassign it if hostname is not null
+                try {
+                    taskQueue.getUsageTracker().removeTask(r.taskId, r.qAttributes);
+                } catch (TaskQueueException e1) {
+                    // shouldn't happen since we're calling outside of scheduling iteration
+                    logger.warn("Unexpected to get exception outside of scheduling iteration: " + e1.getMessage(), e1);
                 }
+                if (r.hostname != null)
+                    taskScheduler.getTaskUnAssigner().call(r.taskId, r.hostname);
             }
         }
     }
@@ -288,12 +300,13 @@ public void initializeRunningTask(QueuableTask task, String hostname) {
      * not the task is already running. If the task is running, the <code>hostname</code> parameter must be set, otherwise,
      * it can be <code>null</code>. The actual remove operation is performed before the start of the next scheduling
      * iteration.
-     * @param task The task to be removed.
+     * @param taskId The Id of the task to be removed.
+     * @param qAttributes The queue attributes of the queue that the task belongs to
      * @param hostname The name of the VM where the task was assigned resources from, or, <code>null</code> if it was
      *                 not assigned any resources.
      */
-    public void removeTask(QueuableTask task, String hostname) {
-        removeTasksQueue.offer(Collections.singletonMap(task, hostname));
+    public void removeTask(String taskId, QAttributes qAttributes, String hostname) {
+        removeTasksQueue.offer(new RemoveTaskRequest(taskId, qAttributes, hostname));
     }
 
     public final static class Builder {

fenzo-core/src/main/java/com/netflix/fenzo/queues/InternalTaskQueue.java

@@ -27,9 +27,8 @@
  * the queue while a scheduling iteration using this queue is in progress. Implementations must handle this. Note that,
  * it may not be sufficient for the implementations to use concurrent versions of collection classes for queue of tasks.
  * The queue must be consistent throughout the scheduling iteration. One recommended way to achieve such consistency is
- * to place the {@link #queueTask(QueuableTask)} and {@link #remove(String, QAttributes)}
- * operations as requests in a holding area within the implementation and return immediately. Later, actually carry
- * them out during the {@link #reset()} method.
+ * to place the {@link #queueTask(QueuableTask)} operations as requests in a holding area within the implementation and
+ * return immediately. Later, actually carry them out during the {@link #reset()} method.
  */
 public interface InternalTaskQueue extends TaskQueue {
     /**
@@ -41,7 +40,7 @@ public interface InternalTaskQueue extends TaskQueue {
      * adding to or removing from the queue.
      * @throws TaskQueueMultiException If any exceptions that may have occurred during resetting the pointer to the head
      * of the queue. Or, this may include exceptions that arose when applying any deferred operations from
-     * {@link #queueTask(QueuableTask)} and {@link #remove(String, QAttributes)} methods.
+     * {@link #queueTask(QueuableTask)} method.
      */
     boolean reset() throws TaskQueueMultiException;
 

fenzo-core/src/main/java/com/netflix/fenzo/samples/SampleQbasedScheduling.java

@@ -78,7 +78,8 @@ public void statusUpdate(SchedulerDriver driver, Protos.TaskStatus status) {
                 case TASK_LOST:
                 case TASK_FINISHED:
                     System.out.println("Task status for " + status.getTaskId().getValue() + ": " + status.getState());
-                    schedSvcGetter.get().removeTask(allTasks.get(status.getTaskId().getValue()),
+                    schedSvcGetter.get().removeTask(status.getTaskId().getValue(),
+                            allTasks.get(status.getTaskId().getValue()).getQAttributes(),
                             tasksToHostnameMap.get(status.getTaskId().getValue()));
                     numTasksCompleted.incrementAndGet();
             }

fenzo-core/src/test/java/com/netflix/fenzo/TaskSchedulingServiceTest.java

@@ -428,7 +428,7 @@ public void call(List<VirtualMachineCurrentState> states) {
             Assert.fail("Didn't get vm states in time");
         }
         Assert.assertTrue("Did not find task on vm", found.get());
-        schedulingService.removeTask(task, leases.get(0).hostname());
+        schedulingService.removeTask(task.getId(), task.getQAttributes(), leases.get(0).hostname());
         found.set(false);
         final CountDownLatch latch3 = new CountDownLatch(1);
         schedulingService.requestVmCurrentStates(new Action1<List<VirtualMachineCurrentState>>() {