Adapt to vertx 4 and improve the implementation.

Author: vietj

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

@@ -669,6 +669,9 @@ 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();
+    if (executor == null) {
+      throw new IllegalStateException();
+    }
     CountDownLatch latch = executor.suspend(cont -> future.onComplete(ar -> cont.resume()));
     if (latch != null) {
       try {

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

@@ -51,9 +51,9 @@ public CloseFuture(Logger log) {
    *
    * @param hook the hook to add
    */
-  public synchronized void add(Closeable hook) {
+  public synchronized boolean add(Closeable hook) {
     if (closed) {
-      throw new IllegalStateException();
+      return false;
     }
     if (hook instanceof CloseFuture) {
       // Close future might be closed independently, so we optimize and remove the hooks when
@@ -72,6 +72,7 @@ public synchronized void add(Closeable hook) {
       }
       hooks.put(hook, this);
     }
+    return true;
   }
 
   /**

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

@@ -56,7 +56,7 @@ static <T> void setResultHandler(ContextInternal ctx, Future<T> fut, Handler<Asy
   final TaskQueue internalOrderedTasks;
   final WorkerPool internalWorkerPool;
   final WorkerPool workerPool;
-  final TaskQueue orderedTasks;
+  final WorkerTaskQueue orderedTasks;
 
   public ContextImpl(VertxInternal vertx,
                      int localsLength,
@@ -65,7 +65,7 @@ public ContextImpl(VertxInternal vertx,
                      EventExecutor executor,
                      WorkerPool internalWorkerPool,
                      WorkerPool workerPool,
-                     TaskQueue orderedTasks,
+                     WorkerTaskQueue orderedTasks,
                      Deployment deployment,
                      CloseFuture closeFuture,
                      ClassLoader tccl) {
@@ -84,6 +84,14 @@ public ContextImpl(VertxInternal vertx,
     this.internalOrderedTasks = new TaskQueue();
   }
 
+  public Future<Void> close() {
+    if (closeFuture == owner.closeFuture()) {
+      return Future.future(p -> orderedTasks.shutdown(eventLoop, p));
+    } else {
+      return closeFuture.close().eventually(() -> Future.<Void>future(p -> orderedTasks.shutdown(eventLoop, p)));
+    }
+  }
+
   public Deployment getDeployment() {
     return deployment;
   }
@@ -201,29 +209,29 @@ private static <T> Future<T> internalExecuteBlocking(ContextInternal context, Ha
     Object queueMetric = metrics != null ? metrics.submitted() : null;
     Promise<T> promise = context.promise();
     Future<T> fut = promise.future();
-    try {
-      Runnable command = () -> {
-        Object execMetric = null;
-        if (metrics != null) {
-          execMetric = metrics.begin(queueMetric);
-        }
+    WorkerTask task = new WorkerTask(metrics, queueMetric) {
+      @Override
+      protected void execute() {
         context.dispatch(promise, blockingCodeHandler);
+      }
+      @Override
+      void reject() {
         if (metrics != null) {
-          metrics.end(execMetric, fut.succeeded());
+          metrics.rejected(queueMetric);
         }
-      };
+        promise.fail(new RejectedExecutionException());
+      }
+    };
+    try {
       Executor exec = workerPool.executor();
       if (queue != null) {
-        queue.execute(command, exec);
+        queue.execute(task, exec);
       } else {
-        exec.execute(command);
+        exec.execute(task);
       }
     } catch (RejectedExecutionException e) {
       // Pool is already shut down
-      if (metrics != null) {
-        metrics.rejected(queueMetric);
-      }
-      throw e;
+      task.reject();
     }
     return fut;
   }

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

@@ -275,6 +275,13 @@ default <E> void dispatch(E event, Handler<E> handler) {
     }
   }
 
+  /**
+   * Close this context, cleanup close future hooks then dispose pending ordered task queue.
+   *
+   * @return a future signalling close completion
+   */
+  Future<Void> close();
+
   /**
    * Begin the execution of a task on this context.
    * <p>

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

@@ -200,7 +200,7 @@ private Future<Deployment> doDeploy(String identifier,
           context = vertx.createVirtualThreadContext(deployment, closeFuture, tccl);
           break;
       }
-      VerticleHolder holder = new VerticleHolder(verticle, context, closeFuture);
+      VerticleHolder holder = new VerticleHolder(verticle, context);
       deployment.addVerticle(holder);
       context.runOnContext(v -> {
         try {
@@ -241,16 +241,14 @@ static class VerticleHolder {
 
     final Verticle verticle;
     final ContextImpl context;
-    final CloseFuture closeFuture;
 
-    VerticleHolder(Verticle verticle, ContextImpl context, CloseFuture closeFuture) {
+    VerticleHolder(Verticle verticle, ContextImpl context) {
       this.verticle = verticle;
       this.context = context;
-      this.closeFuture = closeFuture;
     }
 
     void close(Handler<AsyncResult<Void>> completionHandler) {
-      closeFuture.close().onComplete(completionHandler);
+      context.close().onComplete(completionHandler);
     }
   }
 

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

@@ -193,4 +193,9 @@ public ContextInternal unwrap() {
   public boolean isDuplicate() {
     return true;
   }
+
+  @Override
+  public Future<Void> close() {
+    return Future.succeededFuture();
+  }
 }

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

@@ -38,10 +38,11 @@ public class TaskQueue {
 
   // @protectedby tasks
   private final LinkedList<Task> tasks = new LinkedList<>();
-  private final Set<ContinuationTask> suspendedTasks = new HashSet<>();
+  private final Set<ContinuationTask> continuations = new HashSet<>();
   private boolean closed;
   private Executor currentExecutor;
   private Thread currentThread;
+  private ExecuteTask currentTask;
 
   private final Runnable runner;
 
@@ -65,6 +66,7 @@ private void run() {
           ContinuationTask resume = (ContinuationTask) task;
           currentExecutor = resume.executor;
           currentThread = resume.thread;
+          currentTask = resume.task;
           resume.latch.run();
           return;
         }
@@ -78,11 +80,13 @@ private void run() {
       }
       try {
         currentThread = Thread.currentThread();
+        currentTask = execute;
         execute.runnable.run();
       } catch (Throwable t) {
         log.error("Caught unexpected Throwable", t);
       } finally {
         currentThread = null;
+        currentTask = null;
       }
     }
   }
@@ -100,6 +104,7 @@ private interface Task {
    * @throws IllegalStateException if the current thread is not currently being executed by the queue
    */
   private ContinuationTask continuationTask() {
+    ExecuteTask task;
     Thread thread;
     Executor executor;
     synchronized (tasks) {
@@ -108,8 +113,9 @@ private ContinuationTask continuationTask() {
       }
       thread = currentThread;
       executor = currentExecutor;
+      task = currentTask;
     }
-    return new ContinuationTask(thread, executor);
+    return new ContinuationTask(task, thread, executor);
   }
 
   /**
@@ -152,12 +158,20 @@ public boolean isEmpty() {
   public final static class CloseResult {
 
     private final Thread activeThread;
-    private final List<Runnable> pendingTasks;
+    private final Runnable activeTask;
+    private final List<Runnable> suspendedTasks;
     private final List<Thread> suspendedThreads;
+    private final List<Runnable> pendingTasks;
 
-    private CloseResult(Thread activeThread, List<Thread> suspendedThreads, List<Runnable> pendingTasks) {
+    private CloseResult(Thread activeThread,
+                        Runnable activeTask,
+                        List<Thread> suspendedThreads,
+                        List<Runnable> suspendedTasks,
+                        List<Runnable> pendingTasks) {
       this.activeThread = activeThread;
+      this.activeTask = activeTask;
       this.suspendedThreads = suspendedThreads;
+      this.suspendedTasks = suspendedTasks;
       this.pendingTasks = pendingTasks;
     }
 
@@ -168,6 +182,10 @@ public Thread activeThread() {
       return activeThread;
     }
 
+    public Runnable activeTask() {
+      return activeTask;
+    }
+
     /**
      * @return the list of pending tasks
      */
@@ -181,6 +199,13 @@ public List<Runnable> pendingTasks() {
     public List<Thread> suspendedThreads() {
       return suspendedThreads;
     }
+
+    /**
+     * @return the list of suspended tasks
+     */
+    public List<Runnable> suspendedTasks() {
+      return suspendedTasks;
+    }
   }
 
   /**
@@ -191,12 +216,15 @@ public List<Thread> suspendedThreads() {
   public CloseResult close() {
     List<Runnable> pendingTasks = Collections.emptyList();
     List<Thread> suspendedThreads;
-    Thread currentThread;
+    List<Runnable> suspendedTasks;
+    Thread activeThread;
+    Runnable activeTask;
     synchronized (tasks) {
       if (closed) {
         throw new IllegalStateException("Already closed");
       }
-      suspendedThreads = new ArrayList<>(suspendedTasks.size() + 1);
+      suspendedThreads = new ArrayList<>(continuations.size());
+      suspendedTasks = new ArrayList<>(continuations.size());
       for (Task t : tasks) {
         if (t instanceof ExecuteTask) {
           if (pendingTasks.isEmpty()) {
@@ -206,31 +234,36 @@ public CloseResult close() {
         } else if (t instanceof ContinuationTask) {
           ContinuationTask rt = (ContinuationTask) t;
           suspendedThreads.add(rt.thread);
+          suspendedTasks.add(rt.task.runnable);
         }
       }
       tasks.clear();
-      for (ContinuationTask task : suspendedTasks) {
-        suspendedThreads.add(task.thread);
+      for (ContinuationTask cont : continuations) {
+        suspendedThreads.add(cont.thread);
+        suspendedTasks.add(cont.task.runnable);
       }
-      suspendedTasks.clear();
-      currentThread = this.currentThread;
+      continuations.clear();
+      activeThread = currentThread;
+      activeTask = currentTask != null ? currentTask.runnable : null;
       currentExecutor = null;
       closed = true;
     }
-    return new CloseResult(currentThread, suspendedThreads, pendingTasks);
+    return new CloseResult(activeThread, activeTask, suspendedThreads, suspendedTasks, 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 ExecuteTask task;
     private final Thread thread;
     private final Executor executor;
     private int status;
     private Runnable latch;
 
-    public ContinuationTask(Thread thread, Executor executor) {
+    public ContinuationTask(ExecuteTask task, Thread thread, Executor executor) {
       super(1);
+      this.task = task;
       this.thread = thread;
       this.executor = executor;
       this.status = ST_CREATED;
@@ -244,7 +277,7 @@ public void resume(Runnable callback) {
         }
         switch (status) {
           case ST_SUSPENDED:
-            boolean removed = suspendedTasks.remove(this);
+            boolean removed = continuations.remove(this);
             assert removed;
             latch = () -> {
               callback.run();
@@ -256,11 +289,13 @@ public void resume(Runnable callback) {
             }
             currentExecutor = executor;
             currentThread = thread;
+            currentTask = task;
             break;
           case ST_CREATED:
             // The current task still owns the queue
             assert currentExecutor == executor;
             assert currentThread == thread;
+            assert currentTask == task;
             latch = callback;
             break;
           default:
@@ -290,9 +325,10 @@ public boolean suspend() {
             throw new IllegalStateException();
         }
         status = ST_SUSPENDED;
-        boolean added = suspendedTasks.add(this);
+        boolean added = continuations.add(this);
         assert added;
         currentThread = null;
+        currentTask = null;
       }
       executor.execute(runner);
       return true;