Virtual context improvements

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

@@ -738,9 +738,12 @@ default T await(long timeout, TimeUnit unit) throws TimeoutException {
     }
     if (succeeded()) {
       return result();
-    } else {
+    } else if (failed()) {
       Utils.throwAsUnchecked(cause());
       return null;
+    } else {
+      Utils.throwAsUnchecked(new InterruptedException("Context closed"));
+      return null;
     }
   }
 

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

@@ -54,9 +54,6 @@ private void run() {
     for (; ; ) {
       final ExecuteTask execute;
       synchronized (tasks) {
-        if (closed) {
-          return;
-        }
         Task task = tasks.poll();
         if (task == null) {
           currentExecutor = null;
@@ -125,9 +122,6 @@ private ContinuationTask continuationTask() {
    */
   public void execute(Runnable task, Executor executor) throws RejectedExecutionException {
     synchronized (tasks) {
-      if (closed) {
-        throw new RejectedExecutionException("Closed");
-      }
       if (currentExecutor == null) {
         currentExecutor = executor;
         try {
@@ -161,18 +155,15 @@ public final static class CloseResult {
     private final Runnable activeTask;
     private final List<Runnable> suspendedTasks;
     private final List<Thread> suspendedThreads;
-    private final List<Runnable> pendingTasks;
 
     private CloseResult(Thread activeThread,
                         Runnable activeTask,
                         List<Thread> suspendedThreads,
-                        List<Runnable> suspendedTasks,
-                        List<Runnable> pendingTasks) {
+                        List<Runnable> suspendedTasks) {
       this.activeThread = activeThread;
       this.activeTask = activeTask;
       this.suspendedThreads = suspendedThreads;
       this.suspendedTasks = suspendedTasks;
-      this.pendingTasks = pendingTasks;
     }
 
     /**
@@ -186,13 +177,6 @@ public Runnable activeTask() {
       return activeTask;
     }
 
-    /**
-     * @return the list of pending tasks
-     */
-    public List<Runnable> pendingTasks() {
-      return pendingTasks;
-    }
-
     /**
      * @return the list of suspended threads
      */
@@ -214,7 +198,6 @@ public List<Runnable> suspendedTasks() {
    * @return a structure of suspended threads and pending tasks
    */
   public CloseResult close() {
-    List<Runnable> pendingTasks = Collections.emptyList();
     List<Thread> suspendedThreads;
     List<Runnable> suspendedTasks;
     Thread activeThread;
@@ -225,19 +208,16 @@ public CloseResult close() {
       }
       suspendedThreads = new ArrayList<>(continuations.size());
       suspendedTasks = new ArrayList<>(continuations.size());
-      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);
-          suspendedTasks.add(rt.task.runnable);
+      Iterator<Task> it = tasks.iterator();
+      while (it.hasNext()) {
+        Task task = it.next();
+        if (task instanceof ContinuationTask) {
+          ContinuationTask continuationTask = (ContinuationTask) task;
+          suspendedThreads.add(continuationTask.thread);
+          suspendedTasks.add(continuationTask.task.runnable);
+          it.remove();
         }
       }
-      tasks.clear();
       for (ContinuationTask cont : continuations) {
         suspendedThreads.add(cont.thread);
         suspendedTasks.add(cont.task.runnable);
@@ -248,7 +228,7 @@ public CloseResult close() {
       currentExecutor = null;
       closed = true;
     }
-    return new CloseResult(activeThread, activeTask, suspendedThreads, suspendedTasks, pendingTasks);
+    return new CloseResult(activeThread, activeTask, suspendedThreads, suspendedTasks);
   }
 
   private class ContinuationTask extends CountDownLatch implements WorkerExecutor.Continuation, Task {

vertx-core/src/main/java/io/vertx/core/impl/WorkerTaskQueue.java

@@ -34,13 +34,6 @@ public WorkerTaskQueue() {
   void shutdown(EventLoop executor, Promise<Void> completion) {
     TaskQueue.CloseResult closeResult = close();
 
-    // Reject all pending tasks
-    List<Runnable> pendingTasks = closeResult.pendingTasks();
-    for (Runnable pendingTask : pendingTasks) {
-      WorkerTask pendingWorkerTask = (WorkerTask) pendingTask;
-      pendingWorkerTask.reject();
-    }
-
     // Maintain context invariant: serialize task execution while interrupting tasks
     class InterruptSequence {
 

vertx-core/src/test/java/io/vertx/tests/context/ContextTest.java

@@ -250,15 +250,11 @@ public void testExecuteBlockingClose() {
       latch.await();
       return "";
     });
-    Future<String> fut2 = ctx.executeBlocking(() -> "");
     assertWaitUntil(() -> thread.get() != null && thread.get().getState() == Thread.State.WAITING);
     ctx.close();
     assertWaitUntil(fut1::isComplete);
     assertTrue(fut1.failed());
     assertTrue(fut1.cause() instanceof InterruptedException);
-    assertWaitUntil(fut2::isComplete);
-    assertTrue(fut2.failed());
-    assertTrue(fut2.cause() instanceof RejectedExecutionException);
   }
 
   @Test

vertx-core/src/test/java/io/vertx/tests/context/TaskQueueTest.java

@@ -202,20 +202,6 @@ public void shouldNotHaveTaskInQueueWhenTaskHasBeenRejected() {
     Assertions.assertThat(taskQueue.isEmpty()).isTrue();
   }
 
-  @Test
-  public void testClosePendingTasks() {
-    TaskQueue taskQueue = new TaskQueue();
-    Deque<Runnable> pending = new ConcurrentLinkedDeque<>();
-    Executor executor = pending::add;
-    Runnable task = () -> {
-    };
-    taskQueue.execute(task, executor);
-    assertEquals(1, pending.size());
-    TaskQueue.CloseResult result = taskQueue.close();
-    assertEquals(1, result.pendingTasks().size());
-    assertSame(task, result.pendingTasks().get(0));
-  }
-
   @Test
   public void testCloseSuspendedTasks() {
     TaskQueue taskQueue = new TaskQueue();
@@ -325,14 +311,10 @@ public void testSubmitAfterClose() {
     taskQueue.close();
     Deque<Runnable> pending = new ConcurrentLinkedDeque<>();
     Executor exec = pending::add;
-    try {
-      taskQueue.execute(() -> {
+    taskQueue.execute(() -> {
 
-      }, exec);
-      fail();
-    } catch (RejectedExecutionException expected) {
-    }
-    assertEquals(0, pending.size());
+    }, exec);
+    assertEquals(1, pending.size());
   }
 
   @Test

vertx-core/src/test/java/io/vertx/tests/virtualthread/VirtualThreadContextTest.java

@@ -353,4 +353,42 @@ public void testContextCloseContextSerialization() throws Exception {
     }
     f.await();
   }
+
+  @Test
+  public void testAwaitWhenClosed() throws Exception {
+    Assume.assumeTrue(isVirtualThreadAvailable());
+    ContextInternal ctx = vertx.createVirtualThreadContext();
+    CountDownLatch latch = new CountDownLatch(1);
+    ctx.runOnContext(v -> {
+      latch.countDown();
+      try {
+        new CountDownLatch(1).await();
+        fail();
+      } catch (InterruptedException expected) {
+        assertFalse(Thread.currentThread().isInterrupted());
+      }
+      try {
+        Promise.promise().future().await();
+        fail();
+      } catch (Exception e) {
+        assertEquals(InterruptedException.class, e.getClass());
+        testComplete();
+      }
+    });
+    awaitLatch(latch);
+    // Interrupts virtual thread
+    ctx.close();
+    await();
+  }
+
+  @Test
+  public void testSubmitAfterClose() {
+    Assume.assumeTrue(isVirtualThreadAvailable());
+    ContextInternal ctx = vertx.createVirtualThreadContext();
+    ctx.close();
+    ctx.runOnContext(v -> {
+      testComplete();
+    });
+    await();
+  }
 }