Merge pull request #23436 from manovotn/contextStateInvalidation

Allow to invalidate ContextState of shareable CDI context and make use of it in ArcContextProvider

extensions/arc/runtime/src/main/java/io/quarkus/arc/runtime/context/ArcContextProvider.java

@@ -108,7 +108,7 @@ public ThreadContextController begin() {
                 return new ThreadContextController() {
                     @Override
                     public void endContext() throws IllegalStateException {
-                        requestContext.activate(toRestore);
+                        requestContext.activate(toRestore.isValid() ? toRestore : null);
                     }
                 };
             } else {
@@ -141,11 +141,13 @@ public ThreadContextController begin() {
             if (toRestore != null) {
                 // context active, store current state, feed it new one and restore state afterwards
                 // it is not necessary to deactivate the context first - just overwrite the previous state
-                requestContext.activate(state);
+                // if the context state is invalid (i.e. the context was destroyed by Arc), we instead create new state
+                requestContext.activate(state.isValid() ? state : null);
                 return new RestoreContextController(requestContext, toRestore);
             } else {
                 // context not active, activate and pass it new instance, deactivate afterwards
-                requestContext.activate(state);
+                // if the context state is invalid (i.e. the context was destroyed by Arc), we instead create new state
+                requestContext.activate(state.isValid() ? state : null);
                 return requestContext::deactivate;
             }
         }
@@ -165,7 +167,7 @@ private static final class RestoreContextController implements ThreadContextCont
         @Override
         public void endContext() throws IllegalStateException {
             // it is not necessary to deactivate the context first - just overwrite the previous state
-            requestContext.activate(stateToRestore);
+            requestContext.activate(stateToRestore.isValid() ? stateToRestore : null);
         }
 
     }

independent-projects/arc/runtime/src/main/java/io/quarkus/arc/InjectableContext.java

@@ -95,5 +95,15 @@ interface ContextState {
          */
         Map<InjectableBean<?>, Object> getContextualInstances();
 
+        /**
+         * Context state is typically invalidated once the context to which is belongs is being destroyed.
+         * This flag is then used by context propagation to indicate that the given state shouldn't be reused anymore.
+         * 
+         * @return true if the context state is valid, false otherwise
+         */
+        default boolean isValid() {
+            return true;
+        }
+
     }
 }

independent-projects/arc/runtime/src/main/java/io/quarkus/arc/ManagedContext.java

