Merge pull request #34524 from cescoffier/runOnVirtualThread-duplicat…

…ed-context
quarkusio · Jul 5, 2023 · e153bcc · e153bcc
2 parents e0de475 + ba2e83c
commit e153bcc
Showing 1 changed file with 28 additions and 2 deletions.
diff --git a/...e/src/main/java/io/quarkus/resteasy/reactive/server/runtime/ResteasyReactiveRecorder.java b/...e/src/main/java/io/quarkus/resteasy/reactive/server/runtime/ResteasyReactiveRecorder.java
@@ -65,6 +65,8 @@
 import io.quarkus.vertx.http.runtime.security.HttpSecurityRecorder.DefaultAuthFailureHandler;
 import io.quarkus.vertx.http.runtime.security.QuarkusHttpUser;
 import io.vertx.core.Handler;
+import io.vertx.core.Vertx;
+import io.vertx.core.impl.ContextInternal;
 import io.vertx.ext.web.RoutingContext;
 
 @Recorder
@@ -106,15 +108,39 @@ private Executor setVirtualThreadCustomScheduler(Executor executor) throws Class
          * change --release, --source, --target flags and to enable previews.
          * Since we try to load the "Loom-preview" classes/methods at runtime, the application can even be compiled
          * using java 11 and executed with a loom-compliant JDK.
+         * <p>
+         * IMPORTANT: we still need to use a duplicated context to have all the propagation working.
+         * Thus, the context is captured and applied/terminated in the virtual thread.
          */
         @Override
         public Executor get() {
             if (current == null) {
                 try {
-                    current = (Executor) Executors.class.getMethod("newVirtualThreadPerTaskExecutor")
+                    var virtual = (Executor) Executors.class.getMethod("newVirtualThreadPerTaskExecutor")
                             .invoke(this);
+                    current = new Executor() {
+                        @Override
+                        public void execute(Runnable command) {
+                            var context = Vertx.currentContext();
+                            if (!(context instanceof ContextInternal)) {
+                                virtual.execute(command);
+                            } else {
+                                virtual.execute(new Runnable() {
+                                    @Override
+                                    public void run() {
+                                        final var previousContext = ((ContextInternal) context).beginDispatch();
+                                        try {
+                                            command.run();
+                                        } finally {
+                                            ((ContextInternal) context).endDispatch(previousContext);
+                                        }
+                                    }
+                                });
+                            }
+                        }
+                    };
                 } catch (InvocationTargetException | IllegalAccessException | NoSuchMethodException e) {
-                    System.err.println(e);
+                    logger.debug("Unable to invoke java.util.concurrent.Executors#newVirtualThreadPerTaskExecutor", e);
                     //quite ugly but works
                     logger.warnf("You weren't able to create an executor that spawns virtual threads, the default" +
                             " blocking executor will be used, please check that your JDK is compatible with " +