Skip to content

support reactor 3.4 ContextView in library instrumentation #11156

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ public class ContextPropagationOperatorInstrumentation implements TypeInstrument
@Override
public ElementMatcher<TypeDescription> typeMatcher() {
return named(
"application.io.opentelemetry.instrumentation.reactor.v3_1.ContextPropagationOperator");
"application.io.opentelemetry.instrumentation.reactor.v3.common.ContextPropagationOperator");
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,10 @@ plugins {

dependencies {
library("io.projectreactor:reactor-core:3.1.0.RELEASE")

api(project(":instrumentation:reactor:reactor-common:library"))
implementation(project(":instrumentation-annotations-support"))

testLibrary("io.projectreactor:reactor-test:3.1.0.RELEASE")

testImplementation(project(":instrumentation:reactor:reactor-3.1:testing"))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -22,68 +22,12 @@

package io.opentelemetry.instrumentation.reactor.v3_1;

import static java.lang.invoke.MethodType.methodType;

import io.opentelemetry.context.Context;
import io.opentelemetry.context.Scope;
import io.opentelemetry.instrumentation.api.annotation.support.async.AsyncOperationEndStrategies;
import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.util.function.BiFunction;
import java.util.function.Function;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.annotation.Nullable;
import org.reactivestreams.Publisher;
import reactor.core.CoreSubscriber;
import reactor.core.Fuseable;
import reactor.core.Scannable;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Hooks;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Operators;
import reactor.core.scheduler.Schedulers;

/** Based on Spring Sleuth's Reactor instrumentation. */
public final class ContextPropagationOperator {
private static final Logger logger = Logger.getLogger(ContextPropagationOperator.class.getName());

private static final Object VALUE = new Object();

@Nullable
private static final MethodHandle MONO_CONTEXT_WRITE_METHOD = getContextWriteMethod(Mono.class);

@Nullable
private static final MethodHandle FLUX_CONTEXT_WRITE_METHOD = getContextWriteMethod(Flux.class);
public class ContextPropagationOperator
extends io.opentelemetry.instrumentation.reactor.v3.common.ContextPropagationOperator {

@Nullable private static final MethodHandle SCHEDULERS_HOOK_METHOD = getSchedulersHookMethod();

@Nullable
private static MethodHandle getContextWriteMethod(Class<?> type) {
MethodHandles.Lookup lookup = MethodHandles.publicLookup();
try {
return lookup.findVirtual(type, "contextWrite", methodType(type, Function.class));
} catch (NoSuchMethodException | IllegalAccessException e) {
// ignore
}
try {
return lookup.findVirtual(type, "subscriberContext", methodType(type, Function.class));
} catch (NoSuchMethodException | IllegalAccessException e) {
// ignore
}
return null;
}

@Nullable
private static MethodHandle getSchedulersHookMethod() {
MethodHandles.Lookup lookup = MethodHandles.publicLookup();
try {
return lookup.findStatic(
Schedulers.class, "onScheduleHook", methodType(void.class, String.class, Function.class));
} catch (NoSuchMethodException | IllegalAccessException e) {
// ignore
}
return null;
public ContextPropagationOperator(boolean captureExperimentalSpanAttributes) {
super(captureExperimentalSpanAttributes);
}

public static ContextPropagationOperator create() {
Expand All @@ -93,270 +37,4 @@ public static ContextPropagationOperator create() {
public static ContextPropagationOperatorBuilder builder() {
return new ContextPropagationOperatorBuilder();
}

private final ReactorAsyncOperationEndStrategy asyncOperationEndStrategy;

private static final Object TRACE_CONTEXT_KEY =
new Object() {
@Override
public String toString() {
return "otel-trace-context";
}
};

private static final Object lock = new Object();

private static volatile boolean enabled = false;

/**
* Stores Trace {@link io.opentelemetry.context.Context} in Reactor {@link
* reactor.util.context.Context}.
*
* @param context Reactor's context to store trace context in.
* @param traceContext Trace context to be stored.
*/
public static reactor.util.context.Context storeOpenTelemetryContext(
reactor.util.context.Context context, Context traceContext) {
return context.put(TRACE_CONTEXT_KEY, traceContext);
}

/**
* Gets Trace {@link io.opentelemetry.context.Context} from Reactor {@link
* reactor.util.context.Context}.
*
* @param context Reactor's context to get trace context from.
* @param defaultTraceContext Default value to be returned if no trace context is found on Reactor
* context.
* @return Trace context or default value.
*/
public static Context getOpenTelemetryContext(
reactor.util.context.Context context, Context defaultTraceContext) {
return context.getOrDefault(TRACE_CONTEXT_KEY, defaultTraceContext);
}

ContextPropagationOperator(boolean captureExperimentalSpanAttributes) {
this.asyncOperationEndStrategy =
ReactorAsyncOperationEndStrategy.builder()
.setCaptureExperimentalSpanAttributes(captureExperimentalSpanAttributes)
.build();
}

/**
* Registers a hook that applies to every operator, propagating {@link Context} to downstream
* callbacks to ensure spans in the {@link Context} are available throughout the lifetime of a
* reactive stream. This should generally be called in a static initializer block in your
* application.
*/
public void registerOnEachOperator() {
synchronized (lock) {
if (enabled) {
return;
}
Hooks.onEachOperator(
TracingSubscriber.class.getName(), tracingLift(asyncOperationEndStrategy));
AsyncOperationEndStrategies.instance().registerStrategy(asyncOperationEndStrategy);
registerScheduleHook(RunnableWrapper.class.getName(), RunnableWrapper::new);
enabled = true;
}
}

private static void registerScheduleHook(String key, Function<Runnable, Runnable> function) {
if (SCHEDULERS_HOOK_METHOD == null) {
return;
}
try {
SCHEDULERS_HOOK_METHOD.invoke(key, function);
} catch (Throwable throwable) {
logger.log(Level.WARNING, "Failed to install scheduler hook", throwable);
}
}

/** Unregisters the hook registered by {@link #registerOnEachOperator()}. */
public void resetOnEachOperator() {
synchronized (lock) {
if (!enabled) {
return;
}
Hooks.resetOnEachOperator(TracingSubscriber.class.getName());
AsyncOperationEndStrategies.instance().unregisterStrategy(asyncOperationEndStrategy);
enabled = false;
}
}

private static <T> Function<? super Publisher<T>, ? extends Publisher<T>> tracingLift(
ReactorAsyncOperationEndStrategy asyncOperationEndStrategy) {
return Operators.lift(
ContextPropagationOperator::shouldInstrument, new Lifter<>(asyncOperationEndStrategy));
}

/** Forces Mono to run in traceContext scope. */
@SuppressWarnings("unchecked")
public static <T> Mono<T> runWithContext(Mono<T> publisher, Context tracingContext) {
if (!enabled || MONO_CONTEXT_WRITE_METHOD == null) {
return publisher;
}

// this hack forces 'publisher' to run in the onNext callback of `TracingSubscriber`
// (created for this publisher) and with current() span that refers to span created here
// without the hack, publisher runs in the onAssembly stage, before traceContext is made current
try {
return (Mono<T>)
MONO_CONTEXT_WRITE_METHOD.invoke(
ScalarPropagatingMono.create(publisher),
new StoreOpenTelemetryContext(tracingContext));
} catch (Throwable t) {
// rethrowing without any wrapping to avoid any change to the underlying application behavior
throw sneakyThrow(t);
}
}

/** Forces Flux to run in traceContext scope. */
@SuppressWarnings("unchecked")
public static <T> Flux<T> runWithContext(Flux<T> publisher, Context tracingContext) {
if (!enabled || FLUX_CONTEXT_WRITE_METHOD == null) {
return publisher;
}

// this hack forces 'publisher' to run in the onNext callback of `TracingSubscriber`
// (created for this publisher) and with current() span that refers to span created here
// without the hack, publisher runs in the onAssembly stage, before traceContext is made current
try {
return (Flux<T>)
FLUX_CONTEXT_WRITE_METHOD.invoke(
ScalarPropagatingFlux.create(publisher),
new StoreOpenTelemetryContext(tracingContext));
} catch (Throwable t) {
// rethrowing without any wrapping to avoid any change to the underlying application behavior
throw sneakyThrow(t);
}
}

@SuppressWarnings({"TypeParameterUnusedInFormals", "unchecked"})
private static <T extends Throwable> T sneakyThrow(Throwable t) throws T {
throw (T) t;
}

private static class StoreOpenTelemetryContext
implements Function<reactor.util.context.Context, reactor.util.context.Context> {

private final Context tracingContext;

private StoreOpenTelemetryContext(Context tracingContext) {
this.tracingContext = tracingContext;
}

@Override
public reactor.util.context.Context apply(reactor.util.context.Context context) {
return storeOpenTelemetryContext(context, tracingContext);
}
}

private static boolean shouldInstrument(Scannable publisher) {
// skip if Flux/Mono #just, #empty, #error
return !(publisher instanceof Fuseable.ScalarCallable);
}

private static class Lifter<T>
implements BiFunction<Scannable, CoreSubscriber<? super T>, CoreSubscriber<? super T>> {

/** Holds reference to strategy to prevent it from being collected. */
@SuppressWarnings({"FieldCanBeLocal", "UnusedVariable"})
private final ReactorAsyncOperationEndStrategy asyncOperationEndStrategy;

public Lifter(ReactorAsyncOperationEndStrategy asyncOperationEndStrategy) {
this.asyncOperationEndStrategy = asyncOperationEndStrategy;
}

@Override
public CoreSubscriber<? super T> apply(Scannable publisher, CoreSubscriber<? super T> sub) {
return new TracingSubscriber<>(sub, sub.currentContext());
}
}

static void subscribeInActiveSpan(CoreSubscriber<? super Object> actual, Object value) {
Context tracingContextInReactor =
ContextPropagationOperator.getOpenTelemetryContext(actual.currentContext(), null);
if (tracingContextInReactor == null || tracingContextInReactor == Context.current()) {
actual.onSubscribe(Operators.scalarSubscription(actual, value));
} else {
try (Scope ignored = tracingContextInReactor.makeCurrent()) {
actual.onSubscribe(Operators.scalarSubscription(actual, value));
}
}
}

static class ScalarPropagatingMono extends Mono<Object> implements Scannable {

static <T> Mono<T> create(Mono<T> source) {
return new ScalarPropagatingMono(source).flatMap(unused -> source);
}

private final Mono<?> source;

private ScalarPropagatingMono(Mono<?> source) {
this.source = source;
}

@Override
public void subscribe(CoreSubscriber<? super Object> actual) {
subscribeInActiveSpan(actual, VALUE);
}

@Override
@Nullable
// Interface method doesn't have type parameter so we can't add it either.
@SuppressWarnings("rawtypes")
public Object scanUnsafe(Attr attr) {
if (attr == Attr.PARENT) {
return source;
}
return null;
}
}

static class ScalarPropagatingFlux extends Flux<Object> implements Scannable {

static <T> Flux<T> create(Flux<T> source) {
return new ScalarPropagatingFlux(source).flatMap(unused -> source);
}

private final Flux<?> source;

private ScalarPropagatingFlux(Flux<?> source) {
this.source = source;
}

@Override
public void subscribe(CoreSubscriber<? super Object> actual) {
subscribeInActiveSpan(actual, VALUE);
}

@Override
@Nullable
// Interface method doesn't have type parameter so we can't add it either.
@SuppressWarnings("rawtypes")
public Object scanUnsafe(Scannable.Attr attr) {
if (attr == Scannable.Attr.PARENT) {
return source;
}
return null;
}
}

private static class RunnableWrapper implements Runnable {
private final Runnable delegate;
private final Context context;

RunnableWrapper(Runnable delegate) {
this.delegate = delegate;
context = Context.current();
}

@Override
public void run() {
try (Scope ignore = context.makeCurrent()) {
delegate.run();
}
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

import static org.assertj.core.api.Assertions.assertThat;

import io.opentelemetry.instrumentation.reactor.v3.common.TracingSubscriber;
import java.util.concurrent.Callable;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
Expand Down
15 changes: 15 additions & 0 deletions instrumentation/reactor/reactor-3.4/library/build.gradle.kts
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
plugins {
id("otel.library-instrumentation")
}

dependencies {
api(project(":instrumentation:reactor:reactor-common:library"))
library("io.projectreactor:reactor-core:3.4.0")
implementation(project(":instrumentation-annotations-support"))
testLibrary("io.projectreactor:reactor-test:3.4.0")

testImplementation(project(":instrumentation:reactor:reactor-3.1:testing"))

latestDepTestLibrary("io.projectreactor:reactor-core:3.4.+")
latestDepTestLibrary("io.projectreactor:reactor-test:3.4.+")
}
Loading