Add context wrapping methods for arguments found in CompletableFuture APIs. (#4154)
This commit is contained in:
		
							parent
							
								
									3f5778e647
								
							
						
					
					
						commit
						cffbd3249c
					
				| 
						 | 
				
			
			@ -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();
 | 
			
		||||
      }
 | 
			
		||||
    };
 | 
			
		||||
  }
 | 
			
		||||
}
 | 
			
		||||
| 
						 | 
				
			
			
 | 
			
		|||
| 
						 | 
				
			
			@ -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 {
 | 
			
		||||
 | 
			
		||||
| 
						 | 
				
			
			
 | 
			
		|||
| 
						 | 
				
			
			@ -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<>();
 | 
			
		||||
| 
						 | 
				
			
			
 | 
			
		|||
| 
						 | 
				
			
			@ -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)
 | 
			
		||||
| 
						 | 
				
			
			
 | 
			
		|||
		Loading…
	
		Reference in New Issue