diff --git a/context-propagation/src/main/java/io/micrometer/context/ContextSnapshot.java b/context-propagation/src/main/java/io/micrometer/context/ContextSnapshot.java index 43228c0..a506c2a 100644 --- a/context-propagation/src/main/java/io/micrometer/context/ContextSnapshot.java +++ b/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 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. + *

+ * 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. + *

+ * 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 keyPredicate, ContextRegistry registry) { return DefaultContextSnapshot.captureFromContext(keyPredicate, registry, context, null); } + /** + * Create a {@link ContextSnapshot} by reading values from the given context objects. + *

+ * 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 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. diff --git a/context-propagation/src/main/java/io/micrometer/context/DefaultContextSnapshot.java b/context-propagation/src/main/java/io/micrometer/context/DefaultContextSnapshot.java index f8655f5..002a964 100644 --- a/context-propagation/src/main/java/io/micrometer/context/DefaultContextSnapshot.java +++ b/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 implements ContextSnapshot { @@ -151,6 +152,15 @@ private static DefaultContextSnapshot captureFromThreadLocals(Predicate return snapshot; } + static ContextSnapshot captureFromContexts(Predicate 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 keyPredicate, ContextRegistry contextRegistry, Object context, @Nullable DefaultContextSnapshot snapshot) { diff --git a/context-propagation/src/test/java/io/micrometer/context/DefaultContextSnapshotTests.java b/context-propagation/src/test/java/io/micrometer/context/DefaultContextSnapshotTests.java index d572166..ddeb26a 100644 --- a/context-propagation/src/test/java/io/micrometer/context/DefaultContextSnapshotTests.java +++ b/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 firstContext = Collections.singletonMap(key, "hello"); + Map 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());