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

This commit is contained in:
Anuraag Agrawal 2022-02-10 08:09:36 +09:00 committed by GitHub
parent 3f5778e647
commit cffbd3249c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 187 additions and 1 deletions

View File

@ -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();
}
};
}
}

View File

@ -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 {

View File

@ -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<>();

View File

@ -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)