Rework the TaskQueue implementation as well as implementing a close state.

Motivation:

The TaskQueue does not implement closeability, since now such queue can hold many suspended virtual thread tasks, we should provide a way to deal those threads (interrupt) when the context holding the queue is closed, e.g.  undeploying a verticle or closing a vertx instance.

Changes:

Implement a TaskQueue close method that returns the list of current thread being suspended. The context holding the queue can close the queue when the context close future is destroyed.

Author: vietj

src/main/java/io/vertx/core/Future.java

@@ -22,6 +22,7 @@
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -668,13 +669,14 @@ static <T> Future<T> fromCompletionStage(CompletionStage<T> completionStage, Con
    */
   static <T> T await(Future<T> future) {
     io.vertx.core.impl.WorkerExecutor executor = io.vertx.core.impl.WorkerExecutor.unwrapWorkerExecutor();
-    io.vertx.core.impl.WorkerExecutor.TaskController cont = executor.current();
-    future.onComplete(ar -> cont.resume());
-    try {
-      cont.suspendAndAwaitResume();
-    } catch (InterruptedException e) {
-      Utils.throwAsUnchecked(e);
-      return null;
+    CountDownLatch latch = executor.suspend(cont -> future.onComplete(ar -> cont.resume()));
+    if (latch != null) {
+      try {
+        latch.await();
+      } catch (InterruptedException e) {
+        Utils.throwAsUnchecked(e);
+        return null;
+      }
     }
     if (future.succeeded()) {
       return future.result();

src/main/java/io/vertx/core/impl/TaskQueue.java

@@ -14,11 +14,11 @@
 import io.vertx.core.impl.logging.Logger;
 import io.vertx.core.impl.logging.LoggerFactory;
 
-import java.util.LinkedList;
+import java.util.*;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executor;
 import java.util.concurrent.RejectedExecutionException;
-import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
 
 /**
  * A task queue that always run all tasks in order. The executor to run the tasks is passed
@@ -38,6 +38,8 @@ public class TaskQueue {
 
   // @protectedby tasks
   private final LinkedList<Task> tasks = new LinkedList<>();
+  private final Set<ContinuationTask> suspendedTasks = new HashSet<>();
+  private boolean closed;
   private Executor currentExecutor;
   private Thread currentThread;
 
@@ -51,13 +53,16 @@ private void run() {
     for (; ; ) {
       final ExecuteTask execute;
       synchronized (tasks) {
+        if (closed) {
+          return;
+        }
         Task task = tasks.poll();
         if (task == null) {
           currentExecutor = null;
           return;
         }
-        if (task instanceof ResumeTask) {
-          ResumeTask resume = (ResumeTask) task;
+        if (task instanceof ContinuationTask) {
+          ContinuationTask resume = (ContinuationTask) task;
           currentExecutor = resume.executor;
           currentThread = resume.thread;
           resume.latch.run();
@@ -83,12 +88,18 @@ private void run() {
   }
 
   /**
-   * Return a controller for the current task.
+   * A task of this queue.
+   */
+  private interface Task {
+  }
+
+  /**
+   * Return a continuation task for the current task execution.
    *
    * @return the controller
    * @throws IllegalStateException if the current thread is not currently being executed by the queue
    */
-  public WorkerExecutor.TaskController current() {
+  private ContinuationTask continuationTask() {
     Thread thread;
     Executor executor;
     synchronized (tasks) {
@@ -98,51 +109,19 @@ public WorkerExecutor.TaskController current() {
       thread = currentThread;
       executor = currentExecutor;
     }
-    return new WorkerExecutor.TaskController() {
-
-      final CountDownLatch latch = new CountDownLatch(1);
-
-      @Override
-      public void resume(Runnable callback) {
-        Runnable task = () -> {
-          callback.run();
-          latch.countDown();
-        };
-        synchronized (tasks) {
-          if (currentExecutor != null) {
-            tasks.addFirst(new ResumeTask(task, executor, thread));
-            return;
-          }
-          currentExecutor = executor;
-          currentThread = thread;
-        }
-        task.run();
-      }
-
-      @Override
-      public CountDownLatch suspend() {
-        if (Thread.currentThread() != thread) {
-          throw new IllegalStateException();
-        }
-        synchronized (tasks) {
-          if (currentThread == null || currentThread != Thread.currentThread()) {
-            throw new IllegalStateException();
-          }
-          currentThread = null;
-        }
-        executor.execute(runner);
-        return latch;
-      }
-    };
+    return new ContinuationTask(thread, executor);
   }
 
   /**
    * Run a task.
    *
    * @param task the task to run.
    */
-  public void execute(Runnable task, Executor executor) {
+  public void execute(Runnable task, Executor executor) throws RejectedExecutionException {
     synchronized (tasks) {
+      if (closed) {
+        throw new RejectedExecutionException("Closed");
+      }
       if (currentExecutor == null) {
         currentExecutor = executor;
         try {
@@ -168,9 +147,171 @@ public boolean isEmpty() {
   }
 
   /**
-   * A task of this queue.
+   * Structure holding the queue state at close time.
    */
-  private interface Task {
+  public final static class CloseResult {
+
+    private final Thread activeThread;
+    private final List<Runnable> pendingTasks;
+    private final List<Thread> suspendedThreads;
+
+    private CloseResult(Thread activeThread, List<Thread> suspendedThreads, List<Runnable> pendingTasks) {
+      this.activeThread = activeThread;
+      this.suspendedThreads = suspendedThreads;
+      this.pendingTasks = pendingTasks;
+    }
+
+    /**
+     * @return the thread that was active
+     */
+    public Thread activeThread() {
+      return activeThread;
+    }
+
+    /**
+     * @return the list of pending tasks
+     */
+    public List<Runnable> pendingTasks() {
+      return pendingTasks;
+    }
+
+    /**
+     * @return the list of suspended threads
+     */
+    public List<Thread> suspendedThreads() {
+      return suspendedThreads;
+    }
+  }
+
+  /**
+   * Close the queue.
+   *
+   * @return a structure of suspended threads and pending tasks
+   */
+  public CloseResult close() {
+    List<Runnable> pendingTasks = Collections.emptyList();
+    List<Thread> suspendedThreads;
+    Thread currentThread;
+    synchronized (tasks) {
+      if (closed) {
+        throw new IllegalStateException("Already closed");
+      }
+      suspendedThreads = new ArrayList<>(suspendedTasks.size() + 1);
+      for (Task t : tasks) {
+        if (t instanceof ExecuteTask) {
+          if (pendingTasks.isEmpty()) {
+            pendingTasks = new LinkedList<>();
+          }
+          pendingTasks.add(((ExecuteTask)t).runnable);
+        } else if (t instanceof ContinuationTask) {
+          ContinuationTask rt = (ContinuationTask) t;
+          suspendedThreads.add(rt.thread);
+        }
+      }
+      tasks.clear();
+      for (ContinuationTask task : suspendedTasks) {
+        suspendedThreads.add(task.thread);
+      }
+      suspendedTasks.clear();
+      currentThread = this.currentThread;
+      currentExecutor = null;
+      closed = true;
+    }
+    return new CloseResult(currentThread, suspendedThreads, pendingTasks);
+  }
+
+  private class ContinuationTask extends CountDownLatch implements WorkerExecutor.Continuation, Task {
+
+    private static final int ST_CREATED = 0, ST_SUSPENDED = 1, ST_RESUMED = 2;
+
+    private final Thread thread;
+    private final Executor executor;
+    private int status;
+    private Runnable latch;
+
+    public ContinuationTask(Thread thread, Executor executor) {
+      super(1);
+      this.thread = thread;
+      this.executor = executor;
+      this.status = ST_CREATED;
+    }
+
+    @Override
+    public void resume(Runnable callback) {
+      synchronized (tasks) {
+        if (closed) {
+          return;
+        }
+        switch (status) {
+          case ST_SUSPENDED:
+            boolean removed = suspendedTasks.remove(this);
+            assert removed;
+            latch = () -> {
+              callback.run();
+              countDown();
+            };
+            if (currentExecutor != null) {
+              tasks.addFirst(this);
+              return;
+            }
+            currentExecutor = executor;
+            currentThread = thread;
+            break;
+          case ST_CREATED:
+            // The current task still owns the queue
+            assert currentExecutor == executor;
+            assert currentThread == thread;
+            latch = callback;
+            break;
+          default:
+            throw new IllegalStateException();
+        }
+        status = ST_RESUMED;
+      }
+      latch.run();
+    }
+
+    public boolean suspend() {
+      if (Thread.currentThread() != thread) {
+        throw new IllegalStateException();
+      }
+      synchronized (tasks) {
+        if (closed) {
+          return false;
+        }
+        if (currentThread == null || currentThread != thread) {
+          throw new IllegalStateException();
+        }
+        switch (status) {
+          case ST_RESUMED:
+            countDown();
+            return false;
+          case ST_SUSPENDED:
+            throw new IllegalStateException();
+        }
+        status = ST_SUSPENDED;
+        boolean added = suspendedTasks.add(this);
+        assert added;
+        currentThread = null;
+      }
+      executor.execute(runner);
+      return true;
+    }
+  }
+
+  public CountDownLatch suspend() {
+    return suspend(cont -> {});
+  }
+
+  public CountDownLatch suspend(Consumer<WorkerExecutor.Continuation> abc) {
+    ContinuationTask continuationTask = continuationTask();
+    abc.accept(continuationTask);
+    if (continuationTask.suspend()) {
+      return continuationTask;
+    } else {
+      // Closed
+      return null;
+    }
   }
 
   /**
@@ -184,18 +325,4 @@ public ExecuteTask(Runnable runnable, Executor exec) {
       this.exec = exec;
     }
   }
-
-  /**
-   * Resume an existing task blocked on a thread
-   */
-  private static class ResumeTask implements Task {
-    private final Runnable latch;
-    private final Executor executor;
-    private final Thread thread;
-    ResumeTask(Runnable latch, Executor executor, Thread thread) {
-      this.latch = latch;
-      this.executor = executor;
-      this.thread = thread;
-    }
-  }
 }