TOMEE-4703 / TOMEE-4704 capture thread context on the submitting thread

container/openejb-core/src/main/java/org/apache/openejb/core/ThreadContext.java

@@ -25,6 +25,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.CopyOnWriteArraySet;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Collectors;
@@ -151,13 +152,99 @@ public ThreadContext(final BeanContext beanContext, final Object primaryKey, fin
         this.currentOperation = operation;
     }
 
+    /**
+     * Copy constructor. Must be called on the thread that owns <code>that</code>, since a
+     * ThreadContext is confined to its thread. Use {@link #capture()} to pass a context to
+     * another thread.
+     */
     public ThreadContext(final ThreadContext that) {
         this.beanContext = that.beanContext;
         this.primaryKey = that.primaryKey;
-        this.data.putAll(that.data);
+        synchronized (that.data) {
+            this.data.putAll(that.data);
+        }
         this.oldClassLoader = that.oldClassLoader;
     }
 
+    /**
+     * Returns an immutable copy of the calling thread's context, which may be passed to other
+     * threads. Must be called on the thread that owns the context.
+     *
+     * @return the capture, or <code>null</code> if no context is entered on this thread
+     */
+    public static Capture capture() {
+        final ThreadContext current = threadStorage.get();
+        return current == null ? null : new Capture(current);
+    }
+
+    /**
+     * Immutable copy of the state a {@link ThreadContext} propagates: bean context, primary key and
+     * context data. Per-thread state such as the class loader to restore, the entered flag and the
+     * current operation is not included.
+     * <p>
+     * A capture may be applied to any number of threads, including concurrently.
+     * {@link #newThreadContext()} returns a separate mutable {@link ThreadContext} for each caller,
+     * since {@link ThreadContext#enter(ThreadContext)} modifies its argument and fails if that
+     * context was already entered.
+     */
+    public static final class Capture {
+
+        /**
+         * Context data tied to the invocation a capture is taken from, listed by class name to avoid
+         * a dependency on the types. It is not propagated:
+         * <ul>
+         *   <li><code>InvocationContext</code> is part of the interceptor chain the calling thread is
+         *       still in. It is single use, and {@link BaseContext#getContextData()} exposes its
+         *       unsynchronized map to application code.</li>
+         *   <li><code>DestroyContext</code> references the captured context and would keep it
+         *       reachable for the lifetime of the capture. A new one is created when the context is
+         *       entered on another thread.</li>
+         * </ul>
+         */
+        private static final Set<String> NON_PROPAGATED = Set.of(
+            "jakarta.interceptor.InvocationContext",
+            "org.apache.openejb.cdi.RequestScopedThreadContextListener$DestroyContext");
+
+        private final BeanContext beanContext;
+        private final Object primaryKey;
+        private final Map<Class, Object> data;
+
+        private Capture(final ThreadContext that) {
+            this.beanContext = that.beanContext;
+            this.primaryKey = that.primaryKey;
+
+            final Map<Class, Object> copy = new HashMap<>();
+            synchronized (that.data) {
+                for (final Map.Entry<Class, Object> entry : that.data.entrySet()) {
+                    if (NON_PROPAGATED.contains(entry.getKey().getName())) {
+                        continue;
+                    }
+                    copy.put(entry.getKey(), entry.getValue());
+                }
+            }
+            this.data = Collections.unmodifiableMap(copy);
+        }
+
+        /**
+         * @return a new mutable {@link ThreadContext} with the captured state, for the calling thread
+         *         to pass to {@link ThreadContext#enter(ThreadContext)}
+         */
+        public ThreadContext newThreadContext() {
+            final ThreadContext context = new ThreadContext(beanContext, primaryKey);
+            context.data.putAll(data);
+            return context;
+        }
+
+        @Override
+        public String toString() {
+            return "ThreadContext.Capture{" +
+                "beanContext=" + beanContext.getId() +
+                ", primaryKey=" + primaryKey +
+                ", data=" + dataToString(data) +
+                '}';
+        }
+    }
+
     public BeanContext getBeanContext() {
         return beanContext;
     }
