Merge pull request #4915 from eclipse-vertx/backport-worker-improvements

Backport worker improvements

Author: vietj

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

@@ -15,9 +15,10 @@
 import io.vertx.core.impl.logging.LoggerFactory;
 
 import java.util.LinkedList;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executor;
 import java.util.concurrent.RejectedExecutionException;
-import java.util.function.Consumer;
+import java.util.concurrent.TimeUnit;
 
 /**
  * A task queue that always run all tasks in order. The executor to run the tasks is passed
@@ -82,17 +83,12 @@ private void run() {
   }
 
   /**
-   * Unschedule the current task from execution, the next task in the queue will be executed
-   * when there is one.
+   * Return a controller for the current task.
    *
-   * <p>When the current task wants to be resumed, it should call the returned consumer with a command
-   * to unpark the thread (e.g most likely yielding a latch), this task will be executed immediately if there
-   * is no tasks being executed, otherwise it will be added first in the queue.
-   *
-   * @return a mean to signal to resume the thread when it shall be resumed
+   * @return the controller
    * @throws IllegalStateException if the current thread is not currently being executed by the queue
    */
-  public Consumer<Runnable> unschedule() {
+  public WorkerExecutor.TaskController current() {
     Thread thread;
     Executor executor;
     synchronized (tasks) {
@@ -101,20 +97,42 @@ public Consumer<Runnable> unschedule() {
       }
       thread = currentThread;
       executor = currentExecutor;
-      currentThread = null;
     }
-    executor.execute(runner);
-    return r -> {
-      synchronized (tasks) {
-        if (currentExecutor != null) {
-          tasks.addFirst(new ResumeTask(r, executor, thread));
-          return;
-        } else {
+    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;
       }
-      r.run();
     };
   }
 

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

@@ -535,11 +535,20 @@ public ContextImpl createEventLoopContext() {
     return createEventLoopContext(null, closeFuture, null, Thread.currentThread().getContextClassLoader());
   }
 
-  @Override
-  public ContextImpl createWorkerContext(Deployment deployment, CloseFuture closeFuture, WorkerPool workerPool, ClassLoader tccl) {
+  private ContextImpl createWorkerContext(EventLoop eventLoop, CloseFuture closeFuture, WorkerPool workerPool, Deployment deployment, ClassLoader tccl) {
     TaskQueue orderedTasks = new TaskQueue();
     WorkerPool wp = workerPool != null ? workerPool : this.workerPool;
-    return new ContextImpl(this, false, eventLoopGroup.next(), new WorkerExecutor(wp, orderedTasks), internalWorkerPool, wp, orderedTasks, deployment, closeFuture, disableTCCL ? null : tccl);
+    return new ContextImpl(this, false, eventLoop, new WorkerExecutor(wp, orderedTasks), internalWorkerPool, wp, orderedTasks, deployment, closeFuture, disableTCCL ? null : tccl);
+  }
+
+  @Override
+  public ContextInternal createWorkerContext(EventLoop eventLoop, WorkerPool workerPool, ClassLoader tccl) {
+    return createWorkerContext(eventLoop, closeFuture, workerPool, null, tccl);
+  }
+
+  @Override
+  public ContextImpl createWorkerContext(Deployment deployment, CloseFuture closeFuture, WorkerPool workerPool, ClassLoader tccl) {
+    return createWorkerContext(eventLoopGroup.next(), closeFuture, workerPool, deployment, tccl);
   }
 
   @Override

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

@@ -118,15 +118,29 @@ default <C> C createSharedClient(String clientKey, String clientName, CloseFutur
    */
   ContextInternal createEventLoopContext(Deployment deployment, CloseFuture closeFuture, WorkerPool workerPool, ClassLoader tccl);
 
+  /**
+   * @return event loop context
+   */
   ContextInternal createEventLoopContext(EventLoop eventLoop, WorkerPool workerPool, ClassLoader tccl);
 
+  /**
+   * @return event loop context
+   */
   ContextInternal createEventLoopContext();
 
   /**
-   * @return worker loop context
+   * @return worker context
    */
-  ContextInternal createWorkerContext(Deployment deployment, CloseFuture closeFuture, WorkerPool pool, ClassLoader tccl);
+  ContextInternal createWorkerContext(Deployment deployment, CloseFuture closeFuture, WorkerPool workerPool, ClassLoader tccl);
 
+  /**
+   * @return worker context
+   */
+  ContextInternal createWorkerContext(EventLoop eventLoop, WorkerPool workerPool, ClassLoader tccl);
+
+  /**
+   * @return worker context
+   */
   ContextInternal createWorkerContext();
 
   @Override

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

@@ -454,8 +454,13 @@ public ContextInternal createEventLoopContext() {
   }
 
   @Override
-  public ContextInternal createWorkerContext(Deployment deployment, CloseFuture closeFuture, WorkerPool pool, ClassLoader tccl) {
-    return delegate.createWorkerContext(deployment, closeFuture, pool, tccl);
+  public ContextInternal createWorkerContext(EventLoop eventLoop, WorkerPool workerPool, ClassLoader tccl) {
+    return delegate.createWorkerContext(eventLoop, workerPool, tccl);
+  }
+
+  @Override
+  public ContextInternal createWorkerContext(Deployment deployment, CloseFuture closeFuture, WorkerPool workerPool, ClassLoader tccl) {
+    return delegate.createWorkerContext(deployment, closeFuture, workerPool, tccl);
   }
 
   @Override

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

@@ -12,7 +12,8 @@
 
 import io.vertx.core.spi.metrics.PoolMetrics;
 
-import java.util.function.Consumer;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 
 /**
  * Execute events on a worker pool.
@@ -60,9 +61,46 @@ public void execute(Runnable command) {
   }
 
   /**
-   * See {@link TaskQueue#unschedule()}.
+   * See {@link TaskQueue#current()}.
    */
-  public Consumer<Runnable> unschedule() {
-    return orderedTasks.unschedule();
+  public TaskController current() {
+    return orderedTasks.current();
+  }
+
+  public interface TaskController {
+
+    /**
+     * Resume the task, the {@code callback} will be executed when the task is resumed, before the task thread
+     * is unparked.
+     *
+     * @param callback called when the task is resumed
+     */
+    void resume(Runnable callback);
+
+    /**
+     * Like {@link #resume(Runnable)}.
+     */
+    default void resume() {
+      resume(() -> {});
+    }
+
+    /**
+     * Suspend the task execution and park the current thread until the task is resumed.
+     * The next task in the queue will be executed, when there is one.
+     *
+     * <p>When the task wants to be resumed, it should call {@link #resume}, this will be executed immediately if there
+     * is no other tasks being executed, otherwise it will be added first in the queue.
+     */
+    default void suspendAndAwaitResume() throws InterruptedException {
+      suspend().await();
+    }
+
+    /**
+     * Like {@link #suspendAndAwaitResume()} but does not await the task to be resumed.
+     *
+     * @return the latch to await
+     */
+    CountDownLatch suspend();
+
   }
 }