Skip to content
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.

Sign up for GitHub

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

Use new reactor contextWrite when available (from reactor 3.4.0) #7538

Merged
merged 11 commits into from
Jan 12, 2023
Merged

Use new reactor contextWrite when available (from reactor 3.4.0) #7538

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
0bfb698
Use new contextWrite when available
trask Jan 10, 2023
bc08b73
In tests too
trask Jan 10, 2023
49259e3
Revert "Add muzzle limits to fix CI (#7106)"
trask Jan 10, 2023
8f79487
Webflux
trask Jan 10, 2023
12450e3
switch
trask Jan 11, 2023
f1cf68c
lambda
trask Jan 11, 2023
8840f13
Looking ahead
trask Jan 11, 2023
311936b
Merge remote-tracking branch 'upstream/main' into project-reactor-3.5
trask Jan 11, 2023
949b222
Update docs
trask Jan 11, 2023
c116838
Test against WebFlux 6
trask Jan 11, 2023
9232dac
spotless
trask Jan 12, 2023
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
2 changes: 1 addition & 1 deletion instrumentation/reactor/reactor-3.1/javaagent/build.gradle.kts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ muzzle {
pass {
group.set("io.projectreactor")
module.set("reactor-core")
versions.set("[3.1.0.RELEASE,3.5.0)")
versions.set("[3.1.0.RELEASE,)")
extraDependency("io.opentelemetry:opentelemetry-api:1.0.0")
assertInverse.set(true)
excludeInstrumentationName("opentelemetry-api")
Expand Down
70 changes: 58 additions & 12 deletions ...etry/javaagent/instrumentation/reactor/ContextPropagationOperatorInstrumentationTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
package io.opentelemetry.javaagent.instrumentation.reactor;

import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.attributeEntry;
import static java.lang.invoke.MethodType.methodType;
import static org.assertj.core.api.Assertions.assertThat;

import io.opentelemetry.api.trace.Span;
Expand All @@ -14,7 +15,11 @@
import io.opentelemetry.instrumentation.reactor.ContextPropagationOperator;
import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.time.Duration;
import java.util.function.Function;
import javax.annotation.Nullable;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
import reactor.core.publisher.Flux;
Expand All @@ -26,6 +31,25 @@ class ContextPropagationOperatorInstrumentationTest {
@RegisterExtension
static final InstrumentationExtension testing = AgentInstrumentationExtension.create();

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

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

@Test
void storeAndGetContext() {
reactor.util.context.Context reactorContext = reactor.util.context.Context.empty();
Expand Down Expand Up @@ -136,18 +160,25 @@ void storeContextForcesItToBecomeCurrent() {
Span span =
testing.getOpenTelemetry().getTracer("test").spanBuilder("parent").startSpan();

return Mono.delay(Duration.ofMillis(1))
.flatMap(
unused -> {
// usual trick to force this to run under new TracingSubscriber with context
// written in the next call
return new ExtensionAnnotationsTracedWithSpan().mono(Mono.just("Value"));
})
.subscriberContext(
ctx ->
ContextPropagationOperator.storeOpenTelemetryContext(
ctx, Context.current().with(span)))
.doFinally(unused -> span.end());
Mono<String> interim =
Mono.delay(Duration.ofMillis(1))
.flatMap(
unused -> {
// usual trick to force this to run under new TracingSubscriber with
// context
// written in the next call
return new ExtensionAnnotationsTracedWithSpan()
.mono(Mono.just("Value"));
});
try {
interim =
(Mono<String>)
MONO_CONTEXT_WRITE_METHOD.invoke(
interim, new StoreOpenTelemetryContext(Context.current().with(span)));
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the previous version of this code, wasn't Context.current().with(span) executed as a part of the lambda? I mean it was lazy before, now it's eager

} catch (Throwable e) {
throw new RuntimeException(e);
}
return interim.doFinally(unused -> span.end());
});

StepVerifier.create(result).expectNext("Value").verifyComplete();
Expand All @@ -161,4 +192,19 @@ void storeContextForcesItToBecomeCurrent() {
.hasKind(SpanKind.INTERNAL)
.hasParent(trace.getSpan(0))));
}

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 ContextPropagationOperator.storeOpenTelemetryContext(context, tracingContext);
}
}
}
74 changes: 68 additions & 6 deletions ...ry/src/main/java/io/opentelemetry/instrumentation/reactor/ContextPropagationOperator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,13 @@

package io.opentelemetry.instrumentation.reactor;

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 javax.annotation.Nullable;
Expand All @@ -42,6 +46,28 @@ public final class ContextPropagationOperator {

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

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

public static ContextPropagationOperator create() {
return builder().build();
}
Expand Down Expand Up @@ -133,29 +159,65 @@ public void resetOnEachOperator() {
}

/** Forces Mono to run in traceContext scope. */
@SuppressWarnings("unchecked")
public static <T> Mono<T> runWithContext(Mono<T> publisher, Context tracingContext) {
if (!enabled) {
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
return ScalarPropagatingMono.create(publisher)
.subscriberContext(ctx -> storeOpenTelemetryContext(ctx, tracingContext));
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) {
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
return ScalarPropagatingFlux.create(publisher)
.subscriberContext(ctx -> storeOpenTelemetryContext(ctx, tracingContext));
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);
}
}

public static class Lifter<T>
Expand Down
Loading