@@ -226,8 +313,7 @@ public String toString() {
         return "ThreadContext{" +
             "beanContext=" + beanContext.getId() +
             ", primaryKey=" + primaryKey +
-            ", data(" + data.size() +
-            ")=" + dataToString(data) +
+            ", data=" + dataToString(data) +
             ", oldClassLoader=" + oldClassLoader +
             ", currentOperation=" + currentOperation +
             ", invokedInterface=" + invokedInterface +
@@ -236,10 +322,17 @@ public String toString() {
             '}';
     }
 
-    private String dataToString(final Map<Class, Object> data) {
-        return data.entrySet().stream()
+    private static String dataToString(final Map<Class, Object> data) {
+        // iterating a synchronized map requires its monitor, see TOMEE-4699. Copy under the monitor
+        // and format outside of it, so that application hashCode() implementations do not run while
+        // a lock that is taken on every invocation is held.
+        final Map<Class, Object> copy;
+        synchronized (data) {
+            copy = new HashMap<>(data);
+        }
+
+        return "(" + copy.size() + ")=" + copy.entrySet().stream()
                 .map(entry -> entry.getKey() + "=" + (entry.getValue() == null ? "null" : entry.getValue().hashCode()))
                 .collect(Collectors.joining(", "));
-
     }
 }

container/openejb-core/src/main/java/org/apache/openejb/threads/impl/ApplicationThreadContextProvider.java

@@ -39,7 +39,9 @@ public ThreadContextSnapshot currentContext(final Map<String, String> props) {
             return clearedContext(props);
         }
 
-        return new ApplicationThreadContextSnapshot(appContext.getId(), ThreadContext.getThreadContext());
+        // capture on the thread that owns the ThreadContext, see TOMEE-4699. A ThreadContext is
+        // confined to its thread; reading it from the thread running the task races with the owner.
+        return new ApplicationThreadContextSnapshot(appContext.getId(), ThreadContext.capture());
     }
 
     @Override
