WorkQueue: Prevent premature removal of tasks on cancellation

Previously, interrupted tasks were removed from the WorkQueue
immediately. This caused them to disappear from the show-queue
output even though they were still running in the background and
occupying an executor thread.

This behavior was especially confusing when using server-side
quotas via the `quota` plugin, as tasks appeared gone while still
affecting resource usage.

Ensure that a task is removed immediately from the executor only
if its state never reached to `STARTED`. If it has already begun
execution, we rely on the finally block in `run()` to remove it
at the correct time.

Release-Notes: show-queue now shows cancelled tasks until completion
Change-Id: I9e9b129fb66a39de9bc280f8df623006daf72b32

java/com/google/gerrit/server/git/WorkQueue.java

diff --git a/java/com/google/gerrit/server/git/WorkQueue.java b/java/com/google/gerrit/server/git/WorkQueue.java
index 93608ee..5160e3e 100644
--- a/java/com/google/gerrit/server/git/WorkQueue.java
+++ b/java/com/google/gerrit/server/git/WorkQueue.java
@@ -621,10 +621,7 @@
     }
 
     void remove(Task<?> task) {
-      boolean isRemoved = all.remove(task.getTaskId(), task);
-      if (isRemoved) {
-        cancelIfParked(task);
-      }
+      all.remove(task.getTaskId(), task);
     }
 
     void cancelIfParked(Task<?> task) {
@@ -853,6 +850,7 @@
     @CanIgnoreReturnValue
     public boolean cancel(boolean mayInterruptIfRunning) {
       if (task.cancel(mayInterruptIfRunning)) {
+        boolean isSetRunningDuringCancellation = false;
         // Tiny abuse of runningState: if the task needs to know it
         // was canceled (to clean up resources) and it hasn't started
         // yet the task's run method won't execute. So we tag it
@@ -861,6 +859,7 @@
         //
         if (runnable instanceof CancelableRunnable) {
           if (runningState.compareAndSet(null, State.RUNNING)) {
+            isSetRunningDuringCancellation = true;
             ((CancelableRunnable) runnable).cancel();
           } else if (runnable instanceof CanceledWhileRunning) {
             ((CanceledWhileRunning) runnable).setCanceledWhileRunning();
@@ -875,8 +874,11 @@
           ((Future<?>) runnable).cancel(mayInterruptIfRunning);
         }
 
-        executor.remove(this);
-        executor.purge();
+        if (isSetRunningDuringCancellation || runningState.get() == null) {
+          executor.remove(this);
+          executor.purge();
+        }
+        executor.cancelIfParked(this);
         return true;
       }
       return false;

javatests/com/google/gerrit/acceptance/server/util/WorkQueueIT.java

diff --git a/javatests/com/google/gerrit/acceptance/server/util/WorkQueueIT.java b/javatests/com/google/gerrit/acceptance/server/util/WorkQueueIT.java
index 21a4d96..6c86b68 100644
--- a/javatests/com/google/gerrit/acceptance/server/util/WorkQueueIT.java
+++ b/javatests/com/google/gerrit/acceptance/server/util/WorkQueueIT.java
@@ -18,11 +18,16 @@
 
 import com.google.gerrit.acceptance.AbstractDaemonTest;
 import com.google.gerrit.extensions.annotations.Exports;
+import com.google.gerrit.server.config.ConfigResource;
 import com.google.gerrit.server.git.WorkQueue;
+import com.google.gerrit.server.git.WorkQueue.Task.State;
+import com.google.gerrit.server.restapi.config.ListTasks;
 import com.google.inject.AbstractModule;
 import com.google.inject.Inject;
 import com.google.inject.Module;
+import java.util.concurrent.Callable;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
@@ -45,12 +50,14 @@
   }
 
   private static final Integer FIXED_RATE_SCHEDULE_INITIAL_DELAY = 0;
-  private static final Integer FIXED_RATE_SCHEDULE_INTERVAL_MILLI_SEC = 1000;
+  private static final Integer FIXED_RATE_SCHEDULE_INTERVAL_MILLI_SEC = 200;
   private static final Integer POOL_CORE_SIZE = 8;
   private static final String QUEUE_NAME = "test-Queue";
   private static final Integer EXCEPT_RUN_TIMES = 2;
+  private static final Integer TIMEOUT_MILLIS = 500;
   private final CountDownLatch downLatch = new CountDownLatch(EXCEPT_RUN_TIMES);
   @Inject private WorkQueue workQueue;
+  @Inject private ListTasks listTasks;
   private TestListener testListener;
 
   @Override
@@ -82,4 +89,50 @@
     assertThat(ifRunMoreThanOnce).isTrue();
     testExecutor.shutdownNow();
   }
+
+  @Test
+  public void testCanceledTaskStaysUntilFinished() throws Exception {
+    ScheduledExecutorService testExecutor = workQueue.createQueue(POOL_CORE_SIZE, QUEUE_NAME);
+    CountDownLatch latch = new CountDownLatch(1);
+    Future taskFuture =
+        testExecutor.submit(
+            () -> {
+              try {
+                latch.await();
+              } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+              }
+            });
+    assertTasksInStateEventually(QUEUE_NAME, State.RUNNING, 1);
+
+    taskFuture.cancel(false);
+    assertTasksInStateEventually(QUEUE_NAME, State.CANCELLED, 1);
+
+    latch.countDown();
+    // task is now removed after completion
+    assertEventually(
+        () ->
+            listTasks.apply(new ConfigResource()).value().stream()
+                .noneMatch(t -> t.queueName.equals(QUEUE_NAME)));
+    testExecutor.shutdownNow();
+  }
+
+  public void assertTasksInStateEventually(String queue, State expectedState, int expectedCount)
+      throws Exception {
+    assertEventually(
+        () ->
+            expectedCount
+                == listTasks.apply(new ConfigResource()).value().stream()
+                    .filter(t -> t.queueName.equals(queue))
+                    .filter(t -> t.state.equals(expectedState))
+                    .count());
+  }
+
+  public void assertEventually(Callable<Boolean> r) throws Exception {
+    long ms = 0;
+    while (r.call() != true) {
+      assertThat(ms++).isLessThan(TIMEOUT_MILLIS);
+      TimeUnit.MILLISECONDS.sleep(1);
+    }
+  }
 }