@@ -17,6 +17,8 @@ default void activate() {
 
     /**
      * Activate the context.
+     * If invoked with {@code null} parameter, a fresh {@link io.quarkus.arc.InjectableContext.ContextState} is
+     * automatically created.
      * 
      * @param initialState The initial state, may be {@code null}
      */

independent-projects/arc/runtime/src/main/java/io/quarkus/arc/impl/RequestContext.java

@@ -11,6 +11,7 @@
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import javax.enterprise.context.BeforeDestroyed;
@@ -33,7 +34,7 @@ class RequestContext implements ManagedContext {
     private static final Logger LOGGER = Logger.getLogger(RequestContext.class.getPackage().getName());
 
     // It's a normal scope so there may be no more than one mapped instance per contextual type per thread
-    private final ThreadLocal<ConcurrentMap<Contextual<?>, ContextInstanceHandle<?>>> currentContext = new ThreadLocal<>();
+    private final ThreadLocal<RequestContextState> currentContext = new ThreadLocal<>();
 
     private final LazyValue<Notifier<Object>> initializedNotifier;
     private final LazyValue<Notifier<Object>> beforeDestroyedNotifier;
@@ -55,18 +56,19 @@ public Class<? extends Annotation> getScope() {
     public <T> T getIfActive(Contextual<T> contextual, Function<Contextual<T>, CreationalContext<T>> creationalContextFun) {
         Objects.requireNonNull(contextual, "Contextual must not be null");
         Objects.requireNonNull(creationalContextFun, "CreationalContext supplier must not be null");
-        Map<Contextual<?>, ContextInstanceHandle<?>> ctx = currentContext.get();
-        if (ctx == null) {
+        RequestContextState ctxState = currentContext.get();
+        if (ctxState == null) {
             // Thread local not set - context is not active!
             return null;
         }
-        ContextInstanceHandle<T> instance = (ContextInstanceHandle<T>) ctx.get(contextual);
+        Map<Contextual<?>, ContextInstanceHandle<?>> ctxMap = currentContext.get().value;
+        ContextInstanceHandle<T> instance = (ContextInstanceHandle<T>) ctxMap.get(contextual);
         if (instance == null) {
             CreationalContext<T> creationalContext = creationalContextFun.apply(contextual);
             // Bean instance does not exist - create one if we have CreationalContext
             instance = new ContextInstanceHandleImpl<T>((InjectableBean<T>) contextual,
                     contextual.create(creationalContext), creationalContext);
-            ctx.put(contextual, instance);
+            ctxMap.put(contextual, instance);
         }
         return instance.get();
     }
@@ -86,7 +88,7 @@ public <T> T get(Contextual<T> contextual, CreationalContext<T> creationalContex
     @Override
     public <T> T get(Contextual<T> contextual) {
         Objects.requireNonNull(contextual, "Contextual must not be null");
-        Map<Contextual<?>, ContextInstanceHandle<?>> ctx = currentContext.get();
+        Map<Contextual<?>, ContextInstanceHandle<?>> ctx = currentContext.get().value;
         if (ctx == null) {
             // Thread local not set - context is not active!
             throw new ContextNotActiveException();
@@ -102,7 +104,7 @@ public boolean isActive() {
 
     @Override
     public void destroy(Contextual<?> contextual) {
-        Map<Contextual<?>, ContextInstanceHandle<?>> ctx = currentContext.get();
+        Map<Contextual<?>, ContextInstanceHandle<?>> ctx = currentContext.get().value;
         if (ctx == null) {
             // Thread local not set - context is not active!
             throw new ContextNotActiveException();
@@ -116,12 +118,12 @@ public void destroy(Contextual<?> contextual) {
     @Override
     public void activate(ContextState initialState) {
         if (initialState == null) {
-            currentContext.set(new ConcurrentHashMap<>());
+            currentContext.set(new RequestContextState(new ConcurrentHashMap<>()));
             // Fire an event with qualifier @Initialized(RequestScoped.class) if there are any observers for it
             fireIfNotEmpty(initializedNotifier);
         } else {
             if (initialState instanceof RequestContextState) {
-                currentContext.set(((RequestContextState) initialState).value);
+                currentContext.set(((RequestContextState) initialState));
             } else {
                 throw new IllegalArgumentException("Invalid initial state: " + initialState.getClass().getName());
             }
@@ -130,20 +132,20 @@ public void activate(ContextState initialState) {
 
     @Override
     public ContextState getState() {
-        ConcurrentMap<Contextual<?>, ContextInstanceHandle<?>> ctx = currentContext.get();
+        RequestContextState ctx = currentContext.get();
         if (ctx == null) {
             // Thread local not set - context is not active!
             throw new ContextNotActiveException();
         }
-        return new RequestContextState(ctx);
+        return ctx;
     }
 
     public ContextState getStateIfActive() {
-        ConcurrentMap<Contextual<?>, ContextInstanceHandle<?>> ctx = currentContext.get();
+        RequestContextState ctx = currentContext.get();
         if (ctx == null) {
             return null;
         }
-        return new RequestContextState(ctx);
+        return ctx;
     }
 
     @Override
@@ -158,8 +160,14 @@ public void destroy() {
 
     @Override
     public void destroy(ContextState state) {
+        if (state == null) {
+            // nothing to destroy
+            return;
+        }
         if (state instanceof RequestContextState) {
-            destroy(((RequestContextState) state).value);
+            RequestContextState reqState = ((RequestContextState) state);
+            reqState.isValid.set(false);
+            destroy(reqState.value);
         } else {
             throw new IllegalArgumentException("Invalid state: " + state.getClass().getName());
         }
@@ -223,9 +231,11 @@ private static Notifier<Object> createDestroyedNotifier() {
     static class RequestContextState implements ContextState {
 
         private final ConcurrentMap<Contextual<?>, ContextInstanceHandle<?>> value;
+        private final AtomicBoolean isValid;
 
         RequestContextState(ConcurrentMap<Contextual<?>, ContextInstanceHandle<?>> value) {
             this.value = value;
+            this.isValid = new AtomicBoolean(true);
         }
 
         @Override
@@ -234,6 +244,11 @@ public Map<InjectableBean<?>, Object> getContextualInstances() {
                     .collect(Collectors.toUnmodifiableMap(ContextInstanceHandle::getBean, ContextInstanceHandle::get));
         }
 
+        @Override
+        public boolean isValid() {
+            return isValid.get();
+        }
+
     }
 
 }