Add more intuitive method to wrap ExecutorService (#305)

The `ContextExecutorService#wrap` method accepting a
`Supplier<ContextSnapshot>` can be a mystery for users as it's not
immediately obvious that a `ContextSnapshotFactory` needs to be used to
supply a `ContextSnapshot`. This change adds an overload that accepts
the `ContextSnapshotFactory` explicitly and only works with ThreadLocal
values as a sane default. For more advanced use cases the variant
accepting `Supplier` can be used.

The same variant is also added to `ContextScheduledExecutorService`.

Author: chemicL

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

@@ -139,12 +139,35 @@ protected ContextSnapshot capture() {
         return this.contextSnapshot.get();
     }
 
+    /**
+     * Wrap the given {@code ExecutorService} in order to propagate context to any
+     * executed task through the given {@link ContextSnapshotFactory}.
+     * <p>
+     * This method only captures ThreadLocal value. To work with other types of contexts,
+     * use {@link #wrap(ExecutorService, Supplier)}.
+     * </p>
+     * @param service the executorService to wrap
+     * @param contextSnapshotFactory {@link ContextSnapshotFactory} for capturing a
+     * {@link ContextSnapshot} at the point when tasks are scheduled
+     * @return {@code ExecutorService} wrapper
+     * @since 1.1.2
+     */
+    public static ExecutorService wrap(ExecutorService service, ContextSnapshotFactory contextSnapshotFactory) {
+        return new ContextExecutorService<>(service, contextSnapshotFactory::captureAll);
+    }
+
     /**
      * Wrap the given {@code ExecutorService} in order to propagate context to any
      * executed task through the given {@link ContextSnapshot} supplier.
+     * <p>
+     * Typically, a {@link ContextSnapshotFactory} can be used to supply the snapshot. In
+     * the case that only ThreadLocal values are to be captured, the
+     * {@link #wrap(ExecutorService, ContextSnapshotFactory)} variant can be used.
+     * </p>
      * @param service the executorService to wrap
      * @param snapshotSupplier supplier for capturing a {@link ContextSnapshot} at the
      * point when tasks are scheduled
+     * @return {@code ExecutorService} wrapper
      */
     public static ExecutorService wrap(ExecutorService service, Supplier<ContextSnapshot> snapshotSupplier) {
         return new ContextExecutorService<>(service, snapshotSupplier);
@@ -154,6 +177,7 @@ 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
+     * @return {@code ExecutorService} wrapper
      * @deprecated use {@link #wrap(ExecutorService, Supplier)}
      */
     @Deprecated

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

@@ -64,12 +64,37 @@ public ScheduledFuture<?> scheduleWithFixedDelay(Runnable command, long initialD
         return getExecutorService().scheduleWithFixedDelay(capture().wrap(command), initialDelay, delay, unit);
     }
 
+    /**
+     * Wrap the given {@code ScheduledExecutorService} in order to propagate context to
+     * any executed task through the given {@link ContextSnapshotFactory}.
+     * <p>
+     * This method only captures ThreadLocal value. To work with other types of contexts,
+     * use {@link #wrap(ScheduledExecutorService, Supplier)}.
+     * </p>
+     * @param service the executorService to wrap
+     * @param contextSnapshotFactory {@link ContextSnapshotFactory} for capturing a
+     * {@link ContextSnapshot} at the point when tasks are scheduled
+     * @return {@code ScheduledExecutorService} wrapper
+     * @since 1.1.2
+     */
+    public static ScheduledExecutorService wrap(ScheduledExecutorService service,
+            ContextSnapshotFactory contextSnapshotFactory) {
+        return new ContextScheduledExecutorService(service, contextSnapshotFactory::captureAll);
+    }
+
     /**
      * Wrap the given {@code ScheduledExecutorService} in order to propagate context to
      * any executed task through the given {@link ContextSnapshot} supplier.
+     * <p>
+     * Typically, a {@link ContextSnapshotFactory} can be used to supply the snapshot. In
+     * the case that only ThreadLocal values are to be captured, the
+     * {@link #wrap(ScheduledExecutorService, ContextSnapshotFactory)} variant can be
+     * used.
+     * </p>
      * @param service the executorService to wrap
      * @param supplier supplier for capturing a {@link ContextSnapshot} at the point when
      * tasks are scheduled
+     * @return {@code ScheduledExecutorService} wrapper
      */
     public static ScheduledExecutorService wrap(ScheduledExecutorService service, Supplier<ContextSnapshot> supplier) {
         return new ContextScheduledExecutorService(service, supplier);
@@ -79,6 +104,7 @@ 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
+     * @return {@code ScheduledExecutorService} wrapper
      * @deprecated use {@link #wrap(ScheduledExecutorService, Supplier)}
      */
     @Deprecated

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

@@ -16,6 +16,7 @@
 package io.micrometer.context;
 
 import java.util.Collections;
+import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
@@ -39,7 +40,8 @@
 class ContextWrappingTests {
 
     private final ContextRegistry registry = new ContextRegistry()
-        .registerThreadLocalAccessor(new StringThreadLocalAccessor());
+        .registerThreadLocalAccessor(new StringThreadLocalAccessor())
+        .registerContextAccessor(new TestContextAccessor());
 
     private final ContextSnapshotFactory defaultSnapshotFactory = ContextSnapshotFactory.builder()
         .contextRegistry(registry)
@@ -140,6 +142,43 @@ void should_instrument_scheduled_executor_service()
         }
     }
 
+    @Test
+    void should_instrument_scheduled_executor_service_with_snapshot_supplier()
+            throws InterruptedException, ExecutionException, TimeoutException {
+        Map<String, String> sourceContext = Collections.singletonMap(StringThreadLocalAccessor.KEY, "hello from map");
+
+        ScheduledExecutorService executorService = Executors.newSingleThreadScheduledExecutor();
+        try {
+            StringThreadLocalHolder.setValue("hello at time of creation of the executor");
+            AtomicReference<String> valueInNewThread = new AtomicReference<>();
+            runInNewThread(executorService, valueInNewThread,
+                    atomic -> then(atomic.get()).as("By default thread local information should not be propagated")
+                        .isNull());
+
+            StringThreadLocalHolder.setValue("hello at time of creation of the executor");
+            runInNewThread(
+                    ContextExecutorService
+                        .wrap(executorService, () -> defaultSnapshotFactory.captureAll(sourceContext)),
+                    valueInNewThread,
+                    atomic -> then(atomic.get())
+                        .as("With context container the thread local information should be propagated")
+                        .isEqualTo("hello from map"));
+
+            StringThreadLocalHolder.setValue("hello at time of creation of the executor");
+            runInNewThread(
+                    ContextScheduledExecutorService
+                        .wrap(executorService, () -> defaultSnapshotFactory.captureAll(sourceContext)),
+                    valueInNewThread,
+                    atomic -> then(atomic.get())
+                        .as("With context container the thread local information should be propagated")
+                        .isEqualTo("hello from map"));
+        }
+        finally {
+            executorService.shutdown();
+        }
+
+    }
+
     private void runInNewThread(Runnable runnable) throws InterruptedException, TimeoutException {
         CountDownLatch latch = new CountDownLatch(1);
         Thread thread = new Thread(countDownWhenDone(runnable, latch));