TaskHandle should not count forced splits as running splits

presto-main/src/main/java/com/facebook/presto/execution/TaskExecutor.java

@@ -191,7 +191,7 @@ public synchronized List<ListenableFuture<?>> enqueueSplits(TaskHandle taskHandl
                 // Note: we do not record queued time for forced splits
                 startSplit(prioritizedSplitRunner);
                 // add the runner to the handle so it can be destroyed if the task is canceled
-                taskHandle.recordRunningSplit(prioritizedSplitRunner);
+                taskHandle.recordForcedRunningSplit(prioritizedSplitRunner);
             }
             else {
                 // add this to the work queue for the task
@@ -287,6 +287,7 @@ public static class TaskHandle
         private final TaskId taskId;
         private final Queue<PrioritizedSplitRunner> queuedSplits = new ArrayDeque<>(10);
         private final List<PrioritizedSplitRunner> runningSplits = new ArrayList<>(10);
+        private final List<PrioritizedSplitRunner> forcedRunningSplits = new ArrayList<>(10);
         private final AtomicLong taskThreadUsageNanos = new AtomicLong();
 
         private final AtomicInteger nextSplitId = new AtomicInteger();
@@ -308,6 +309,11 @@ private TaskId getTaskId()
 
         private void destroy()
         {
+            for (PrioritizedSplitRunner runningSplit : forcedRunningSplits) {
+                runningSplit.destroy();
+            }
+            forcedRunningSplits.clear();
+
             for (PrioritizedSplitRunner runningSplit : runningSplits) {
                 runningSplit.destroy();
             }
@@ -324,12 +330,13 @@ private void enqueueSplit(PrioritizedSplitRunner split)
             queuedSplits.add(split);
         }
 
-        private void recordRunningSplit(PrioritizedSplitRunner split)
+        private void recordForcedRunningSplit(PrioritizedSplitRunner split)
         {
-            runningSplits.add(split);
+            forcedRunningSplits.add(split);
         }
 
-        private int getRunningSplits()
+        @VisibleForTesting
+        int getRunningSplits()
         {
             return runningSplits.size();
         }
@@ -350,6 +357,7 @@ private PrioritizedSplitRunner pollNextSplit()
 
         private void splitComplete(PrioritizedSplitRunner split)
         {
+            forcedRunningSplits.remove(split);
             runningSplits.remove(split);
         }
 

presto-main/src/test/java/com/facebook/presto/execution/TaskExecutorTest.java

@@ -113,6 +113,35 @@ public void test()
         }
     }
 
+    @Test
+    public void testTaskHandle()
+            throws Exception
+    {
+        TaskExecutor taskExecutor = new TaskExecutor(4, 8);
+        taskExecutor.start();
+
+        try {
+            TaskHandle taskHandle = taskExecutor.addTask(new TaskId("test", "test", "test"));
+
+            Phaser beginPhase = new Phaser();
+            beginPhase.register();
+            Phaser verificationComplete = new Phaser();
+            verificationComplete.register();
+            TestingJob driver = new TestingJob(beginPhase, verificationComplete, 10);
+
+            // force enqueue a split
+            taskExecutor.enqueueSplits(taskHandle, true, ImmutableList.of(driver));
+            assertEquals(taskHandle.getRunningSplits(), 0);
+
+            // normal enqueue a split
+            taskExecutor.enqueueSplits(taskHandle, false, ImmutableList.of(driver));
+            assertEquals(taskHandle.getRunningSplits(), 1);
+        }
+        finally {
+            taskExecutor.stop();
+        }
+    }
+
     private static class TestingJob
             implements SplitRunner
     {