Add context wrapping methods for arguments found in CompletableFuture APIs. (#4154)

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

@@ -27,6 +27,11 @@ import java.util.concurrent.Callable;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.function.BiConsumer;
+import java.util.function.BiFunction;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
 /**
@@ -245,4 +250,64 @@ public interface Context {
   default ScheduledExecutorService wrap(ScheduledExecutorService executor) {
     return new ContextScheduledExecutorService(this, executor);
   }
+
+  /**
+   * Returns a {@link Function} that makes this the {@linkplain Context#current() current context}
+   * and then invokes the input {@link Function}.
+   */
+  default <T, U> Function<T, U> wrapFunction(Function<T, U> function) {
+    return t -> {
+      try (Scope ignored = makeCurrent()) {
+        return function.apply(t);
+      }
+    };
+  }
+
+  /**
+   * Returns a {@link BiFunction} that makes this the {@linkplain Context#current() current context}
+   * and then invokes the input {@link BiFunction}.
+   */
+  default <T, U, V> BiFunction<T, U, V> wrapFunction(BiFunction<T, U, V> function) {
+    return (t, u) -> {
+      try (Scope ignored = makeCurrent()) {
+        return function.apply(t, u);
+      }
+    };
+  }
+
+  /**
+   * Returns a {@link Consumer} that makes this the {@linkplain Context#current() current context}
+   * and then invokes the input {@link Consumer}.
+   */
+  default <T> Consumer<T> wrapConsumer(Consumer<T> consumer) {
+    return t -> {
+      try (Scope ignored = makeCurrent()) {
+        consumer.accept(t);
+      }
+    };
+  }
+
+  /**
+   * Returns a {@link BiConsumer} that makes this the {@linkplain Context#current() current context}
+   * and then invokes the input {@link BiConsumer}.
+   */
+  default <T, U> BiConsumer<T, U> wrapConsumer(BiConsumer<T, U> consumer) {
+    return (t, u) -> {
+      try (Scope ignored = makeCurrent()) {
+        consumer.accept(t, u);
+      }
+    };
+  }
+
+  /**
+   * Returns a {@link Supplier} that makes this the {@linkplain Context#current() current context}
+   * and then invokes the input {@link Supplier}.
+   */
+  default <T> Supplier<T> wrapSupplier(Supplier<T> supplier) {
+    return () -> {
+      try (Scope ignored = makeCurrent()) {
+        return supplier.get();
+      }
+    };
+  }
 }

context/src/strictContextEnabledTest/java/io/opentelemetry/context/StrictContextStorageTest.java

@@ -29,6 +29,7 @@ import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.api.trace.TraceFlags;
 import io.opentelemetry.api.trace.TraceState;
+import io.opentelemetry.internal.testing.slf4j.SuppressLogger;
 import java.time.Duration;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
@@ -38,6 +39,7 @@ import org.junit.jupiter.api.extension.RegisterExtension;
 import org.slf4j.event.Level;
 import org.slf4j.event.LoggingEvent;
 
