Introducing ContextSnapshotFactory (#105)

This change addresses these two goals:

* To allow clearing ThreadLocal values not present in source context when
  opening a scope.
* To clean the API from accumulated static methods in ContextSnapshot and make
  the use more intuitive and comprehensible.

The responsibility of ContextSnapshot has been reduced. ContextSnapshotFactory
is introduced to capture the shared concerns:

* ContextRegistry to use when capturing and setting ThreadLocal values.
* Predicate for keys used when capturing state to a ContextSnpashot.
* Flag whether to clear ThreadLocal values missing in the used context object.

Multiple deprecations have been introduced with pointers on how to combine the
new building blocks to achieve the desired effect.

ContextSnapshotFactory contains a Builder interface with a default
implementation provided.

---------

Co-authored-by: Rossen Stoyanchev <[email protected]>

Author: Dariusz Jędrzejczyk

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

@@ -154,9 +154,12 @@ public static ExecutorService wrap(ExecutorService service, Supplier<ContextSnap
      * Variant of {@link #wrap(ExecutorService, Supplier)} that uses
      * {@link ContextSnapshot#captureAll(Object...)} to create the context snapshot.
      * @param service the executorService to wrap
+     * @deprecated use {@link #wrap(ExecutorService, Supplier)}
      */
+    @Deprecated
     public static ExecutorService wrap(ExecutorService service) {
-        return new ContextExecutorService<>(service, ContextSnapshot::captureAll);
+        return wrap(service,
+                () -> DefaultContextSnapshotFactory.captureAll(ContextRegistry.getInstance(), key -> true, false));
     }
 
 }

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

@@ -122,7 +122,7 @@ public void setValue(V value) {
             }
 
             @Override
-            public void reset() {
+            public void setValue() {
                 resetTask.run();
             }
         });

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

@@ -79,9 +79,12 @@ public static ScheduledExecutorService wrap(ScheduledExecutorService service, Su
      * Variant of {@link #wrap(ScheduledExecutorService, Supplier)} that uses
      * {@link ContextSnapshot#captureAll(Object...)} to create the context snapshot.
      * @param service the executorService to wrap
+     * @deprecated use {@link #wrap(ScheduledExecutorService, Supplier)}
      */
+    @Deprecated
     public static ScheduledExecutorService wrap(ScheduledExecutorService service) {
-        return new ContextScheduledExecutorService(service, ContextSnapshot::captureAll);
+        return wrap(service,
+                () -> DefaultContextSnapshotFactory.captureAll(ContextRegistry.getInstance(), key -> true, false));
     }
 
 }

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

@@ -24,7 +24,9 @@
 /**
  * Holds values extracted from {@link ThreadLocal} and other types of context and exposes
  * methods to propagate those values.
- *
+ * <p>
+ * Use {@link ContextSnapshotFactory#builder()} to configure a factory to work with
+ * snapshots.
  * <p>
  * Implementations are disallowed to store {@code null} values. If a {@link ThreadLocal}
  * is not set, or it's value is {@code null}, there is no way of distinguishing one from
@@ -34,11 +36,9 @@
  * {@link ContextAccessor#readValues(Object, Predicate, Map)}, and should likewise ignore
  * {@code null} values from {@link ContextAccessor#readValue(Object, Object)}.
  *
- * <p>
- * Use static factory methods on this interface to create a snapshot.
- *
  * @author Rossen Stoyanchev
  * @author Brian Clozel
+ * @author Dariusz Jędrzejczyk
  * @since 1.0.0
  */
 public interface ContextSnapshot {
@@ -72,6 +72,7 @@ public interface ContextSnapshot {
     /**
      * Variant of {@link #setThreadLocals()} with a predicate to select context values by
      * key.
+     * @param keyPredicate selects keys for use when setting {@link ThreadLocal} values
      * @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.
@@ -82,6 +83,7 @@ public interface ContextSnapshot {
      * Return a new {@code Runnable} that sets {@code ThreadLocal} values from the
      * snapshot around the invocation of the given {@code Runnable}.
      * @param runnable the runnable to instrument
+     * @return wrapped instance
      */
     default Runnable wrap(Runnable runnable) {
         return () -> {
@@ -96,6 +98,7 @@ default Runnable wrap(Runnable runnable) {
      * snapshot around the invocation of the given {@code Callable}.
      * @param callable the callable to instrument
      * @param <T> the type of value produced by the {@code Callable}
+     * @return wrapped instance
      */
     default <T> Callable<T> wrap(Callable<T> callable) {
         return () -> {
@@ -110,6 +113,7 @@ default <T> Callable<T> wrap(Callable<T> callable) {
      * snapshot around the invocation of the given {@code Consumer}.
      * @param consumer the callable to instrument
      * @param <T> the type of value produced by the {@code Callable}
+     * @return wrapped instance
      */
     default <T> Consumer<T> wrap(Consumer<T> consumer) {
         return value -> {
@@ -123,6 +127,7 @@ default <T> Consumer<T> wrap(Consumer<T> consumer) {
      * Return a new {@code Executor} that sets {@code ThreadLocal} values from the
      * snapshot around the invocation of any executed, {@code Runnable}.
      * @param executor the executor to instrument
+     * @return wrapped instance
      * @see ContextExecutorService
      * @see ContextScheduledExecutorService
      */
@@ -139,7 +144,10 @@ default Executor wrapExecutor(Executor executor) {
      * instance.
      * @param contexts one more context objects to extract values from
      * @return a snapshot with saved context values
+     * @deprecated use {@link ContextSnapshotFactory#captureAll(Object...)} on a factory
+     * obtained via a {@link ContextSnapshotFactory#builder()}.
      */
+    @Deprecated
     static ContextSnapshot captureAll(Object... contexts) {
         return captureAll(ContextRegistry.getInstance(), contexts);
     }
@@ -151,7 +159,11 @@ static ContextSnapshot captureAll(Object... contexts) {
      * @param registry the registry to use
      * @param contexts one more context objects to extract values from
      * @return a snapshot with saved context values
+     * @deprecated use {@link ContextSnapshotFactory#captureAll(Object...)} on a factory
+     * obtained via a {@link ContextSnapshotFactory#builder()} combined with
+     * {@link ContextSnapshotFactory.Builder#contextRegistry(ContextRegistry)}.
      */
+    @Deprecated
     static ContextSnapshot captureAll(ContextRegistry registry, Object... contexts) {
         return captureAllUsing(key -> true, registry, contexts);
     }
@@ -163,61 +175,23 @@ static ContextSnapshot captureAll(ContextRegistry registry, Object... contexts)
      * @param registry the registry with the accessors to use
      * @param contexts one more context objects to extract values from
      * @return a snapshot with saved context values
+     * @deprecated use {@link ContextSnapshotFactory#captureAll(Object...)} on a factory
+     * obtained via a {@link ContextSnapshotFactory#builder()} and configure
+     * {@link ContextSnapshotFactory.Builder#captureKeyPredicate(Predicate)} and
+     * {@link ContextSnapshotFactory.Builder#contextRegistry(ContextRegistry)}.
      */
+    @Deprecated
     static ContextSnapshot captureAllUsing(Predicate<Object> keyPredicate, ContextRegistry registry,
             Object... contexts) {
-
-        return DefaultContextSnapshot.captureAll(registry, keyPredicate, contexts);
-    }
-
-    /**
-     * 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}
-     * @since 1.0.3
-     */
-    static ContextSnapshot captureFromContext(Object... contexts) {
-        return DefaultContextSnapshot.captureFromContext(key -> true, ContextRegistry.getInstance(), null, contexts);
-    }
-
-    /**
-     * 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}
-     * @since 1.0.3
-     */
-    static ContextSnapshot captureFromContext(ContextRegistry registry, Object... contexts) {
-        return DefaultContextSnapshot.captureFromContext(key -> true, registry, null, contexts);
-    }
-
-    /**
-     * 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}
-     * @since 1.0.3
-     */
-    static ContextSnapshot captureFromContext(Predicate<Object> keyPredicate, ContextRegistry registry,
-            Object... contexts) {
-        return DefaultContextSnapshot.captureFromContext(keyPredicate, registry, null, contexts);
+        return DefaultContextSnapshotFactory.captureAll(registry, keyPredicate, false, contexts);
     }
 
     /**
      * 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 use {@link ContextSnapshotFactory#captureFrom(Object...)} on a factory
+     * obtained via a {@link ContextSnapshotFactory#builder()}.
      */
     @Deprecated
     static ContextSnapshot captureFrom(Object context) {
@@ -229,12 +203,13 @@ static ContextSnapshot captureFrom(Object context) {
      * @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 use {@link ContextSnapshotFactory#captureFrom(Object...)} on a factory
+     * obtained via a {@link ContextSnapshotFactory#builder()} combined with
+     * {@link ContextSnapshotFactory.Builder#contextRegistry(ContextRegistry)}.
      */
     @Deprecated
     static ContextSnapshot captureFrom(Object context, ContextRegistry registry) {
-        return DefaultContextSnapshot.captureFromContext(key -> true, registry, null, context);
+        return DefaultContextSnapshotFactory.captureFromContext(key -> true, false, registry, null, context);
     }
 
     /**
@@ -243,12 +218,14 @@ static ContextSnapshot captureFrom(Object context, ContextRegistry registry) {
      * @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 use {@link ContextSnapshotFactory#captureFrom(Object...)} on a factory
+     * obtained via a {@link ContextSnapshotFactory#builder()} and configure
+     * {@link ContextSnapshotFactory.Builder#captureKeyPredicate(Predicate)} and
+     * {@link ContextSnapshotFactory.Builder#contextRegistry(ContextRegistry)}.
      */
     @Deprecated
     static ContextSnapshot captureFrom(Object context, Predicate<Object> keyPredicate, ContextRegistry registry) {
-        return DefaultContextSnapshot.captureFromContext(keyPredicate, registry, null, context);
+        return DefaultContextSnapshotFactory.captureFromContext(keyPredicate, false, registry, null, context);
     }
 
     /**
@@ -258,9 +235,14 @@ static ContextSnapshot captureFrom(Object context, Predicate<Object> keyPredicat
      * @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.
+     * @deprecated use
+     * {@link ContextSnapshotFactory#setThreadLocalsFrom(Object, String...)} with no keys
+     * on a factory obtained via a {@link ContextSnapshotFactory#builder()}.
      */
+    @Deprecated
     static Scope setAllThreadLocalsFrom(Object sourceContext) {
-        return DefaultContextSnapshot.setAllThreadLocalsFrom(sourceContext, ContextRegistry.getInstance());
+        return DefaultContextSnapshotFactory.setAllThreadLocalsFrom(sourceContext, ContextRegistry.getInstance(),
+                false);
     }
 
     /**
@@ -271,9 +253,14 @@ static Scope setAllThreadLocalsFrom(Object sourceContext) {
      * @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.
+     * @deprecated use
+     * {@link ContextSnapshotFactory#setThreadLocalsFrom(Object, String...)} with no keys
+     * on a factory obtained via a {@link ContextSnapshotFactory#builder()} combined with
+     * {@link ContextSnapshotFactory.Builder#contextRegistry(ContextRegistry)}.
      */
+    @Deprecated
     static Scope setAllThreadLocalsFrom(Object sourceContext, ContextRegistry contextRegistry) {
-        return DefaultContextSnapshot.setAllThreadLocalsFrom(sourceContext, contextRegistry);
+        return DefaultContextSnapshotFactory.setAllThreadLocalsFrom(sourceContext, contextRegistry, false);
     }
 
     /**
@@ -286,7 +273,11 @@ static Scope setAllThreadLocalsFrom(Object sourceContext, ContextRegistry contex
      * @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.
+     * @deprecated use
+     * {@link ContextSnapshotFactory#setThreadLocalsFrom(Object, String...)} on a factory
+     * obtained via a {@link ContextSnapshotFactory#builder()}.
      */
+    @Deprecated
     static Scope setThreadLocalsFrom(Object sourceContext, String... keys) {
         return setThreadLocalsFrom(sourceContext, ContextRegistry.getInstance(), keys);
     }
@@ -300,9 +291,14 @@ static Scope setThreadLocalsFrom(Object sourceContext, String... keys) {
      * @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.
+     * @deprecated use
+     * {@link ContextSnapshotFactory#setThreadLocalsFrom(Object, String...)} on a factory
+     * obtained via a {@link ContextSnapshotFactory#builder()} combined with
+     * {@link ContextSnapshotFactory.Builder#contextRegistry(ContextRegistry)}.
      */
+    @Deprecated
     static Scope setThreadLocalsFrom(Object sourceContext, ContextRegistry contextRegistry, String... keys) {
-        return DefaultContextSnapshot.setThreadLocalsFrom(sourceContext, contextRegistry, keys);
+        return DefaultContextSnapshotFactory.setThreadLocalsFrom(sourceContext, contextRegistry, false, keys);
     }
 
     /**