Block contextual operations from lazily enlisting CDI beans

dev/com.ibm.ws.cdi.mp.context/bnd.bnd

@@ -23,7 +23,7 @@ Require-Capability: osgi.ee;filter:="(&(osgi.ee=JavaSE)(version=1.8))"
 WS-TraceGroup: JCDI
 
 Private-Package:\
-  com.ibm.ws.cdi.mp.context
+  com.ibm.ws.cdi.mp.context.*
 
 -dsannotations: \
   com.ibm.ws.cdi.mp.context.WeldContextProvider

dev/com.ibm.ws.cdi.mp.context/resources/com/ibm/ws/cdi/mp/context/resources/CDI_MP_CONTEXT.nlsprops

@@ -0,0 +1,33 @@
+###############################################################################
+# Copyright (c) 2019 IBM Corporation and others.
+# All rights reserved. This program and the accompanying materials
+# are made available under the terms of the Eclipse Public License v1.0
+# which accompanies this distribution, and is available at
+# http://www.eclipse.org/legal/epl-v10.html
+#
+# Contributors:
+#     IBM Corporation - initial API and implementation
+###############################################################################
+#CMVCPATHNAME com.ibm.ws.concurrent/resources/com/ibm/ws/concurrent/resources/CWWKCMessages.nlsprops
+#ISMESSAGEFILE TRUE
+#NLS_ENCODING=UNICODE
+#
+#COMPONENTPREFIX CWWKC
+#COMPONENTNAMEFOR CWWKC MicroProfile Context Propagation
+#
+# NLS_MESSAGEFORMAT_VAR
+#
+#   Strings in this file which contain replacement variables are processed by the MessageFormat 
+#   class (single quote must be coded as 2 consecutive single quotes ''). Strings in this file 
+#   which do NOT contain replacement variables are NOT processed by the MessageFormat class 
+#   (single quote must be coded as one single quote '). 
+
+# This message bundle does not have its own message range, and instead reserves specific error messages from the 
+# com.ibm.ws.concurrent.mp.1.0 message bundle
+
+# do not translate: CDI, CompletionStage
+CWWKC1158.cannot.lazy.enlist.beans=CWWKC1158E: CDI beans used in contextual operations, such as a CompletionStage, must be accessed using a CDI bean operation, such as a public \
+method, before context for the operation is captured. The CDI beans that were not accessed before context was captured are: {0}  
+CWWKC1158.cannot.lazy.enlist.beans.explanation=When CDI context is captured, only CDI beans that have been accessed at that point in time have been created by the CDI \
+runtime. As a result, CDI beans that are first accessed in contextual operations cannot propagate their state correctly to subsequent contextual operations.
+CWWKC1158.cannot.lazy.enlist.beans.useraction=The application must be changed to invoke some CDI bean operation, such as a public method, before context for the operation is captured.

dev/com.ibm.ws.cdi.mp.context/src/com/ibm/ws/cdi/mp/context/ContextualInstanceSnapshot.java

@@ -73,9 +73,16 @@ public ContextualInstanceSnapshot(WeldManager manager) {
         this.conInstances = conInstances == null ? emptySet() : conInstances;
     }
 
