Add helper class to capture context using ScheduledExecutorService

context/src/main/java/io/opentelemetry/context/Context.java

@@ -138,6 +138,24 @@ static ExecutorService taskWrapping(ExecutorService executorService) {
     return new CurrentContextExecutorService(executorService);
   }
 
+  /**
+   * Returns an {@link ScheduledExecutorService} which delegates to the provided {@code
+   * executorService}, wrapping all invocations of {@link ExecutorService} methods such as {@link
+   * ExecutorService#execute(Runnable)} or {@link ExecutorService#submit(Runnable)} with the
+   * {@linkplain Context#current() current context} at the time of invocation.
+   *
+   * <p>This is generally used to create an {@link ScheduledExecutorService} which will forward the
+   * {@link Context} during an invocation to another thread. For example, you may use something like
+   * {@code ScheduledExecutorService dbExecutor = Context.wrapTasks(threadPool)} to ensure calls
+   * like {@code dbExecutor.execute(() -> database.query())} have {@link Context} available on the
+   * thread executing database queries.
+   *
+   * @since 1.43.0
+   */
+  static ScheduledExecutorService taskWrapping(ScheduledExecutorService executorService) {
+    return new CurrentContextScheduledExecutorService(executorService);
+  }
+
   /**
    * Returns the value stored in this {@link Context} for the given {@link ContextKey}, or {@code
    * null} if there is no value for the key in this context.

context/src/main/java/io/opentelemetry/context/CurrentContextScheduledExecutorService.java

@@ -0,0 +1,92 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.context;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+final class CurrentContextScheduledExecutorService extends ForwardingScheduledExecutorService {
+
+  CurrentContextScheduledExecutorService(ScheduledExecutorService delegate) {
+    super(delegate);
+  }
+
+  @Override
+  public <T> Future<T> submit(Callable<T> task) {
+    return delegate().submit(Context.current().wrap(task));
+  }
+
+  @Override
+  public <T> Future<T> submit(Runnable task, T result) {
+    return delegate().submit(Context.current().wrap(task), result);
+  }
+
+  @Override
+  public Future<?> submit(Runnable task) {
+    return delegate().submit(Context.current().wrap(task));
+  }
+
+  @Override
+  public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks)
+      throws InterruptedException {
+    return delegate().invokeAll(wrap(Context.current(), tasks));
+  }
+
+  @Override
+  public <T> List<Future<T>> invokeAll(
+      Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit)
+      throws InterruptedException {
+    return delegate().invokeAll(wrap(Context.current(), tasks), timeout, unit);
+  }
+
+  @Override
+  public <T> T invokeAny(Collection<? extends Callable<T>> tasks)
+      throws InterruptedException, ExecutionException {
+    return delegate().invokeAny(wrap(Context.current(), tasks));
+  }
+
+  @Override
+  public <T> T invokeAny(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit)
+      throws InterruptedException, ExecutionException, TimeoutException {
+    return delegate().invokeAny(wrap(Context.current(), tasks), timeout, unit);
+  }
+
+  @Override
+  public void execute(Runnable command) {
+    delegate().execute(Context.current().wrap(command));
+  }
+
+  @Override
+  public ScheduledFuture<?> schedule(Runnable command, long delay, TimeUnit unit) {
+    return delegate().schedule(Context.current().wrap(command), delay, unit);
+  }
+
+  @Override
+  public <V> ScheduledFuture<V> schedule(Callable<V> callable, long delay, TimeUnit unit) {
+    return delegate().schedule(Context.current().wrap(callable), delay, unit);
+  }
+
+  @Override
+  public ScheduledFuture<?> scheduleAtFixedRate(
+      Runnable command, long initialDelay, long period, TimeUnit unit) {
+    return delegate()
+        .scheduleAtFixedRate(Context.current().wrap(command), initialDelay, period, unit);
+  }
+
+  @Override
+  public ScheduledFuture<?> scheduleWithFixedDelay(
+      Runnable command, long initialDelay, long delay, TimeUnit unit) {
+    return delegate()
+        .scheduleWithFixedDelay(Context.current().wrap(command), initialDelay, delay, unit);
+  }
+}

context/src/main/java/io/opentelemetry/context/ForwardingScheduledExecutorService.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.context;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+
+/** A {@link ScheduledExecutorService} that implements methods that don't need {@link Context}. */
+abstract class ForwardingScheduledExecutorService implements ScheduledExecutorService {
+
+  private final ScheduledExecutorService delegate;
+
+  protected ForwardingScheduledExecutorService(ScheduledExecutorService delegate) {
+    this.delegate = delegate;
+  }
+
+  ScheduledExecutorService delegate() {
+    return delegate;
+  }
+
+  @Override
+  public final void shutdown() {
+    delegate.shutdown();
+  }
+
+  @Override
+  public final List<Runnable> shutdownNow() {
+    return delegate.shutdownNow();
+  }
+
+  @Override
+  public final boolean isShutdown() {
+    return delegate.isShutdown();
+  }
+
+  @Override
+  public final boolean isTerminated() {
+    return delegate.isTerminated();
+  }
+
+  @Override
+  public final boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
+    return delegate.awaitTermination(timeout, unit);
+  }
+
+  protected static <T> Collection<? extends Callable<T>> wrap(
+      Context context, Collection<? extends Callable<T>> tasks) {
+    List<Callable<T>> wrapped = new ArrayList<>();
+    for (Callable<T> task : tasks) {
+      wrapped.add(context.wrap(task));
+    }
+    return wrapped;
+  }
+}

context/src/test/java/io/opentelemetry/context/ContextTest.java

@@ -366,6 +366,10 @@ class WrapExecutorService {
     protected ExecutorService wrapped;
     protected AtomicReference<String> value;
 
+    protected ExecutorService wrap(ScheduledExecutorService executorService) {
+      return CAT.wrap(executorService);
+    }
+
     protected ExecutorService wrap(ExecutorService executorService) {
       return CAT.wrap(executorService);
     }
@@ -525,6 +529,30 @@ void close() {
     }
   }
 
+  @Nested
+  @TestInstance(Lifecycle.PER_CLASS)
+  class CurrentContextWrappingScheduledExecutorService extends WrapExecutorService {
+    @Override
+    protected ExecutorService wrap(ScheduledExecutorService executorService) {
+      return Context.taskWrapping(executorService);
+    }
+
+    private Scope scope;
+
+    @BeforeEach
+    // Closed in AfterEach
+    @SuppressWarnings("MustBeClosedChecker")
+    void makeCurrent() {
+      scope = CAT.makeCurrent();
+    }
+
+    @AfterEach
+    void close() {
+      scope.close();
+      scope = null;
+    }
+  }
+
   @Test
   void keyToString() {
     assertThat(ANIMAL.toString()).isEqualTo("animal");

docs/apidiffs/current_vs_latest/opentelemetry-context.txt

@@ -1,2 +1,4 @@
 Comparing source compatibility of opentelemetry-context-1.43.0-SNAPSHOT.jar against opentelemetry-context-1.42.1.jar
-No changes.
+***  MODIFIED INTERFACE: PUBLIC ABSTRACT io.opentelemetry.context.Context  (not serializable)
+	===  CLASS FILE FORMAT VERSION: 52.0 <- 52.0
+	+++  NEW METHOD: PUBLIC(+) STATIC(+) java.util.concurrent.ScheduledExecutorService taskWrapping(java.util.concurrent.ScheduledExecutorService)