Micro performance improvements

application/src/main/java/io/smallrye/context/application/context/propagation/ApplicationContextProvider.java

@@ -5,6 +5,7 @@
 import java.util.Map;
 
 import org.eclipse.microprofile.context.ThreadContext;
+import org.eclipse.microprofile.context.spi.ThreadContextController;
 import org.eclipse.microprofile.context.spi.ThreadContextProvider;
 import org.eclipse.microprofile.context.spi.ThreadContextSnapshot;
 
@@ -35,28 +36,40 @@ private static ClassLoader calculateSystemClassLoader() {
     @Override
     public ThreadContextSnapshot currentContext(Map<String, String> props) {
         ClassLoader capturedTCCL = Thread.currentThread().getContextClassLoader();
-        return () -> {
-            ClassLoader movedTCCL = Thread.currentThread().getContextClassLoader();
-            if (capturedTCCL != movedTCCL)
-                Thread.currentThread().setContextClassLoader(capturedTCCL);
-            return () -> {
-                if (Thread.currentThread().getContextClassLoader() != movedTCCL)
-                    Thread.currentThread().setContextClassLoader(movedTCCL);
-            };
+        return new ThreadContextSnapshot() {
+            @Override
+            public ThreadContextController begin() {
+                ClassLoader movedTCCL = Thread.currentThread().getContextClassLoader();
+                if (capturedTCCL != movedTCCL)
+                    Thread.currentThread().setContextClassLoader(capturedTCCL);
+                return new ThreadContextController() {
+                    @Override
+                    public void endContext() throws IllegalStateException {
+                        if (Thread.currentThread().getContextClassLoader() != movedTCCL)
+                            Thread.currentThread().setContextClassLoader(movedTCCL);
+                    }
+                };
+            }
         };
     }
 
     @Override
     public ThreadContextSnapshot clearedContext(Map<String, String> props) {
         ClassLoader capturedTCCL = SYSTEM_CL;
-        return () -> {
-            ClassLoader movedTCCL = Thread.currentThread().getContextClassLoader();
-            if (capturedTCCL != movedTCCL)
-                Thread.currentThread().setContextClassLoader(capturedTCCL);
-            return () -> {
-                if (Thread.currentThread().getContextClassLoader() != movedTCCL)
-                    Thread.currentThread().setContextClassLoader(movedTCCL);
-            };
+        return new ThreadContextSnapshot() {
+            @Override
+            public ThreadContextController begin() {
+                ClassLoader movedTCCL = Thread.currentThread().getContextClassLoader();
+                if (capturedTCCL != movedTCCL)
+                    Thread.currentThread().setContextClassLoader(capturedTCCL);
+                return new ThreadContextController() {
+                    @Override
+                    public void endContext() throws IllegalStateException {
+                        if (Thread.currentThread().getContextClassLoader() != movedTCCL)
+                            Thread.currentThread().setContextClassLoader(movedTCCL);
+                    }
+                };
+            }
         };
     }
 

core/src/main/java/io/smallrye/context/SmallRyeContextManager.java

@@ -69,8 +69,7 @@ public String[] getAllProviderTypes() {
     }
 
     public CapturedContextState captureContext(SmallRyeThreadContext context) {
-        Map<String, String> props = Collections.emptyMap();
-        return new CapturedContextState(context, context.getPlan(), props);
+        return new CapturedContextState(context, context.getPlan());
     }
 
     // for tests

core/src/main/java/io/smallrye/context/SmallRyeThreadContext.java

@@ -23,6 +23,12 @@
 public class SmallRyeThreadContext implements ThreadContext {
 
     private final static ThreadLocal<SmallRyeThreadContext> currentThreadContext = new ThreadLocal<>();
+    private final static CleanAutoCloseable NULL_THREAD_STATE = new CleanAutoCloseable() {
+        @Override
+        public void close() {
+            currentThreadContext.remove();
+        }
+    };
 
     /**
      * Updates the current @{link SmallRyeThreadContext} in use by the current thread, and returns an
@@ -34,9 +40,18 @@ public class SmallRyeThreadContext implements ThreadContext {
     public static CleanAutoCloseable withThreadContext(SmallRyeThreadContext threadContext) {
         SmallRyeThreadContext oldValue = currentThreadContext.get();
         currentThreadContext.set(threadContext);
-        return () -> {
-            currentThreadContext.set(oldValue);
-        };
+        if (oldValue == null) {
+            //For restoring null values we can optimise this a little:
+            return NULL_THREAD_STATE;
+        } else {
+            return new CleanAutoCloseable() {
+                @Override
+                public void close() {
+                    currentThreadContext.set(oldValue);
+                }
+            };
+        }
+
     }
 
     /**
@@ -47,8 +62,16 @@ public static CleanAutoCloseable withThreadContext(SmallRyeThreadContext threadC
      * @param f the @{link Runnable} to invoke
      */
     public static void withThreadContext(SmallRyeThreadContext threadContext, Runnable f) {
-        try (CleanAutoCloseable foo = withThreadContext(threadContext)) {
+        final SmallRyeThreadContext oldValue = currentThreadContext.get();
+        currentThreadContext.set(threadContext);
+        try {
             f.run();
+        } finally {
+            if (oldValue == null) {
+                currentThreadContext.remove();
+            } else {
+                currentThreadContext.set(oldValue);
+            }
         }
     }
 

core/src/main/java/io/smallrye/context/impl/ActiveContextState.java

@@ -1,6 +1,5 @@
 package io.smallrye.context.impl;
 
-import java.util.ArrayList;
 import java.util.List;
 
 import org.eclipse.microprofile.context.spi.ThreadContextController;
@@ -11,21 +10,22 @@
 
 public class ActiveContextState implements AutoCloseable {
 
-    private List<ThreadContextController> activeContext;
+    private ThreadContextController[] activeContext;
     private CleanAutoCloseable activeThreadContext;
 
     public ActiveContextState(SmallRyeThreadContext threadContext, List<ThreadContextSnapshot> threadContextSnapshots) {
-        activeContext = new ArrayList<>(threadContextSnapshots.size());
+        activeContext = new ThreadContextController[threadContextSnapshots.size()];
+        int i = 0;
         for (ThreadContextSnapshot threadContextSnapshot : threadContextSnapshots) {
-            activeContext.add(threadContextSnapshot.begin());
+            activeContext[i++] = threadContextSnapshot.begin();
         }
         activeThreadContext = SmallRyeThreadContext.withThreadContext(threadContext);
     }
 
     public void close() {
         // restore in reverse order
-        for (int i = activeContext.size() - 1; i >= 0; i--) {
-            activeContext.get(i).endContext();
+        for (int i = activeContext.length - 1; i >= 0; i--) {
+            activeContext[i].endContext();
         }
         activeThreadContext.close();
     }

core/src/main/java/io/smallrye/context/impl/CapturedContextState.java

@@ -1,34 +1,19 @@
 package io.smallrye.context.impl;
 
-import java.util.LinkedList;
 import java.util.List;
-import java.util.Map;
 
-import org.eclipse.microprofile.context.spi.ThreadContextProvider;
 import org.eclipse.microprofile.context.spi.ThreadContextSnapshot;
 
 import io.smallrye.context.SmallRyeThreadContext;
 
 public class CapturedContextState {
 
-    private List<ThreadContextSnapshot> threadContextSnapshots = new LinkedList<>();
+    private List<ThreadContextSnapshot> threadContextSnapshots;
     private SmallRyeThreadContext threadContext;
 
-    public CapturedContextState(SmallRyeThreadContext threadContext, ThreadContextProviderPlan plan,
-            Map<String, String> props) {
+    public CapturedContextState(SmallRyeThreadContext threadContext, ThreadContextProviderPlan plan) {
         this.threadContext = threadContext;
-        for (ThreadContextProvider provider : plan.propagatedProviders) {
-            ThreadContextSnapshot snapshot = provider.currentContext(props);
-            if (snapshot != null) {
-                threadContextSnapshots.add(snapshot);
-            }
-        }
-        for (ThreadContextProvider provider : plan.clearedProviders) {
-            ThreadContextSnapshot snapshot = provider.clearedContext(props);
-            if (snapshot != null) {
-                threadContextSnapshots.add(snapshot);
-            }
-        }
+        this.threadContextSnapshots = plan.takeThreadContextSnapshots();
     }
 
     public ActiveContextState begin() {

core/src/main/java/io/smallrye/context/impl/ThreadContextProviderPlan.java

@@ -1,19 +1,60 @@
 package io.smallrye.context.impl;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 import org.eclipse.microprofile.context.spi.ThreadContextProvider;
+import org.eclipse.microprofile.context.spi.ThreadContextSnapshot;
 
 public class ThreadContextProviderPlan {
 
     public final Set<ThreadContextProvider> propagatedProviders;
     public final Set<ThreadContextProvider> unchangedProviders;
     public final Set<ThreadContextProvider> clearedProviders;
 
+    /**
+     * Following 3 fields are meant to optimise method #takeThreadContextSnapshots
+     * which is extremely hot at runtime, at cost of little extra memory for this plan.
+     */
+    private final int snapshotInitialSize;
+    private final ThreadContextProvider[] propagatedProvidersFastIterable;
+    private final ThreadContextProvider[] clearedProvidersFastIterable;
+
     public ThreadContextProviderPlan(Set<ThreadContextProvider> propagatedSet, Set<ThreadContextProvider> unchangedSet,
             Set<ThreadContextProvider> clearedSet) {
-        this.propagatedProviders = propagatedSet;
-        this.unchangedProviders = unchangedSet;
-        this.clearedProviders = clearedSet;
+        this.propagatedProviders = Collections.unmodifiableSet(propagatedSet);
+        this.unchangedProviders = Collections.unmodifiableSet(unchangedSet);
+        this.clearedProviders = Collections.unmodifiableSet(clearedSet);
+        this.snapshotInitialSize = propagatedProviders.size() + clearedProviders.size();
+        this.propagatedProvidersFastIterable = propagatedProviders.toArray(new ThreadContextProvider[0]);
+        this.clearedProvidersFastIterable = clearedProviders.toArray(new ThreadContextProvider[0]);
+    }
+
+    /**
+     * This helps to optimise construction of CapturedContextState
+     * without exposing too many implementation details.
+     * Only useful for snapshots with an empty property set.
+     * 
+     * @return a list of snapshots
+     */
+    public List<ThreadContextSnapshot> takeThreadContextSnapshots() {
+        List<ThreadContextSnapshot> threadContextSnapshots = new ArrayList<>(snapshotInitialSize);
+        final Map<String, String> props = Collections.emptyMap();
+        for (ThreadContextProvider provider : propagatedProvidersFastIterable) {
+            ThreadContextSnapshot snapshot = provider.currentContext(props);
+            if (snapshot != null) {
+                threadContextSnapshots.add(snapshot);
+            }
+        }
+        for (ThreadContextProvider provider : clearedProvidersFastIterable) {
+            ThreadContextSnapshot snapshot = provider.clearedContext(props);
+            if (snapshot != null) {
+                threadContextSnapshots.add(snapshot);
+            }
+        }
+        return threadContextSnapshots;
     }
 }