+    public int getBeanCount() {
+        return reqInstances.size() +
+               sesInstances.size() +
+               conInstances.size();
+    }
+
     @Override
     public String toString() {
         StringBuilder sb = new StringBuilder(super.toString())
+                        .append(" beanCount=" + getBeanCount())
                         .append('\n')
                         .append("RequestScoped instances = ")
                         .append(this.reqInstances)

dev/com.ibm.ws.cdi.mp.context/src/com/ibm/ws/cdi/mp/context/WeldContextSnapshot.java

@@ -11,10 +11,13 @@
 package com.ibm.ws.cdi.mp.context;
 
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 
 import org.eclipse.microprofile.context.spi.ThreadContextController;
 import org.eclipse.microprofile.context.spi.ThreadContextSnapshot;
+import org.jboss.weld.context.api.ContextualInstance;
 import org.jboss.weld.context.bound.BoundConversationContext;
 import org.jboss.weld.context.bound.BoundLiteral;
 import org.jboss.weld.context.bound.BoundRequestContext;
@@ -31,11 +34,13 @@ public class WeldContextSnapshot implements ThreadContextSnapshot {
 
     private static final TraceComponent tc = Tr.register(WeldContextSnapshot.class);
 
+    private final boolean propagate;
     private final WeldManager manager;
     private final ContextualInstanceSnapshot contextToApply;
 
     public WeldContextSnapshot(boolean propagate, WeldManager manager) {
         this.manager = manager;
+        this.propagate = propagate;
         contextToApply = propagate ? new ContextualInstanceSnapshot(manager) : ContextualInstanceSnapshot.EMPTY_SNAPSHOT;
         if (TraceComponent.isAnyTracingEnabled() && tc.isDebugEnabled())
             Tr.debug(tc, "Snapshotted contextual instances to apply:", contextToApply);
@@ -84,7 +89,9 @@ public ThreadContextController begin() {
         }
 
         return () -> {
-            // This will run after a task has executed and will clean up the CDI context
+            // This will run after a task has executed and will restore the CDI context that was originally on the thread
+            ContextualInstanceSnapshot afterTaskContexts = propagate ? new ContextualInstanceSnapshot(manager) : null;
+
             if (existingContexts.reqCtx != null)
                 existingContexts.reqCtx.clearAndSet(existingContexts.reqInstances);
             else
@@ -99,6 +106,19 @@ public ThreadContextController begin() {
                 existingContexts.conCtx.clearAndSet(existingContexts.conInstances);
             else
                 conversationCtx.deactivate();
+
+            if (propagate && contextToApply.getBeanCount() != afterTaskContexts.getBeanCount()) {
+                Set<ContextualInstance<?>> lazilyRegisteredBeans = new HashSet<>();
+                lazilyRegisteredBeans.addAll(afterTaskContexts.reqInstances);
+                lazilyRegisteredBeans.addAll(afterTaskContexts.sesInstances);
+                lazilyRegisteredBeans.addAll(afterTaskContexts.conInstances);
+                lazilyRegisteredBeans.removeAll(contextToApply.reqInstances);
+                lazilyRegisteredBeans.removeAll(contextToApply.sesInstances);
+                lazilyRegisteredBeans.removeAll(contextToApply.conInstances);
+
+                Tr.error(tc, "CWWKC1158.cannot.lazy.enlist.beans", lazilyRegisteredBeans);
+                throw new IllegalStateException(Tr.formatMessage(tc, "CWWKC1158.cannot.lazy.enlist.beans", lazilyRegisteredBeans));
+            }
         };
     }
 }

dev/com.ibm.ws.cdi.mp.context/src/com/ibm/ws/cdi/mp/context/package-info.java

@@ -10,7 +10,7 @@
  *******************************************************************************/
 
 @Version("1.0.0")
-@TraceOptions(traceGroup = "JCDI")
+@TraceOptions(traceGroup = "JCDI", messageBundle = "com.ibm.ws.cdi.mp.context.resources.CDI_MP_CONTEXT")
 package com.ibm.ws.cdi.mp.context;
 
 import org.osgi.annotation.versioning.Version;

dev/com.ibm.ws.concurrent.mp.1.0/resources/com/ibm/ws/concurrent/mp/resources/CWWKCMessages.nlsprops

@@ -49,3 +49,5 @@ CWWKC1156.not.supported.useraction=Update the application to use the method that
 CWWKC1157.cannot.propagate.tx=CWWKC1157E: Propagating transactions to contextual actions and tasks is not supported.
 CWWKC1157.cannot.propagate.tx.explanation=A ManagedExecutor or ThreadContext that is configured to propagate transaction contexts can propagate empty transaction contexts only. Therefore, you cannot create contextual actions and tasks within a transaction.
 CWWKC1157.cannot.propagate.tx.useraction=Create the contextual action or task outside of a transaction. Alternatively, configure the ManagedExecutor or ThreadContext to not propagate transaction contexts.
+
+# CWWKC1158E used by com.ibm.ws.cdi.mp.context bundle

dev/com.ibm.ws.concurrent.mp.1.0/src/com/ibm/ws/concurrent/mp/ContextualRunnable.java

@@ -68,6 +68,12 @@ public void run() {
             try {
                 if (contextApplied != null)
                     threadContextDescriptor.taskStopping(contextApplied);
+            } catch (RuntimeException x) {
+                // prioritize surfacing an actual task failure over a failure clearing context
+                if (failure == null) {
+                    failure = x;
+                    throw x;
+                }
             } finally {
                 if (completableFuture != null)
                     if (failure == null)

dev/com.ibm.ws.concurrent.mp_fat/fat/src/com/ibm/ws/concurrent/mp/fat/MPConcurrentCDITest.java

@@ -51,6 +51,6 @@ public static void setUp() throws Exception {
 
     @AfterClass
     public static void tearDown() throws Exception {
-        server.stopServer();
+        server.stopServer("CWWKC1158E"); // expected by testCDIContextPropagationBeanFirstUsedInCompletionStage
     }
 }

dev/com.ibm.ws.concurrent.mp_fat/test-applications/MPConcurrentCDIApp/src/concurrent/mp/fat/cdi/web/MPConcurrentCDITestServlet.java

@@ -12,6 +12,8 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 
 import java.lang.annotation.ElementType;
 import java.lang.annotation.Retention;
@@ -44,6 +46,7 @@
 import org.test.context.location.TestContextTypes;
 
 import componenttest.annotation.AllowedFFDC;
+import componenttest.annotation.ExpectedFFDC;
 import componenttest.app.FATServlet;
 
 @SuppressWarnings("serial")
@@ -209,21 +212,30 @@ public void testCDIContextPropagationAcrossMultipleStagesBeanAccessedFirst() thr
     }
 
     /**
-     * Propagate CDI context across completion stages. Access (and modify) a request scoped bean
-     * within the stages, but not before.
+     * Propagate CDI context to a completion stage. Access (and modify) a request scoped bean
+     * within the stage, but not before the context snapshot is taken. Lazily enlisting CDI beans
+     * should be rejected.
      */
     @Test
-    public void testCDIContextPropagationAcrossMultipleStagesBeanFirstUsedInCompletionStage() throws Exception {
+    @ExpectedFFDC("java.lang.IllegalStateException")
+    public void testCDIContextPropagationBeanFirstUsedInCompletionStage() throws Exception {
+        sessionBean.setState("foo");
+
         CompletableFuture<Void> cf = appCDIExecutor.runAsync(() -> {
-            requestBean.setState(requestBean.getState() + ",A");
-        }).thenRunAsync(() -> {
-            requestBean.setState(requestBean.getState() + ",B");
-        }).thenRunAsync(() -> {
-            requestBean.setState(requestBean.getState() + ",C");
+            requestBean.setState("A");
         });
-        cf.join();
-
-        assertEquals("UNINITIALIZED", requestBean.getState()); // TODO change to: UNINITIALIZED,A,B,C ?
+        try {
+            cf.join();
+            fail("Should not be able to lazily enlist a CDI bean when CDI context is propagated.");
+        } catch (CompletionException x) {
+            if (x.getCause() instanceof IllegalStateException && x.getCause().getMessage().contains("CWWKC1158E")) {
+                System.out.println("Caught expected IllegalStateException");
+            } else {
+                throw x;
+            }
+        }
+        assertTrue("CompletableFuture should have been marked done", cf.isDone());
+        assertTrue("CompletableFuture should have been marked as completex exceptionally", cf.isCompletedExceptionally());
     }
 
     /**