ArC: temporary fix for quarkusio#3186

FroMage · Jul 15, 2019 · 76eb976 · 76eb976
1 parent ad04f97
commit 76eb976
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 7 deletions.
diff --git a/extensions/arc/runtime/src/main/java/io/quarkus/arc/runtime/context/ArcContextProvider.java b/extensions/arc/runtime/src/main/java/io/quarkus/arc/runtime/context/ArcContextProvider.java
@@ -3,6 +3,8 @@
 import java.util.Collection;
 import java.util.Map;
 
+import javax.enterprise.context.spi.Contextual;
+
 import org.eclipse.microprofile.context.ThreadContext;
 import org.eclipse.microprofile.context.spi.ThreadContextController;
 import org.eclipse.microprofile.context.spi.ThreadContextProvider;
@@ -36,7 +38,7 @@ public ThreadContextSnapshot currentContext(Map<String, String> map) {
         }
 
         // capture all instances
-        Collection<ContextInstanceHandle<?>> instancesToPropagate = arc.requestContext().getAll();
+        Map<Contextual<?>, ContextInstanceHandle<?>> capturedContext = arc.requestContext().getContext();
         return () -> {
             // can be called later on, we should retrieve the container again
             ArcContainer arcContainer = Arc.container();
@@ -48,17 +50,14 @@ public ThreadContextSnapshot currentContext(Map<String, String> map) {
             // this is executed on another thread, context can but doesn't need to be active here
             if (isContextActiveOnThisThread(arcContainer)) {
                 // context active, store current state, feed it new one and restore state afterwards
-                Collection<ContextInstanceHandle<?>> instancesToRestore = requestContext.getAll();
-                requestContext.deactivate();
-                requestContext.activate(instancesToPropagate);
+                Map<Contextual<?>, ContextInstanceHandle<?>> contextToRestore = requestContext.setContext(capturedContext);
                 controller = () -> {
                     // clean up, reactivate context with previous values
-                    requestContext.deactivate();
-                    requestContext.activate(instancesToRestore);
+                    requestContext.setContext(contextToRestore);
                 };
             } else {
                 // context not active, activate and pass it new instance, deactivate afterwards
-                requestContext.activate(instancesToPropagate);
+                requestContext.setContext(capturedContext);
                 controller = () -> {
                     requestContext.deactivate();
                 };

diff --git a/independent-projects/arc/runtime/src/main/java/io/quarkus/arc/ManagedContext.java b/independent-projects/arc/runtime/src/main/java/io/quarkus/arc/ManagedContext.java
@@ -1,6 +1,9 @@
 package io.quarkus.arc;
 
 import java.util.Collection;
+import java.util.Map;
+
+import javax.enterprise.context.spi.Contextual;
 
 /**
  *
@@ -36,4 +39,9 @@ default void terminate() {
         destroy();
         deactivate();
     }
+
+
+    Map<Contextual<?>, ContextInstanceHandle<?>> getContext();
+
+    Map<Contextual<?>, ContextInstanceHandle<?>> setContext(Map<Contextual<?>, ContextInstanceHandle<?>> context);
 }
diff --git a/independent-projects/arc/runtime/src/main/java/io/quarkus/arc/RequestContext.java b/independent-projects/arc/runtime/src/main/java/io/quarkus/arc/RequestContext.java
@@ -157,4 +157,26 @@ public void destroy() {
             }
         }
     }
+
+    @Override
+    public Map<Contextual<?>, ContextInstanceHandle<?>> getContext() {
+        return currentContext.get();
+    }
+
+    @Override
+    public Map<Contextual<?>, ContextInstanceHandle<?>> setContext(Map<Contextual<?>, ContextInstanceHandle<?>> context) {
+        if(context != null) {
+            for (ContextInstanceHandle<?> instanceHandle : context.values()) {
+                if (!instanceHandle.getBean().getScope().equals(getScope())) {
+                    throw new IllegalArgumentException("Invalid bean scope: " + instanceHandle.getBean());
+                }
+            }
+        }
+        Map<Contextual<?>, ContextInstanceHandle<?>> previousContext = currentContext.get();
+        if(context == null)
+            currentContext.remove();
+        else
+            currentContext.set(context);
+        return previousContext;
+    }
 }