Merge pull request #98 from spodila/rmTaskFromTrackrFromQ

move the removeTask function from TaskQueue to TaskSchedulingService

Author: spodila

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

@@ -49,8 +49,7 @@
  *         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 com.netflix.fenzo.queues.TaskQueue#remove(String, QAttributes)} when tasks complete or they no
- *         longer need resource assignments.
+ *         {@link #removeTask(QueuableTask, String)} when tasks complete or they no longer need resource assignments.
  *     </LI>
  * </UL>
  */
@@ -64,6 +63,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<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);
@@ -123,6 +123,7 @@ private void scheduleOnce() {
             // check if next scheduling iteration is actually needed right away
             final boolean qModified = taskQueue.reset();
             addPendingRunningTasks();
+            removeTasks();
             final boolean newLeaseExists = leaseBlockingQueue.peek() != null;
             if ( qModified || newLeaseExists || doNextIteration()) {
                 lastSchedIterationAt.set(System.currentTimeMillis());
@@ -158,6 +159,26 @@ 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());
+                }
+            }
+        }
+    }
+
     private void doPendingActions() {
         final Action1<Map<TaskQueue.TaskState, Collection<QueuableTask>>> action = taskMapRequest.poll();
         try {
@@ -262,6 +283,19 @@ public void initializeRunningTask(QueuableTask task, String hostname) {
         addRunningTasksQueue.offer(Collections.singletonMap(hostname, task));
     }
 
+    /**
+     * Mark the task to be removed. This is expected to be called for all tasks that were added to the queue, whether or
+     * 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 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 final static class Builder {
 
         private TaskScheduler taskScheduler = null;

fenzo-core/src/main/java/com/netflix/fenzo/functions/Action0.java

@@ -1,3 +1,19 @@
+/*
+ * Copyright 2016 Netflix, Inc.
+ *
+ * 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 com.netflix.fenzo.functions;
 
 /**

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

@@ -54,16 +54,4 @@ enum TaskState { QUEUED, LAUNCHED }
      * @param task A task to add to the queue.
      */
     void queueTask(QueuableTask task);
-
-    /**
-     * Remove a task from queue with the given {@code taskId} and {@code qAttributes}. The task is removed from the
-     * queue, whether it is marked as queued for assignments or already marked as running. This must be called for all
-     * tasks that either no longer need resource assignments or if previously running tasks complete for any reason.
-     * <P>
-     * This operation is designed to be performed asynchronously, when it is safe to modify the queue. The queue
-     * implementations generally do not modify the queue while a scheduling iteration is in progress.
-     * @param taskId The id of the task to remove.
-     * @param qAttributes The queue attributes of the task to remove.
-     */
-    void remove(String taskId, QAttributes qAttributes);
 }

fenzo-core/src/main/java/com/netflix/fenzo/queues/tiered/QueueBucket.java

@@ -1,3 +1,19 @@
+/*
+ * Copyright 2016 Netflix, Inc.
+ *
+ * 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 com.netflix.fenzo.queues.tiered;
 
 import com.netflix.fenzo.queues.*;

fenzo-core/src/main/java/com/netflix/fenzo/queues/tiered/TieredQueue.java

@@ -40,7 +40,6 @@ public class TieredQueue implements InternalTaskQueue {
     private Iterator<Tier> iterator = null;
     private Tier currTier = null;
     private final BlockingQueue<QueuableTask> tasksToQueue;
-    private final BlockingQueue<QAttributes.TaskIdAttributesTuple> taskIdsToRemove;
 
     /**
      * Construct a tiered queue system with the given number of tiers.
@@ -51,7 +50,6 @@ public TieredQueue(int numTiers) {
         for ( int i=0; i<numTiers; i++ )
             tiers.add(new Tier(i));
         tasksToQueue = new LinkedBlockingQueue<>();
-        taskIdsToRemove = new LinkedBlockingQueue<>();
     }
 
     @Override
@@ -73,16 +71,6 @@ private void addRunningInternal(QueuableTask t) throws TaskQueueException {
         tiers.get(number).launchTask(t);
     }
 
-    @Override
-    public void remove(final String taskId, final QAttributes qAttributes) {
-        taskIdsToRemove.offer(new QAttributes.TaskIdAttributesTuple(taskId, qAttributes));
-    }
-
-    private boolean removeInternalById(String id, QAttributes qAttributes) throws TaskQueueException {
-        final Tier tierBuckets = tiers.get(qAttributes.getTierNumber());
-        return tierBuckets != null && tierBuckets.removeTask(id, qAttributes) != null;
-    }
-
     /**
      * This implementation dynamically picks the next task to consider for resource assignment based on tiers and then
      * based on current dominant resource usage. The usage is updated with each resource assignment during the
@@ -134,23 +122,6 @@ public boolean reset() throws TaskQueueMultiException {
                     }
             }
         }
-        if (taskIdsToRemove.peek() != null) {
-            List<QAttributes.TaskIdAttributesTuple> taskIdTuples = new ArrayList<>();
-            taskIdsToRemove.drainTo(taskIdTuples);
-            if (!taskIdTuples.isEmpty()) {
-                for (QAttributes.TaskIdAttributesTuple tuple : taskIdTuples) {
-                    try {
-                        if (!removeInternalById(tuple.getId(), tuple.getqAttributes())) {
-                            logger.debug("Task with id " + tuple.getId() + " not found to remove");
-                        } else {
-                            queueChanged = true;
-                        }
-                    } catch (TaskQueueException e) {
-                        exceptions.add(e);
-                    }
-                }
-            }
-        }
         if (!exceptions.isEmpty())
             throw new TaskQueueMultiException(exceptions);
         return queueChanged;

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

@@ -1,3 +1,19 @@
+/*
+ * Copyright 2016 Netflix, Inc.
+ *
+ * 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 com.netflix.fenzo.samples;
 
 import com.netflix.fenzo.*;
@@ -15,21 +31,24 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 
 public class SampleQbasedScheduling {
 
     private static class MesosScheduler implements Scheduler {
 
         private final AtomicInteger numTasksCompleted;
-        private final TaskQueue queue;
+        private final AtomicReference<TaskSchedulingService> schedSvcGetter;
         private Action1<List<Protos.Offer>> leaseAction = null;
 
-        MesosScheduler(AtomicInteger numTasksCompleted, TaskQueue queue) {
+        MesosScheduler(AtomicInteger numTasksCompleted, AtomicReference<TaskSchedulingService> schedSvcGetter) {
             this.numTasksCompleted = numTasksCompleted;
-            this.queue = queue;
+            this.schedSvcGetter = schedSvcGetter;
         }
 
         @Override
@@ -59,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());
-                    queue.remove(status.getTaskId().getValue(), qAttribs);
+                    schedSvcGetter.get().removeTask(allTasks.get(status.getTaskId().getValue()),
+                            tasksToHostnameMap.get(status.getTaskId().getValue()));
                     numTasksCompleted.incrementAndGet();
             }
         }
@@ -92,6 +112,9 @@ public void error(SchedulerDriver driver, String message) {
 
     private final static QAttributes qAttribs = new QAttributes.QAttributesAdaptor(0, "onlyBucket");
 
+    private final static ConcurrentMap<String, QueuableTask> allTasks = new ConcurrentHashMap<>();
+    private final static ConcurrentMap<String, String> tasksToHostnameMap = new ConcurrentHashMap<>();
+
     /**
      * This is the main method of this sample framework. It showcases how to use Fenzo queues for scheduling. It creates
      * some number of tasks and launches them into Mesos using the Mesos built-in command executor. The tasks launched
@@ -123,7 +146,8 @@ public void call(VirtualMachineLease v) {
         final TaskQueue queue = TaskQueues.createTieredQueue(2);
 
         // Create our Mesos scheduler callback implementation
-        final MesosScheduler mesosSchedulerCallback = new MesosScheduler(numTasksCompleted, queue);
+        AtomicReference<TaskSchedulingService> schedSvcGetter = new AtomicReference<>();
+        final MesosScheduler mesosSchedulerCallback = new MesosScheduler(numTasksCompleted, schedSvcGetter);
 
         // create Mesos driver
         Protos.FrameworkInfo framework = Protos.FrameworkInfo.newBuilder()
@@ -162,12 +186,14 @@ public void call(SchedulingResult schedulingResult) {
                                 for (VirtualMachineLease l: e.getValue().getLeasesUsed())
                                     offers.add(l.getOffer().getId());
                                 List<Protos.TaskInfo> taskInfos = new ArrayList<Protos.TaskInfo>();
-                                for (TaskAssignmentResult r: e.getValue().getTasksAssigned())
+                                for (TaskAssignmentResult r: e.getValue().getTasksAssigned()) {
                                     taskInfos.add(SampleFramework.getTaskInfo(
                                             e.getValue().getLeasesUsed().iterator().next().getOffer().getSlaveId(),
                                             r.getTaskId(),
                                             "sleep 2"
                                     ));
+                                    tasksToHostnameMap.put(r.getTaskId(), r.getHostname());
+                                }
                                 driver.launchTasks(
                                         offers,
                                         taskInfos
@@ -177,6 +203,7 @@ public void call(SchedulingResult schedulingResult) {
                     }
                 })
                 .build();
+        schedSvcGetter.set(schedulingService);
 
         // set up action in our scheduler callback to send resource offers into our scheduling service
         mesosSchedulerCallback.leaseAction = new Action1<List<Protos.Offer>>() {
@@ -199,8 +226,11 @@ public void run() {
         }.start();
 
         // submit some tasks
-        for (int i=0; i<numTasks; i++)
-            queue.queueTask(getTask(i));
+        for (int i=0; i<numTasks; i++) {
+            final QueuableTask task = getTask(i);
+            allTasks.put(task.getId(), task);
+            queue.queueTask(task);
+        }
 
         // wait for tasks to complete
         while (numTasksCompleted.get() < numTasks) {

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

@@ -29,6 +29,7 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -387,7 +388,68 @@ public void call(Map<TaskQueue.TaskState, Collection<QueuableTask>> stateCollect
 
     @Test
     public void testRemoveFromQueue() throws Exception {
-        // TODO
+        final CountDownLatch latch = new CountDownLatch(1);
+        TaskQueue queue = TaskQueues.createTieredQueue(2);
+        final TaskScheduler scheduler = getScheduler();
+        Action1<SchedulingResult> resultCallback = new Action1<SchedulingResult>() {
+            @Override
+            public void call(SchedulingResult schedulingResult) {
+                //System.out.println("Got scheduling result with " + schedulingResult.getResultMap().size() + " results");
+                if (schedulingResult.getResultMap().size() > 0) {
+                    //System.out.println("Assignment on host " + schedulingResult.getResultMap().values().iterator().next().getHostname());
+                    latch.countDown();
+                }
+            }
+        };
+        final TaskSchedulingService schedulingService = getSchedulingService(queue, scheduler, 100L, 200L, resultCallback);
+        schedulingService.start();
+        final List<VirtualMachineLease> leases = LeaseProvider.getLeases(1, 4, 4000, 1, 10);
+        schedulingService.addLeases(leases);
+        final QueuableTask task = QueuableTaskProvider.wrapTask(tier1bktA, TaskRequestProvider.getTaskRequest(2, 2000, 1));
+        queue.queueTask(task);
+        if (!latch.await(5, TimeUnit.SECONDS))
+            Assert.fail("Did not assign resources in time");
+        final CountDownLatch latch2 = new CountDownLatch(1);
+        final AtomicBoolean found = new AtomicBoolean();
+        schedulingService.requestVmCurrentStates(new Action1<List<VirtualMachineCurrentState>>() {
+            @Override
+            public void call(List<VirtualMachineCurrentState> states) {
+                for (VirtualMachineCurrentState s: states) {
+                    for (TaskRequest t: s.getRunningTasks()) {
+                        if (t.getId().equals(task.getId())) {
+                            found.set(true);
+                            latch2.countDown();
+                        }
+                    }
+                }
+            }
+        });
+        if (!latch2.await(5, TimeUnit.SECONDS)) {
+            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());
+        found.set(false);
+        final CountDownLatch latch3 = new CountDownLatch(1);
+        schedulingService.requestVmCurrentStates(new Action1<List<VirtualMachineCurrentState>>() {
+            @Override
+            public void call(List<VirtualMachineCurrentState> states) {
+                for (VirtualMachineCurrentState s: states) {
+                    for (TaskRequest t: s.getRunningTasks()) {
+                        if (t.getId().equals(task.getId())) {
+                            found.set(true);
+                            latch3.countDown();
+                        }
+                    }
+                }
+                latch3.countDown();
+            }
+        });
+        if (!latch3.await(5, TimeUnit.SECONDS)) {
+            Assert.fail("Timeout waiting for vm states");
+        }
+        Assert.assertFalse("Unexpected to find removed task on vm", found.get());
+        scheduler.shutdown();
     }
 
     @Test