From 76eb976d1499562d3b1b48c6432dbf4a484dc590 Mon Sep 17 00:00:00 2001
From: Stephane Epardaud <stef@epardaud.fr>
Date: Thu, 11 Jul 2019 16:26:35 +0200
Subject: [PATCH] ArC: temporary fix for
 https://github.com/quarkusio/quarkus/issues/3186

---
 .../runtime/context/ArcContextProvider.java   | 13 +++++------
 .../java/io/quarkus/arc/ManagedContext.java   |  8 +++++++
 .../java/io/quarkus/arc/RequestContext.java   | 22 +++++++++++++++++++
 3 files changed, 36 insertions(+), 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
index 006099f5332f5..80035019bbfa6 100644
--- 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
index b29c80c5469ea..9fcb01092f59b 100644
--- 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
index 614d1ac7a2d5a..61c80b08ee982 100644
--- 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;
+    }
 }