Add ContextSnapshot.captureFromMany variants

context-propagation/src/main/java/io/micrometer/context/ContextSnapshot.java

@@ -1,5 +1,5 @@
 /**
- * Copyright 2022 the original author or authors.
+ * Copyright 2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -28,6 +28,7 @@
  * Use static factory methods on this interface to create a snapshot.
  *
  * @author Rossen Stoyanchev
+ * @author Brian Clozel
  * @since 1.0.0
  */
 public interface ContextSnapshot {
@@ -163,30 +164,80 @@ static ContextSnapshot captureAllUsing(Predicate<Object> keyPredicate, ContextRe
      * Create a {@link ContextSnapshot} by reading values from the given context object.
      * @param context the context to read values from
      * @return the created {@link ContextSnapshot}
+     * @deprecated as of 1.0.3 in favor of {@link #captureFromContext(Object...)}
      */
+    @Deprecated
     static ContextSnapshot captureFrom(Object context) {
         return captureFrom(context, ContextRegistry.getInstance());
     }
 
+    /**
+     * Create a {@link ContextSnapshot} by reading values from the given context objects.
+     * <p>
+     * Values captured multiple times are overridden in the snapshot by the order of
+     * contexts given as arguments.
+     * @param contexts the contexts to read values from
+     * @return the created {@link ContextSnapshot}
+     */
+    static ContextSnapshot captureFromContext(Object... contexts) {
+        return DefaultContextSnapshot.captureFromContexts(key -> true, ContextRegistry.getInstance(), contexts);
+    }
+
     /**
      * Create a {@link ContextSnapshot} by reading values from the given context object.
      * @param context the context to read values from
      * @param registry the registry to use
      * @return the created {@link ContextSnapshot}
+     * @deprecated as of 1.0.3 in favor of
+     * {@link #captureFromContext(ContextRegistry, Object...)}
      */
+    @Deprecated
     static ContextSnapshot captureFrom(Object context, ContextRegistry registry) {
-        return captureFrom(context, key -> true, registry);
+        return DefaultContextSnapshot.captureFromContext(key -> true, registry, context, null);
+    }
+
+    /**
+     * Create a {@link ContextSnapshot} by reading values from the given context objects.
+     * <p>
+     * Values captured multiple times are overridden in the snapshot by the order of
+     * contexts given as arguments.
+     * @param registry the registry to use
+     * @param contexts the contexts to read values from
+     * @return the created {@link ContextSnapshot}
+     */
+    static ContextSnapshot captureFromContext(ContextRegistry registry, Object... contexts) {
+        return DefaultContextSnapshot.captureFromContexts(key -> true, registry, contexts);
     }
 
     /**
      * Create a {@link ContextSnapshot} by reading values from the given context object.
      * @param context the context to read values from
+     * @param keyPredicate predicate for context value keys
+     * @param registry the registry to use
      * @return the created {@link ContextSnapshot}
+     * @deprecated as of 1.0.3 in favor of
+     * {@link #captureFromContext(Predicate, ContextRegistry, Object...)}
      */
+    @Deprecated
     static ContextSnapshot captureFrom(Object context, Predicate<Object> keyPredicate, ContextRegistry registry) {
         return DefaultContextSnapshot.captureFromContext(keyPredicate, registry, context, null);
     }
 
+    /**
+     * Create a {@link ContextSnapshot} by reading values from the given context objects.
+     * <p>
+     * Values captured multiple times are overridden in the snapshot by the order of
+     * contexts given as arguments.
+     * @param keyPredicate predicate for context value keys
+     * @param registry the registry to use
+     * @param contexts the contexts to read values from
+     * @return the created {@link ContextSnapshot}
+     */
+    static ContextSnapshot captureFromContext(Predicate<Object> keyPredicate, ContextRegistry registry,
+            Object... contexts) {
+        return DefaultContextSnapshot.captureFromContexts(keyPredicate, registry, contexts);
+    }
+
     /**
      * Variant of {@link #setThreadLocalsFrom(Object, String...)} that sets all
      * {@link ThreadLocal} values for which there is a value in the given source context.

context-propagation/src/main/java/io/micrometer/context/DefaultContextSnapshot.java

@@ -1,5 +1,5 @@
 /**
- * Copyright 2022 the original author or authors.
+ * Copyright 2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -23,6 +23,7 @@
  * Default implementation of {@link ContextSnapshot}.
  *
  * @author Rossen Stoyanchev
+ * @author Brian Clozel
  * @since 1.0.0
  */
 final class DefaultContextSnapshot extends HashMap<Object, Object> implements ContextSnapshot {
@@ -151,6 +152,15 @@ private static DefaultContextSnapshot captureFromThreadLocals(Predicate<Object>
         return snapshot;
     }
 
+    static ContextSnapshot captureFromContexts(Predicate<Object> keyPredicate, ContextRegistry contextRegistry,
+            Object... contexts) {
+        DefaultContextSnapshot snapshot = null;
+        for (Object context : contexts) {
+            snapshot = captureFromContext(keyPredicate, contextRegistry, context, snapshot);
+        }
+        return (snapshot != null ? snapshot : emptyContextSnapshot);
+    }
+
     @SuppressWarnings("unchecked")
     static DefaultContextSnapshot captureFromContext(Predicate<Object> keyPredicate, ContextRegistry contextRegistry,
             Object context, @Nullable DefaultContextSnapshot snapshot) {

context-propagation/src/test/java/io/micrometer/context/DefaultContextSnapshotTests.java

@@ -1,5 +1,5 @@
 /**
- * Copyright 2022 the original author or authors.
+ * Copyright 2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -93,6 +93,25 @@ void should_propagate_all_single_thread_local_value() {
         }
     }
 
+    @Test
+    void should_override_context_values_when_many_contexts() {
+        this.registry.registerContextAccessor(new TestContextAccessor());
+
+        String key = ObservationThreadLocalAccessor.KEY;
+        Map<String, String> firstContext = Collections.singletonMap(key, "hello");
+        Map<String, String> secondContext = Collections.singletonMap(key, "override");
+        try {
+            ContextSnapshot contextSnapshot = ContextSnapshot.captureFromContext(this.registry, firstContext,
+                    secondContext);
+            contextSnapshot.wrap(() -> {
+                then(ObservationThreadLocalHolder.getValue()).isEqualTo("override");
+            });
+        }
+        finally {
+            ObservationThreadLocalHolder.reset();
+        }
+    }
+
     @Test
     void should_throw_an_exception_when_no_keys_are_passed() {
         this.registry.registerContextAccessor(new TestContextAccessor());