Merge pull request quarkusio#38560 from mkouba/vertx-current-context-fix

VertxCurrentContextFactory: support multiple CDI contexts
bpasson · Feb 5, 2024 · 1fd25ca · 1fd25ca
2 parents cf7ae80 + 681b165
commit 1fd25ca
Show file tree

Hide file tree

Showing 4 changed files with 88 additions and 21 deletions.
diff --git a/...s/vertx/deployment/src/main/java/io/quarkus/vertx/core/deployment/VertxCoreProcessor.java b/...s/vertx/deployment/src/main/java/io/quarkus/vertx/core/deployment/VertxCoreProcessor.java
@@ -59,6 +59,7 @@
 import io.quarkus.vertx.core.runtime.VertxLogDelegateFactory;
 import io.quarkus.vertx.core.runtime.config.VertxConfiguration;
 import io.quarkus.vertx.core.runtime.context.SafeVertxContextInterceptor;
+import io.quarkus.vertx.deployment.VertxBuildConfig;
 import io.quarkus.vertx.mdc.provider.LateBoundMDCProvider;
 import io.vertx.core.AbstractVerticle;
 import io.vertx.core.Vertx;
@@ -273,8 +274,8 @@ ThreadFactoryBuildItem createVertxThreadFactory(VertxCoreRecorder recorder, Laun
 
     @BuildStep
     @Record(ExecutionTime.RUNTIME_INIT)
-    ContextHandlerBuildItem createVertxContextHandlers(VertxCoreRecorder recorder) {
-        return new ContextHandlerBuildItem(recorder.executionContextHandler());
+    ContextHandlerBuildItem createVertxContextHandlers(VertxCoreRecorder recorder, VertxBuildConfig buildConfig) {
+        return new ContextHandlerBuildItem(recorder.executionContextHandler(buildConfig.customizeArcContext()));
     }
 
     private void handleBlockingWarningsInDevOrTestMode() {

diff --git a/extensions/vertx/runtime/src/main/java/io/quarkus/vertx/core/runtime/VertxCoreRecorder.java b/extensions/vertx/runtime/src/main/java/io/quarkus/vertx/core/runtime/VertxCoreRecorder.java
@@ -13,6 +13,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Locale;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
@@ -564,7 +565,10 @@ private static void setNewThreadTccl(VertxThread thread) {
         thread.setContextClassLoader(cl);
     }
 
-    public ContextHandler<Object> executionContextHandler() {
+    public ContextHandler<Object> executionContextHandler(boolean customizeArcContext) {
+        VertxCurrentContextFactory currentContextFactory = customizeArcContext
+                ? (VertxCurrentContextFactory) Arc.container().getCurrentContextFactory()
+                : null;
         return new ContextHandler<Object>() {
             @Override
             public Object captureContext() {
@@ -574,17 +578,21 @@ public Object captureContext() {
             @Override
             public void runWith(Runnable task, Object context) {
                 ContextInternal currentContext = (ContextInternal) Vertx.currentContext();
+                // Only do context handling if it's non-null
                 if (context != null && context != currentContext) {
-                    // Only do context handling if it's non-null
                     ContextInternal vertxContext = (ContextInternal) context;
-                    // The CDI request context must not be propagated
-                    ConcurrentMap<Object, Object> local = vertxContext.localContextData();
-                    if (local.containsKey(VertxCurrentContextFactory.LOCAL_KEY)) {
-                        // Duplicate the context, copy the data, remove the request context
-                        vertxContext = vertxContext.duplicate();
-                        vertxContext.localContextData().putAll(local);
-                        vertxContext.localContextData().remove(VertxCurrentContextFactory.LOCAL_KEY);
-                        VertxContextSafetyToggle.setContextSafe(vertxContext, true);
+                    // The CDI contexts must not be propagated
+                    // First test if VertxCurrentContextFactory is actually used
+                    if (currentContextFactory != null) {
+                        List<String> keys = currentContextFactory.keys();
+                        ConcurrentMap<Object, Object> local = vertxContext.localContextData();
+                        if (containsScopeKey(keys, local)) {
+                            // Duplicate the context, copy the data, remove the request context
+                            vertxContext = vertxContext.duplicate();
+                            vertxContext.localContextData().putAll(local);
+                            keys.forEach(vertxContext.localContextData()::remove);
+                            VertxContextSafetyToggle.setContextSafe(vertxContext, true);
+                        }
                     }
                     vertxContext.beginDispatch();
                     try {
@@ -596,6 +604,23 @@ public void runWith(Runnable task, Object context) {
                     task.run();
                 }
             }
+
+            private boolean containsScopeKey(List<String> keys, Map<Object, Object> localContextData) {
+                if (keys.isEmpty()) {
+                    return false;
+                }
+                if (keys.size() == 1) {
+                    // Very often there will be only one key used
+                    return localContextData.containsKey(keys.get(0));
+                } else {
+                    for (String key : keys) {
+                        if (localContextData.containsKey(key)) {
+                            return true;
+                        }
+                    }
+                }
+                return false;
+            }
         };
     }
 

diff --git a/...ions/vertx/runtime/src/main/java/io/quarkus/vertx/runtime/VertxCurrentContextFactory.java b/...ions/vertx/runtime/src/main/java/io/quarkus/vertx/runtime/VertxCurrentContextFactory.java
@@ -1,6 +1,9 @@
 package io.quarkus.vertx.runtime;
 
 import java.lang.annotation.Annotation;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
 
 import io.netty.util.concurrent.FastThreadLocal;
 import io.quarkus.arc.CurrentContext;
@@ -14,22 +17,52 @@
 
 public class VertxCurrentContextFactory implements CurrentContextFactory {
 
-    public static final String LOCAL_KEY = "io.quarkus.vertx.cdi-current-context";
+    private static final String LOCAL_KEY_PREFIX = "io.quarkus.vertx.cdi-current-context";
+
+    private final List<String> keys;
+    private final List<String> unmodifiableKeys;
+
+    public VertxCurrentContextFactory() {
+        // There will be only a few mutative operations max
+        this.keys = new CopyOnWriteArrayList<>();
+        // We do not want to allocate a new object for each VertxCurrentContextFactory#keys() invocation
+        this.unmodifiableKeys = Collections.unmodifiableList(keys);
+    }
 
     @Override
     public <T extends InjectableContext.ContextState> CurrentContext<T> create(Class<? extends Annotation> scope) {
-        return new VertxCurrentContext<>();
+        String key = LOCAL_KEY_PREFIX + scope.getName();
+        if (keys.contains(key)) {
+            throw new IllegalStateException(
+                    "Multiple current contexts for the same scope are not supported. Current context for "
+                            + scope + " already exists!");
+        }
+        keys.add(key);
+        return new VertxCurrentContext<>(key);
+    }
+
+    /**
+     *
+     * @return an unmodifiable list of used keys
+     */
+    public List<String> keys() {
+        return unmodifiableKeys;
     }
 
     private static final class VertxCurrentContext<T extends ContextState> implements CurrentContext<T> {
 
+        private final String key;
         private final FastThreadLocal<T> fallback = new FastThreadLocal<>();
 
+        private VertxCurrentContext(String key) {
+            this.key = key;
+        }
+
         @Override
         public T get() {
             Context context = Vertx.currentContext();
             if (context != null && VertxContext.isDuplicatedContext(context)) {
-                return context.getLocal(LOCAL_KEY);
+                return context.getLocal(key);
             }
             return fallback.get();
         }
@@ -41,9 +74,9 @@ public void set(T state) {
                 VertxContextSafetyToggle.setContextSafe(context, true);
                 // this is racy but should be fine, because DC should not be shared
                 // and never remove the existing mapping
-                var oldState = context.getLocal(LOCAL_KEY);
+                var oldState = context.getLocal(key);
                 if (oldState != state) {
-                    context.putLocal(LOCAL_KEY, state);
+                    context.putLocal(key, state);
                 }
 
             } else {
@@ -56,7 +89,7 @@ public void remove() {
             Context context = Vertx.currentContext();
             if (context != null && VertxContext.isDuplicatedContext(context)) {
                 // NOOP - the DC should not be shared.
-                // context.removeLocal(LOCAL_KEY);
+                // context.removeLocal(key);
             } else {
                 fallback.remove();
             }

diff --git a/independent-projects/arc/runtime/src/main/java/io/quarkus/arc/CurrentContextFactory.java b/independent-projects/arc/runtime/src/main/java/io/quarkus/arc/CurrentContextFactory.java
@@ -5,12 +5,20 @@
 import io.quarkus.arc.InjectableContext.ContextState;
 
 /**
- * This factory can be used to create a new {@link CurrentContext} for a normal scope, e.g. for
- * {@link jakarta.enterprise.context.RequestScoped}. It's usually not necessary for shared contexts, such as
- * {@link jakarta.enterprise.context.ApplicationScoped}.
+ * This factory can be used to create a new {@link CurrentContext} for a normal scope.
+ * <p>
+ * For example, the current context for {@link jakarta.enterprise.context.RequestScoped}. It's usually not necessary for shared
+ * contexts, such as {@link jakarta.enterprise.context.ApplicationScoped}.
  */
 public interface CurrentContextFactory {
 
+    /**
+     *
+     * @param <T>
+     * @param scope
+     * @return the current context
+     * @throws IllegalStateException If the implementation does not support multiple current contexts for the same scope
+     */
     <T extends ContextState> CurrentContext<T> create(Class<? extends Annotation> scope);
 
 }