Handle correctly future await from an execute blocking on a virtual thread.

Motivation:

Future#await ensures that we use a virtual thread context to obtain an execute to suspend the virtual thread, however this check also includes execute blocking tasks.

Changes:

When obtaining the worker executor to suspend the current task, ensure that we are on the context thread, so execute blocking tasks are not included.

Author: vietj

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

@@ -20,10 +20,8 @@
 import io.vertx.core.impl.future.SucceededFuture;
 
 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.Objects;
+import java.util.concurrent.*;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
@@ -668,14 +666,46 @@ static <T> Future<T> fromCompletionStage(CompletionStage<T> completionStage, Con
    * @throws IllegalStateException when called from an event-loop thread or a non Vert.x thread
    */
   static <T> T await(Future<T> future) {
+    try {
+      return await(future, -1, null);
+    } catch (TimeoutException e) {
+      // Not a possible case
+      return null;
+    }
+
+  }
+
+  /**
+   * Like {@link #await(Future)} but with a timeout.
+   *
+   * @param timeout the timeout
+   * @param unit the timeout unit
+   * @return the result
+   * @throws TimeoutException when the timeout fires before the future completes
+   * @throws IllegalStateException when called from a vertx event-loop or worker thread
+   */
+  public static <T> T await(Future<T> future, long timeout, TimeUnit unit) throws TimeoutException {
+    if (timeout >= 0L && unit == null) {
+      throw new NullPointerException();
+    }
     io.vertx.core.impl.WorkerExecutor executor = io.vertx.core.impl.WorkerExecutor.unwrapWorkerExecutor();
-    if (executor == null) {
-      throw new IllegalStateException();
+    CountDownLatch latch;
+    if (executor != null) {
+      latch = executor.suspend(cont -> future.onComplete(ar -> cont.resume()));
+    } else {
+      latch = new CountDownLatch(1);
+      future.onComplete(ar -> latch.countDown());
     }
-    CountDownLatch latch = executor.suspend(cont -> future.onComplete(ar -> cont.resume()));
     if (latch != null) {
       try {
-        latch.await();
+        if (timeout >= 0) {
+          Objects.requireNonNull(unit);
+          if (!latch.await(timeout, unit)) {
+            throw new TimeoutException();
+          }
+        } else {
+          latch.await();
+        }
       } catch (InterruptedException e) {
         Utils.throwAsUnchecked(e);
         return null;
@@ -691,5 +721,4 @@ static <T> T await(Future<T> future) {
       return null;
     }
   }
-
 }

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

@@ -36,7 +36,8 @@ public static io.vertx.core.impl.WorkerExecutor unwrapWorkerExecutor() {
       throw new IllegalStateException(msg);
     }
     ContextInternal ctx = VertxImpl.currentContext(thread);
-    if (ctx != null && ctx.threadingModel() == ThreadingModel.VIRTUAL_THREAD) {
+    if (ctx != null && ctx.inThread()) {
+      // It can only be a Vert.x virtual thread
       return (io.vertx.core.impl.WorkerExecutor) ctx.executor();
     } else {
       return null;

src/test/java/io/vertx/core/FutureAwaitTest.java

@@ -0,0 +1,57 @@
+package io.vertx.core;
+
+import io.vertx.core.impl.ContextInternal;
+import io.vertx.test.core.VertxTestBase;
+import org.junit.Test;
+
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+public class FutureAwaitTest extends VertxTestBase {
+
+  @Test
+  public void testAwaitFromEventLoopThread() {
+    Promise<String> promise = Promise.promise();
+    Future<String> future = promise.future();
+    ContextInternal ctx = (ContextInternal) vertx.getOrCreateContext();
+    ctx.nettyEventLoop().execute(() -> {
+      try {
+        Future.await(future);
+      } catch (IllegalStateException expected) {
+        testComplete();
+      }
+    });
+    await();
+  }
+
+  @Test
+  public void testAwaitFromNonVertxThread() {
+    Promise<String> promise = Promise.promise();
+    Future<String> future = promise.future();
+    Thread current = Thread.currentThread();
+    new Thread(() -> {
+      while (current.getState() != Thread.State.WAITING) {
+        try {
+          Thread.sleep(10);
+        } catch (InterruptedException ignore) {
+        }
+      }
+      promise.complete("the-result");
+    }).start();
+    String res = Future.await(future);
+    assertEquals("the-result", res);
+  }
+
+  @Test
+  public void testAwaitWithTimeout() {
+    Promise<String> promise = Promise.promise();
+    Future<String> future = promise.future();
+    long now = System.currentTimeMillis();
+    try {
+      Future.await(future, 100, TimeUnit.MILLISECONDS);
+      fail();
+    } catch (TimeoutException expected) {
+    }
+    assertTrue((System.currentTimeMillis() - now) >= 100);
+  }
+}

src/test/java/io/vertx/core/FutureTest.java

@@ -1821,15 +1821,6 @@ public void testAndThenCompleteHandlerWithError() {
     await();
   }
 
-  @Test
-  public void testAwaitFromPlainThread() {
-    try {
-      Future.await(Promise.promise().future());
-      fail();
-    } catch (IllegalStateException e) {
-    }
-  }
-
   @Test
   public void contextFutureTimeoutFires() {
     ContextInternal ctx = (ContextInternal) vertx.getOrCreateContext();

src/test/java/io/vertx/core/VirtualThreadContextTest.java

@@ -387,4 +387,30 @@ public void testSubmitAfterClose() {
     });
     await();
   }
+
+  @Test
+  public void testAwaitFromVirtualThreadExecuteBlocking() {
+    Assume.assumeTrue(isVirtualThreadAvailable());
+    Context ctx = vertx.createVirtualThreadContext();
+    ctx.executeBlocking(() -> {
+      Future.await(vertx.timer(20));
+      return "done";
+    }).onComplete(onSuccess(res -> {
+      assertEquals("done", res);
+      testComplete();
+    }));
+    await();
+  }
+
+  @Test
+  public void testAwaitFromWorkerExecuteBlocking() {
+    Context ctx = vertx.getOrCreateContext();
+    ctx.executeBlocking(() -> {
+      Future.await(vertx.timer(20));
+      return "done";
+    }).onComplete(onFailure(res -> {
+      testComplete();
+    }));
+    await();
+  }
 }