From ba2e83c9cbc81b313319a0df9b8761a3b499776c Mon Sep 17 00:00:00 2001
From: Clement Escoffier <clement@apache.org>
Date: Tue, 4 Jul 2023 16:06:10 +0200
Subject: [PATCH] Make @RunOnVirtualThread capture and run in a duplicated
 context

---
 .../runtime/ResteasyReactiveRecorder.java     | 30 +++++++++++++++++--
 1 file changed, 28 insertions(+), 2 deletions(-)

diff --git a/extensions/resteasy-reactive/quarkus-resteasy-reactive/runtime/src/main/java/io/quarkus/resteasy/reactive/server/runtime/ResteasyReactiveRecorder.java b/extensions/resteasy-reactive/quarkus-resteasy-reactive/runtime/src/main/java/io/quarkus/resteasy/reactive/server/runtime/ResteasyReactiveRecorder.java
index fd802f5f30e62..74bdd6d25443e 100644
--- a/extensions/resteasy-reactive/quarkus-resteasy-reactive/runtime/src/main/java/io/quarkus/resteasy/reactive/server/runtime/ResteasyReactiveRecorder.java
+++ b/extensions/resteasy-reactive/quarkus-resteasy-reactive/runtime/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 " +