+@SuppressLogger(StrictContextStorage.class)
 @SuppressWarnings("MustBeClosedChecker")
 class StrictContextStorageTest {
 

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

@@ -26,7 +26,13 @@ import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiConsumer;
+import java.util.function.BiFunction;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import java.util.function.Supplier;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeAll;
@@ -203,6 +209,113 @@ class ContextTest {
     assertThat(value).hasValue(null);
   }
 
+  @Test
+  void wrapFunction() {
+    AtomicReference<String> value = new AtomicReference<>();
+    Function<String, String> callback =
+        (a) -> {
+          value.set(Context.current().get(ANIMAL));
+          return "foo";
+        };
+
+    assertThat(callback.apply("bar")).isEqualTo("foo");
+    assertThat(value).hasValue(null);
+
+    assertThat(CAT.wrapFunction(callback).apply("bar")).isEqualTo("foo");
+    assertThat(value).hasValue("cat");
+
+    assertThat(callback.apply("bar")).isEqualTo("foo");
+    assertThat(value).hasValue(null);
+  }
+
+  @Test
+  void wrapBiFunction() {
+    AtomicReference<String> value = new AtomicReference<>();
+    BiFunction<String, String, String> callback =
+        (a, b) -> {
+          value.set(Context.current().get(ANIMAL));
+          return "foo";
+        };
+
+    assertThat(callback.apply("bar", "baz")).isEqualTo("foo");
+    assertThat(value).hasValue(null);
+
+    assertThat(CAT.wrapFunction(callback).apply("bar", "baz")).isEqualTo("foo");
+    assertThat(value).hasValue("cat");
+
+    assertThat(callback.apply("bar", "baz")).isEqualTo("foo");
+    assertThat(value).hasValue(null);
+  }
+
+  @Test
+  void wrapConsumer() {
+    AtomicReference<String> value = new AtomicReference<>();
+    AtomicBoolean consumed = new AtomicBoolean();
+    Consumer<String> callback =
+        (a) -> {
+          value.set(Context.current().get(ANIMAL));
+          consumed.set(true);
+        };
+
+    callback.accept("bar");
+    assertThat(consumed).isTrue();
+    assertThat(value).hasValue(null);
+
+    consumed.set(false);
+    CAT.wrapConsumer(callback).accept("bar");
+    assertThat(consumed).isTrue();
+    assertThat(value).hasValue("cat");
+
+    consumed.set(false);
+    callback.accept("bar");
+    assertThat(consumed).isTrue();
+    assertThat(value).hasValue(null);
+  }
+
+  @Test
+  void wrapBiConsumer() {
+    AtomicReference<String> value = new AtomicReference<>();
+    AtomicBoolean consumed = new AtomicBoolean();
+    BiConsumer<String, String> callback =
+        (a, b) -> {
+          value.set(Context.current().get(ANIMAL));
+          consumed.set(true);
+        };
+
+    callback.accept("bar", "baz");
+    assertThat(consumed).isTrue();
+    assertThat(value).hasValue(null);
+
+    consumed.set(false);
+    CAT.wrapConsumer(callback).accept("bar", "baz");
+    assertThat(consumed).isTrue();
+    assertThat(value).hasValue("cat");
+
+    consumed.set(false);
+    callback.accept("bar", "baz");
+    assertThat(consumed).isTrue();
+    assertThat(value).hasValue(null);
+  }
+
+  @Test
+  void wrapSupplier() {
+    AtomicReference<String> value = new AtomicReference<>();
+    Supplier<String> callback =
+        () -> {
+          value.set(Context.current().get(ANIMAL));
+          return "foo";
+        };
+
+    assertThat(callback.get()).isEqualTo("foo");
+    assertThat(value).hasValue(null);
+
+    assertThat(CAT.wrapSupplier(callback).get()).isEqualTo("foo");
+    assertThat(value).hasValue("cat");
+
+    assertThat(callback.get()).isEqualTo("foo");
+    assertThat(value).hasValue(null);
+  }
+
   @Test
   void wrapExecutor() {
     AtomicReference<String> value = new AtomicReference<>();

docs/apidiffs/current_vs_latest/opentelemetry-context.txt

@@ -1,2 +1,8 @@
 Comparing source compatibility of  against 
-No changes.
+***! MODIFIED INTERFACE: PUBLIC ABSTRACT io.opentelemetry.context.Context  (not serializable)
+	===  CLASS FILE FORMAT VERSION: 52.0 <- 52.0
+	+++! NEW METHOD: PUBLIC(+) java.util.function.Consumer wrapConsumer(java.util.function.Consumer)
+	+++! NEW METHOD: PUBLIC(+) java.util.function.BiConsumer wrapConsumer(java.util.function.BiConsumer)
+	+++! NEW METHOD: PUBLIC(+) java.util.function.Function wrapFunction(java.util.function.Function)
+	+++! NEW METHOD: PUBLIC(+) java.util.function.BiFunction wrapFunction(java.util.function.BiFunction)
+	+++! NEW METHOD: PUBLIC(+) java.util.function.Supplier wrapSupplier(java.util.function.Supplier)