Add helper class to capture context using ScheduledExecutorService (#6712)

Signed-off-by: Adriano Machado <[email protected]>

Author: ammachado

.gitignore

@@ -1,6 +1,7 @@
 # Gradle
 build
 .gradle
+.kotlin
 local.properties
 out/
 

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

@@ -27,6 +27,7 @@
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
 import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
@@ -135,9 +136,37 @@ static Executor taskWrapping(Executor executor) {
    * @since 1.1.0
    */
   static ExecutorService taskWrapping(ExecutorService executorService) {
+    if (executorService instanceof CurrentContextExecutorService) {
+      return 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.
+   *
+   * <p>Note: The context will not be propagated for {@link
+   * ScheduledExecutorService#scheduleAtFixedRate(Runnable, long, long, TimeUnit)} and {@link
+   * ScheduledExecutorService#scheduleWithFixedDelay(Runnable, long, long, TimeUnit)} calls.
+   *
+   * @since 1.43.0
+   */
+  static ScheduledExecutorService taskWrapping(ScheduledExecutorService executorService) {
+    if (executorService instanceof CurrentContextScheduledExecutorService) {
+      return 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/CurrentContextExecutorService.java

@@ -14,7 +14,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
-final class CurrentContextExecutorService extends ForwardingExecutorService {
+class CurrentContextExecutorService extends ForwardingExecutorService {
 
   CurrentContextExecutorService(ExecutorService delegate) {
     super(delegate);

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

@@ -0,0 +1,44 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.context;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.TimeUnit;
+
+final class CurrentContextScheduledExecutorService extends CurrentContextExecutorService
+    implements ScheduledExecutorService {
+
+  private final ScheduledExecutorService delegate;
+
+  CurrentContextScheduledExecutorService(ScheduledExecutorService delegate) {
+    super(delegate);
+    this.delegate = delegate;
+  }
+
+  @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(command, initialDelay, period, unit);
+  }
+
+  @Override
+  public ScheduledFuture<?> scheduleWithFixedDelay(
+      Runnable command, long initialDelay, long delay, TimeUnit unit) {
+    return delegate.scheduleWithFixedDelay(command, initialDelay, delay, unit);
+  }
+}