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