@@ -54,11 +56,11 @@ public String getThreadContextType() {
 
     public static class ApplicationThreadContextSnapshot implements ThreadContextSnapshot, Serializable {
         private final Object appId;
-        private final ThreadContext threadContext;
+        private final ThreadContext.Capture capturedThreadContext;
 
-        public ApplicationThreadContextSnapshot(final Object appId, final ThreadContext threadContext) {
+        public ApplicationThreadContextSnapshot(final Object appId, final ThreadContext.Capture capturedThreadContext) {
             this.appId = appId;
-            this.threadContext = threadContext;
+            this.capturedThreadContext = capturedThreadContext;
         }
 
         @Override
@@ -71,17 +73,20 @@ public ThreadContextRestorer begin() {
             final ClassLoader oldCl = Thread.currentThread().getContextClassLoader();
             Thread.currentThread().setContextClassLoader(appContext.getClassLoader());
 
-            // Don't touch ThreadContext if it is already correct or none was captured
-            boolean changeThreadContext = threadContext != null && threadContext != ThreadContext.getThreadContext();
-            ThreadContext oldThreadContext = changeThreadContext ? ThreadContext.enter(new ThreadContext(threadContext)) : null;
+            // leave the ThreadContext alone if nothing was captured, otherwise enter a new copy. This
+            // snapshot may be applied to any number of threads, including concurrently, and
+            // ThreadContext.enter modifies the context it is given.
+            final boolean changeThreadContext = capturedThreadContext != null;
+            final ThreadContext oldThreadContext =
+                    changeThreadContext ? ThreadContext.enter(capturedThreadContext.newThreadContext()) : null;
             return new ApplicationThreadContextRestorer(oldCl, oldThreadContext, changeThreadContext);
         }
 
         @Override
         public String toString() {
             return "ApplicationThreadContextSnapshot@" + System.identityHashCode(this) +
                     "{appId=" + appId +
-                    "{threadContext=" + threadContext +
+                    "{capturedThreadContext=" + capturedThreadContext +
                     '}';
         }
 
@@ -100,13 +105,16 @@ public ApplicationThreadContextRestorer(final ClassLoader oldClassLoader, final
 
         @Override
         public void endContext() throws IllegalStateException {
-            if (oldClassLoader != null) {
-                Thread.currentThread().setContextClassLoader(oldClassLoader);
-            }
-
+            // exit before restoring the class loader. ThreadContext.exit sets the loader to the value
+            // the context recorded on entry, which is the application class loader installed by
+            // begin(), so restoring afterwards leaves the thread with the loader it started with.
             if (exitThreadContext) {
                 ThreadContext.exit(oldThreadContext);
             }
+
+            if (oldClassLoader != null) {
+                Thread.currentThread().setContextClassLoader(oldClassLoader);
+            }
         }
 
         @Override

container/openejb-core/src/main/java/org/apache/openejb/threads/impl/ContextServiceImpl.java

@@ -26,6 +26,7 @@
 import org.apache.openejb.OpenEJBRuntimeException;
 import org.apache.openejb.resource.thread.ManagedExecutorServiceImplFactory;
 import org.apache.openejb.threads.future.CUCompletableFuture;
+import org.apache.openejb.threads.task.CURunnable;
 import org.apache.openejb.threads.task.CUTask;
 import org.apache.openejb.util.LogCategory;
 import org.apache.openejb.util.Logger;
@@ -149,7 +150,10 @@ public Object createContextualProxy(final Object instance, final Map<String, Str
 
     @Override
     public Executor currentContextExecutor() {
-        return command -> contextualRunnable(command).run();
+        // ContextService specifies "context that is captured from the thread that invokes
+        // currentContextExecutor", so capture here rather than in execute()
+        final Snapshot snapshot = snapshot(null);
+        return command -> new CURunnable(command, this, snapshot).run();
     }
 
     @Override

container/openejb-core/src/main/java/org/apache/openejb/threads/task/CURunnable.java

@@ -28,6 +28,11 @@ public CURunnable(final Runnable task) {
         super(task, ContextServiceImplFactory.newPropagateEverythingContextService());
         delegate = task;
     }
+    public CURunnable(final Runnable task, final ContextServiceImpl contextService, final ContextServiceImpl.Snapshot snapshot) {
+        super(task, contextService, snapshot);
+        this.delegate = task;
+    }
+
     public CURunnable(final Runnable task, final ContextServiceImpl contextService) {
         super(task, contextService);
         delegate = task;

container/openejb-core/src/main/java/org/apache/openejb/threads/task/CUTask.java

@@ -44,18 +44,24 @@ public static void addContainerListener(final ContainerListener cl) {
     protected final ContextServiceImpl contextService;
     private final ContextServiceImpl.Snapshot snapshot;
     private final Object[] containerListenerStates;
-    private final Context initialContext;
 
     public CUTask(final Object task, final ContextServiceImpl contextService) {
-        this(task, contextService, null);
+        this(task, contextService, (Map<String, String>) null);
     }
 
     public CUTask(final Object task, final ContextServiceImpl contextService, Map<String, String> props) {
+        this(task, contextService, contextService.snapshot(props));
+    }
+
+    /**
+     * Uses a snapshot captured earlier, on the thread the context is taken from.
+     * {@link ContextService#currentContextExecutor()} captures when the executor is created rather
+     * than when a task is submitted to it.
+     */
+    public CUTask(final Object task, final ContextServiceImpl contextService, final ContextServiceImpl.Snapshot snapshot) {
         super(task);
         this.contextService = contextService;
-
-        snapshot = contextService.snapshot(props);
-        initialContext = new Context();
+        this.snapshot = snapshot;
         if (CONTAINER_LISTENERS.length > 0) {
             containerListenerStates = new Object[CONTAINER_LISTENERS.length];
             for (int i = 0; i < CONTAINER_LISTENERS.length; i++) {
@@ -67,27 +73,36 @@ public CUTask(final Object task, final ContextServiceImpl contextService, Map<St
     }
 
     protected T invoke(final Callable<T> call) throws Exception {
-        initialContext.enter();
-        final Object[] oldStates;
-        if (CONTAINER_LISTENERS.length > 0) {
-            oldStates = new Object[CONTAINER_LISTENERS.length];
-            for (int i = 0; i < CONTAINER_LISTENERS.length; i++) {
-                oldStates[i] = CONTAINER_LISTENERS[i].onStart(containerListenerStates[i]);
-            }
-        } else {
-            oldStates = null;
-        }
+        // one per invocation rather than one per task: a contextual proxy runs its task more than
+        // once, and may run it on several threads at the same time
+        final Context invocationContext = new Context();
+
+        // read once: the array is replaced when a listener is registered, and the teardown below
+        // walks the listeners that were started
+        final ContainerListener[] listeners = CONTAINER_LISTENERS;
+        final Object[] oldStates = listeners.length > 0 ? new Object[listeners.length] : null;
+        int started = 0;
 
         ContextServiceImpl.State state = null;
+        boolean entered = false;
+        Throwable throwable = null;
 
-        if (contextService != null && snapshot != null) {
-            state = contextService.enter(snapshot);
-        }
+        // establishing the context can fail, see TOMEE-4699. Keep it inside the try so that the task
+        // listener is notified and the thread is cleaned up in that case as well.
+        try {
+            invocationContext.enter();
+            entered = true;
 
+            for (int i = 0; i < listeners.length; i++) {
+                oldStates[i] = listeners[i].onStart(containerListenerStates[i]);
+                started = i + 1;
+            }
 
-        Throwable throwable = null;
-        try {
-            taskStarting(future, executor, delegate); // do it in try to avoid issues if an exception is thrown
+            if (contextService != null && snapshot != null) {
+                state = contextService.enter(snapshot);
+            }
+
+            taskStarting(future, executor, delegate);
             return call.call();
         } catch (final Throwable t) {
             throwable = t;
@@ -97,15 +112,15 @@ protected T invoke(final Callable<T> call) throws Exception {
             try {
                 taskDone(future, executor, delegate, throwable);
             } finally {
-                if (CONTAINER_LISTENERS.length > 0) {
-                    for (int i = 0; i < CONTAINER_LISTENERS.length; i++) {
-                        CONTAINER_LISTENERS[i].onEnd(oldStates[i]);
-                    }
+                for (int i = 0; i < started; i++) {
+                    listeners[i].onEnd(oldStates[i]);
                 }
                 if (contextService != null && state != null) {
                     contextService.exit(state);
                 }
-                initialContext.exit();
+                if (entered) {
+                    invocationContext.exit();
+                }
             }
         }
     }