Merge pull request #29 from rstoyanchev/shortcuts

marcingrzejszczak · web-flow · commit 6d2424c887ef · 2022-07-22T12:58:13.000+02:00
Add shortcuts to propagate from Reactor context to ThreadLocal values
diff --git a/context-propagation/src/main/java/io/micrometer/context/ContextAccessor.java b/context-propagation/src/main/java/io/micrometer/context/ContextAccessor.java
@@ -47,6 +47,16 @@ public interface ContextAccessor<READ, WRITE> {
      */
     void readValues(READ sourceContext, Predicate<Object> keyPredicate, Map<Object, Object> readValues);
 
+    /**
+     * Read a single value from the source context.
+     * @param sourceContext the context to read from; the context type should be
+     * checked with {@link #canReadFrom(Class)} before this method is called
+     * @param key the key to use to look up the context value
+     * @return the value, if any
+     */
+    @Nullable
+    <T> T readValue(READ sourceContext, Object key);
+
     /**
      * Whether this accessor can restore values to the given type of context.
      * @param contextType the type of external context
diff --git a/context-propagation/src/main/java/io/micrometer/context/ContextSnapshot.java b/context-propagation/src/main/java/io/micrometer/context/ContextSnapshot.java
@@ -155,7 +155,7 @@ static ContextSnapshot captureUsing(Predicate<Object> keyPredicate, Object... co
     /**
      * Variant of {@link #captureUsing(Predicate, Object...)} with a specific
      * {@link ContextRegistry} instead of the global instance.
-     * @param contextRegistry the {@code ContextRegistry} instance to use
+     * @param contextRegistry the registry with the accessors to use
      * @param keyPredicate predicate for context value keys
      * @param contexts one more context objects to extract values from
      * @return a snapshot with saved context values
@@ -166,6 +166,35 @@ static ContextSnapshot captureUsing(
         return DefaultContextSnapshot.capture(contextRegistry, keyPredicate, contexts);
     }
 
+    /**
+     * Read the values specified by from the given source context, and if found,
+     * use them to set {@link ThreadLocal} values. Essentially, a shortcut that
+     * bypasses the need to create of {@link ContextSnapshot} first via
+     * {@link #capture(Object...)}, followed by {@link #setThreadLocalValues()}.
+     * @param sourceContext the source context to read values from
+     * @param keys the keys of the values to read
+     * @return an object that can be used to reset {@link ThreadLocal} values
+     * at the end of the context scope, either removing them or restoring their
+     * previous values, if any.
+     */
+    static Scope setThreadLocalsFrom(Object sourceContext, String... keys) {
+        return setThreadLocalsFrom(sourceContext, ContextRegistry.getInstance(), keys);
+    }
+
+    /**
+     * Variant of {@link #setThreadLocalsFrom(Object, String...)} with a specific
+     * {@link ContextRegistry} instead of the global instance.
+     * @param sourceContext the source context to read values from
+     * @param contextRegistry the registry with the accessors to use
+     * @param keys the keys of the values to read
+     * @return an object that can be used to reset {@link ThreadLocal} values
+     * at the end of the context scope, either removing them or restoring their
+     * previous values, if any.
+     */
+    static Scope setThreadLocalsFrom(Object sourceContext, ContextRegistry contextRegistry, String... keys) {
+        return DefaultContextSnapshot.setThreadLocalsFrom(sourceContext, contextRegistry, keys);
+    }
+
 
     /**
      * An object to use to reset {@link ThreadLocal} values at the end of a
diff --git a/context-propagation/src/main/java/io/micrometer/context/DefaultContextSnapshot.java b/context-propagation/src/main/java/io/micrometer/context/DefaultContextSnapshot.java
@@ -16,9 +16,7 @@
 package io.micrometer.context;
 
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.Map;
-import java.util.Set;
 import java.util.function.Predicate;
 
 /**
@@ -31,11 +29,11 @@ final class DefaultContextSnapshot extends HashMap<Object, Object> implements Co
 
     private static final ContextSnapshot emptyContextSnapshot = new DefaultContextSnapshot(new ContextRegistry());
 
-    private final ContextRegistry accessorRegistry;
+    private final ContextRegistry contextRegistry;
 
 
-    DefaultContextSnapshot(ContextRegistry accessorRegistry) {
-        this.accessorRegistry = accessorRegistry;
+    DefaultContextSnapshot(ContextRegistry contextRegistry) {
+        this.contextRegistry = contextRegistry;
     }
 
 
@@ -48,7 +46,7 @@ public <C> C updateContext(C context) {
     public <C> C updateContext(C context, Predicate<Object> keyPredicate) {
         if (!isEmpty()) {
             Map<Object, Object> valuesToWrite = new HashMap<>();
-            forEach((key, value) -> {
+            this.forEach((key, value) -> {
                 if (keyPredicate.test(key)) {
                     valuesToWrite.put(key, value);
                 }
@@ -61,7 +59,7 @@ public <C> C updateContext(C context, Predicate<Object> keyPredicate) {
     @SuppressWarnings("unchecked")
     private <C> C updateContextInternal(C context, Map<Object, Object> valueContainer) {
         if (!isEmpty()) {
-            ContextAccessor<?, ?> accessor = this.accessorRegistry.getContextAccessorForWrite(context);
+            ContextAccessor<?, ?> accessor = this.contextRegistry.getContextAccessorForWrite(context);
             context = ((ContextAccessor<?, C>) accessor).writeValues(valueContainer, context);
         }
         return context;
@@ -74,27 +72,41 @@ public Scope setThreadLocalValues() {
 
     @Override
     public Scope setThreadLocalValues(Predicate<Object> keyPredicate) {
-        Set<Object> keys = null;
         Map<Object, Object> previousValues = null;
-        for (ThreadLocalAccessor<?> accessor : this.accessorRegistry.getThreadLocalAccessors()) {
+        for (ThreadLocalAccessor<?> accessor : this.contextRegistry.getThreadLocalAccessors()) {
             Object key = accessor.key();
-            if (keyPredicate.test(key) && containsKey(key)) {
-                keys = (keys != null ? keys : new HashSet<>());
-                keys.add(key);
-
-                Object previousValue = accessor.getValue();
-                previousValues = (previousValues != null ? previousValues : new HashMap<>());
-                previousValues.put(key, previousValue);
-
-                setThreadLocalValue(key, accessor);
+            if (keyPredicate.test(key) && this.containsKey(key)) {
+                previousValues = setThreadLocal(key, get(key), accessor, previousValues);
             }
         }
-        return (keys != null ? new DefaultScope(keys, previousValues) : () -> { });
+        return DefaultScope.from(previousValues, this.contextRegistry);
     }
 
     @SuppressWarnings("unchecked")
-    private <V> void setThreadLocalValue(Object key, ThreadLocalAccessor<?> accessor) {
-        ((ThreadLocalAccessor<V>) accessor).setValue((V) get(key));
+    private static <V> Map<Object, Object> setThreadLocal(
+            Object key, V value, ThreadLocalAccessor<?> accessor, @Nullable Map<Object, Object> previousValues) {
+
+        previousValues = (previousValues != null ? previousValues : new HashMap<>());
+        previousValues.put(key, accessor.getValue());
+        ((ThreadLocalAccessor<V>) accessor).setValue(value);
+        return previousValues;
+    }
+
+    @SuppressWarnings("unchecked")
+    static <C> Scope setThreadLocalsFrom(Object context, ContextRegistry registry, String... keys) {
+        ContextAccessor<?, ?> contextAccessor = registry.getContextAccessorForRead(context);
+        Map<Object, Object> previousValues = null;
+        for (String key : keys) {
+            Object value = ((ContextAccessor<C, ?>) contextAccessor).readValue((C) context, key);
+            if (value != null) {
+                for (ThreadLocalAccessor<?> threadLocalAccessor : registry.getThreadLocalAccessors()) {
+                    if (key.equals(threadLocalAccessor.key())) {
+                        previousValues = setThreadLocal(key, value, threadLocalAccessor, previousValues);
+                    }
+                }
+            }
+        }
+        return DefaultScope.from(previousValues, registry);
     }
 
     @SuppressWarnings("unchecked")
@@ -128,21 +140,21 @@ public String toString() {
     /**
      * Default implementation of {@link Scope}.
      */
-    private class DefaultScope implements Scope {
-
-        private final Set<Object> keys;
+    private static class DefaultScope implements Scope {
 
         private final Map<Object, Object> previousValues;
 
-        private DefaultScope(Set<Object> keys, Map<Object, Object> previousValues) {
-            this.keys = keys;
+        private final ContextRegistry contextRegistry;
+
+        private DefaultScope(Map<Object, Object> previousValues, ContextRegistry contextRegistry) {
             this.previousValues = previousValues;
+            this.contextRegistry = contextRegistry;
         }
 
         @Override
         public void close() {
-            for (ThreadLocalAccessor<?> accessor : accessorRegistry.getThreadLocalAccessors()) {
-                if (this.keys.contains(accessor.key())) {
+            for (ThreadLocalAccessor<?> accessor : this.contextRegistry.getThreadLocalAccessors()) {
+                if (this.previousValues.containsKey(accessor.key())) {
                     Object previousValue = this.previousValues.get(accessor.key());
                     resetThreadLocalValue(accessor, previousValue);
                 }
@@ -158,6 +170,11 @@ private <V> void resetThreadLocalValue(ThreadLocalAccessor<?> accessor, @Nullabl
                 accessor.reset();
             }
         }
+
+        public static Scope from(@Nullable Map<Object, Object> previousValues, ContextRegistry registry) {
+            return (previousValues != null ? new DefaultScope(previousValues, registry) : () -> { });
+        }
+
     }
 
 }
diff --git a/context-propagation/src/test/java/io/micrometer/context/DefaultContextSnapshotTests.java b/context-propagation/src/test/java/io/micrometer/context/DefaultContextSnapshotTests.java
@@ -16,6 +16,10 @@
 package io.micrometer.context;
 
 
+import java.util.Collections;
+import java.util.Map;
+
+import io.micrometer.context.ContextSnapshot.Scope;
 import org.junit.jupiter.api.Test;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -35,31 +39,32 @@ public class DefaultContextSnapshotTests {
     void should_propagate_thread_local() {
         this.registry.registerThreadLocalAccessor(new ObservationThreadLocalAccessor());
 
-        then(ObservationThreadLocalHolder.getValue()).isNull();
         ObservationThreadLocalHolder.setValue("hello");
-
         ContextSnapshot snapshot = ContextSnapshot.captureUsing(this.registry, key -> true);
 
-        ObservationThreadLocalHolder.reset();
-        then(ObservationThreadLocalHolder.getValue()).isNull();
-
-        try (ContextSnapshot.Scope scope = snapshot.setThreadLocalValues()) {
-            then(ObservationThreadLocalHolder.getValue()).isEqualTo("hello");
+        ObservationThreadLocalHolder.setValue("hola");
+        try {
+            try (Scope scope = snapshot.setThreadLocalValues()) {
+                then(ObservationThreadLocalHolder.getValue()).isEqualTo("hello");
+            }
+            then(ObservationThreadLocalHolder.getValue()).isEqualTo("hola");
+        }
+        finally {
+            ObservationThreadLocalHolder.reset();
         }
-
-        then(ObservationThreadLocalHolder.getValue()).isNull();
     }
 
     @Test
-    void should_reset_to_thread_local_to_previous_value() {
+    void should_propagate_single_thread_local_value() {
+        this.registry.registerContextAccessor(new TestContextAccessor());
         this.registry.registerThreadLocalAccessor(new ObservationThreadLocalAccessor());
 
-        ObservationThreadLocalHolder.setValue("hello");
-        ContextSnapshot snapshot = ContextSnapshot.captureUsing(this.registry, key -> true);
+        String key = ObservationThreadLocalAccessor.KEY;
+        Map<String, String> sourceContext = Collections.singletonMap(key, "hello");
 
         ObservationThreadLocalHolder.setValue("hola");
         try {
-            try (ContextSnapshot.Scope scope = snapshot.setThreadLocalValues()) {
+            try (Scope scope = ContextSnapshot.setThreadLocalsFrom(sourceContext, this.registry, key)) {
                 then(ObservationThreadLocalHolder.getValue()).isEqualTo("hello");
             }
             then(ObservationThreadLocalHolder.getValue()).isEqualTo("hola");
@@ -80,7 +85,7 @@ void should_not_fail_on_empty_thread_local() {
         ObservationThreadLocalHolder.reset();
         then(ObservationThreadLocalHolder.getValue()).isNull();
 
-        try (ContextSnapshot.Scope scope = snapshot.setThreadLocalValues()) {
+        try (Scope scope = snapshot.setThreadLocalValues()) {
             then(ObservationThreadLocalHolder.getValue()).isNull();
         }
 
@@ -104,7 +109,7 @@ void should_filter_thread_locals_on_capture() {
         fooThreadLocal.remove();
         barThreadLocal.remove();
 
-        try (ContextSnapshot.Scope scope = snapshot.setThreadLocalValues()) {
+        try (Scope scope = snapshot.setThreadLocalValues()) {
             then(fooThreadLocal.get()).isEqualTo("fooValue");
             then(barThreadLocal.get()).isNull();
         }
@@ -130,12 +135,12 @@ void should_filter_thread_locals_on_restore() {
         fooThreadLocal.remove();
         barThreadLocal.remove();
 
-        try (ContextSnapshot.Scope scope = snapshot.setThreadLocalValues(key -> key.equals("foo"))) {
+        try (Scope scope = snapshot.setThreadLocalValues(key -> key.equals("foo"))) {
             then(fooThreadLocal.get()).isEqualTo("fooValue");
             then(barThreadLocal.get()).isNull();
         }
 
-        try (ContextSnapshot.Scope scope = snapshot.setThreadLocalValues(key -> key.equals("bar"))) {
+        try (Scope scope = snapshot.setThreadLocalValues(key -> key.equals("bar"))) {
             then(fooThreadLocal.get()).isNull();
             then(barThreadLocal.get()).isEqualTo("barValue");
         }
diff --git a/context-propagation/src/test/java/io/micrometer/context/TestContextAccessor.java b/context-propagation/src/test/java/io/micrometer/context/TestContextAccessor.java
@@ -36,6 +36,13 @@ public void readValues(Map<?, ?> sourceContext, Predicate<Object> keyPredicate,
         readValues.putAll(sourceContext);
     }
 
+    @SuppressWarnings("unchecked")
+    @Nullable
+    @Override
+    public <T> T readValue(Map<?, ?> sourceContext, Object key) {
+        return (T) sourceContext.get(key);
+    }
+
     @Override
     public boolean canWriteTo(Class<?> contextType) {
         return Map.class.isAssignableFrom(contextType);
