Skip to content

Commit

Permalink
code review comment + reactor-netty 0.9
Browse files Browse the repository at this point in the history
  • Loading branch information
Mateusz Rzeszutek committed Apr 29, 2021
1 parent 015a26d commit 444efb4
Show file tree
Hide file tree
Showing 4 changed files with 305 additions and 10 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,134 @@
/*
* Copyright The OpenTelemetry Authors
* SPDX-License-Identifier: Apache-2.0
*/

package io.opentelemetry.javaagent.instrumentation.reactornetty.v0_9;

import io.netty.channel.Channel;
import io.opentelemetry.context.Context;
import io.opentelemetry.context.Scope;
import io.opentelemetry.instrumentation.netty.v4_1.AttributeKeys;
import java.util.function.BiConsumer;
import org.checkerframework.checker.nullness.qual.Nullable;
import reactor.netty.Connection;
import reactor.netty.http.client.HttpClientRequest;
import reactor.netty.http.client.HttpClientResponse;

public final class DecoratorFunctions {

// ignore our own callbacks - or already decorated functions
public static boolean shouldDecorate(Class<?> callbackClass) {
return !callbackClass.getName().startsWith("io.opentelemetry.javaagent");
}

private abstract static class OnMessageDecorator<M> implements BiConsumer<M, Connection> {
private final BiConsumer<? super M, ? super Connection> delegate;

public OnMessageDecorator(BiConsumer<? super M, ? super Connection> delegate) {
this.delegate = delegate;
}

@Override
public final void accept(M message, Connection connection) {
Context context = getChannelContext(currentContext(message), connection.channel());
if (context == null) {
delegate.accept(message, connection);
} else {
try (Scope ignored = context.makeCurrent()) {
delegate.accept(message, connection);
}
}
}

abstract reactor.util.context.Context currentContext(M message);
}

public static final class OnRequestDecorator extends OnMessageDecorator<HttpClientRequest> {
public OnRequestDecorator(BiConsumer<? super HttpClientRequest, ? super Connection> delegate) {
super(delegate);
}

@Override
reactor.util.context.Context currentContext(HttpClientRequest message) {
return message.currentContext();
}
}

public static final class OnResponseDecorator extends OnMessageDecorator<HttpClientResponse> {
public OnResponseDecorator(
BiConsumer<? super HttpClientResponse, ? super Connection> delegate) {
super(delegate);
}

@Override
reactor.util.context.Context currentContext(HttpClientResponse message) {
return message.currentContext();
}
}

private abstract static class OnMessageErrorDecorator<M> implements BiConsumer<M, Throwable> {
private final BiConsumer<? super M, ? super Throwable> delegate;

public OnMessageErrorDecorator(BiConsumer<? super M, ? super Throwable> delegate) {
this.delegate = delegate;
}

@Override
public final void accept(M message, Throwable throwable) {
Context context = getChannelContext(currentContext(message), null);
if (context == null) {
delegate.accept(message, throwable);
} else {
try (Scope ignored = context.makeCurrent()) {
delegate.accept(message, throwable);
}
}
}

abstract reactor.util.context.Context currentContext(M message);
}

public static final class OnRequestErrorDecorator
extends OnMessageErrorDecorator<HttpClientRequest> {
public OnRequestErrorDecorator(
BiConsumer<? super HttpClientRequest, ? super Throwable> delegate) {
super(delegate);
}

@Override
reactor.util.context.Context currentContext(HttpClientRequest message) {
return message.currentContext();
}
}

public static final class OnResponseErrorDecorator
extends OnMessageErrorDecorator<HttpClientResponse> {
public OnResponseErrorDecorator(
BiConsumer<? super HttpClientResponse, ? super Throwable> delegate) {
super(delegate);
}

@Override
reactor.util.context.Context currentContext(HttpClientResponse message) {
return message.currentContext();
}
}

@Nullable
private static Context getChannelContext(
reactor.util.context.Context reactorContext, @Nullable Channel channel) {
// try to get the client span context from the channel if it's available
if (channel != null) {
Context context = channel.attr(AttributeKeys.CLIENT_CONTEXT).get();
if (context != null) {
return context;
}
}
// otherwise use the parent span context
return reactorContext.getOrDefault(
ReactorNettyInstrumentationModule.MapConnect.CONTEXT_ATTRIBUTE, null);
}

private DecoratorFunctions() {}
}
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,12 @@

import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.ClassLoaderMatcher.hasClassesNamed;
import static java.util.Collections.singletonList;
import static java.util.Collections.singletonMap;
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
import static net.bytebuddy.matcher.ElementMatchers.isStatic;
import static net.bytebuddy.matcher.ElementMatchers.named;
import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;

import com.google.auto.service.AutoService;
import io.netty.bootstrap.Bootstrap;
Expand All @@ -19,6 +21,7 @@
import io.opentelemetry.javaagent.instrumentation.api.CallDepthThreadLocalMap;
import io.opentelemetry.javaagent.tooling.InstrumentationModule;
import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;
Expand All @@ -31,6 +34,7 @@
import reactor.netty.Connection;
import reactor.netty.http.client.HttpClient;
import reactor.netty.http.client.HttpClientRequest;
import reactor.netty.http.client.HttpClientResponse;

/**
* This instrumentation solves the problem of the correct context propagation through the roller
Expand Down Expand Up @@ -64,9 +68,44 @@ public ElementMatcher<TypeDescription> typeMatcher() {

@Override
public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
return singletonMap(
Map<ElementMatcher.Junction<MethodDescription>, String> transformers = new HashMap<>();
transformers.put(
isStatic().and(namedOneOf("create", "newConnection", "from")),
ReactorNettyInstrumentationModule.class.getName() + "$CreateAdvice");

// advice classes below expose current context in doOn*/doAfter* callbacks
transformers.put(
isPublic()
.and(namedOneOf("doOnRequest", "doAfterRequest"))
.and(takesArguments(1))
.and(takesArgument(0, BiConsumer.class)),
ReactorNettyInstrumentationModule.class.getName() + "$OnRequestAdvice");
transformers.put(
isPublic()
.and(named("doOnRequestError"))
.and(takesArguments(1))
.and(takesArgument(0, BiConsumer.class)),
ReactorNettyInstrumentationModule.class.getName() + "$OnRequestErrorAdvice");
transformers.put(
isPublic()
.and(namedOneOf("doOnResponse", "doAfterResponse"))
.and(takesArguments(1))
.and(takesArgument(0, BiConsumer.class)),
ReactorNettyInstrumentationModule.class.getName() + "$OnResponseAdvice");
transformers.put(
isPublic()
.and(named("doOnResponseError"))
.and(takesArguments(1))
.and(takesArgument(0, BiConsumer.class)),
ReactorNettyInstrumentationModule.class.getName() + "$OnResponseErrorAdvice");
transformers.put(
isPublic()
.and(named("doOnError"))
.and(takesArguments(2))
.and(takesArgument(0, BiConsumer.class))
.and(takesArgument(1, BiConsumer.class)),
ReactorNettyInstrumentationModule.class.getName() + "$OnErrorAdvice");
return transformers;
}
}

Expand Down Expand Up @@ -105,4 +144,64 @@ public void accept(HttpClientRequest r, Connection c) {
c.channel().attr(AttributeKeys.WRITE_CONTEXT).set(context);
}
}

public static class OnRequestAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static void onEnter(
@Advice.Argument(value = 0, readOnly = false)
BiConsumer<? super HttpClientRequest, ? super Connection> callback) {
if (DecoratorFunctions.shouldDecorate(callback.getClass())) {
callback = new DecoratorFunctions.OnRequestDecorator(callback);
}
}
}

public static class OnRequestErrorAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static void onEnter(
@Advice.Argument(value = 0, readOnly = false)
BiConsumer<? super HttpClientRequest, ? super Throwable> callback) {
if (DecoratorFunctions.shouldDecorate(callback.getClass())) {
callback = new DecoratorFunctions.OnRequestErrorDecorator(callback);
}
}
}

public static class OnResponseAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static void onEnter(
@Advice.Argument(value = 0, readOnly = false)
BiConsumer<? super HttpClientResponse, ? super Connection> callback) {
if (DecoratorFunctions.shouldDecorate(callback.getClass())) {
callback = new DecoratorFunctions.OnResponseDecorator(callback);
}
}
}

public static class OnResponseErrorAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static void onEnter(
@Advice.Argument(value = 0, readOnly = false)
BiConsumer<? super HttpClientResponse, ? super Throwable> callback) {
if (DecoratorFunctions.shouldDecorate(callback.getClass())) {
callback = new DecoratorFunctions.OnResponseErrorDecorator(callback);
}
}
}

public static class OnErrorAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static void onEnter(
@Advice.Argument(value = 0, readOnly = false)
BiConsumer<? super HttpClientRequest, ? super Throwable> requestCallback,
@Advice.Argument(value = 1, readOnly = false)
BiConsumer<? super HttpClientResponse, ? super Throwable> responseCallback) {
if (DecoratorFunctions.shouldDecorate(requestCallback.getClass())) {
requestCallback = new DecoratorFunctions.OnRequestErrorDecorator(requestCallback);
}
if (DecoratorFunctions.shouldDecorate(responseCallback.getClass())) {
responseCallback = new DecoratorFunctions.OnResponseErrorDecorator(responseCallback);
}
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,14 @@
* SPDX-License-Identifier: Apache-2.0
*/

import static io.opentelemetry.instrumentation.test.utils.TraceUtils.basicSpan
import static io.opentelemetry.instrumentation.test.utils.TraceUtils.runUnderTrace

import io.opentelemetry.api.trace.Span
import io.opentelemetry.instrumentation.test.AgentTestTrait
import io.opentelemetry.instrumentation.test.base.HttpClientTest
import io.opentelemetry.sdk.trace.data.SpanData
import java.util.concurrent.atomic.AtomicReference
import reactor.netty.http.client.HttpClient

abstract class AbstractReactorNettyHttpClientTest extends HttpClientTest<HttpClient.ResponseReceiver> implements AgentTestTrait {
Expand Down Expand Up @@ -52,4 +58,51 @@ abstract class AbstractReactorNettyHttpClientTest extends HttpClientTest<HttpCli
}

abstract HttpClient createHttpClient()

def "should expose context to http client callbacks"() {
given:
def onRequestSpan = new AtomicReference<Span>()
def afterRequestSpan = new AtomicReference<Span>()
def onResponseSpan = new AtomicReference<Span>()
def afterResponseSpan = new AtomicReference<Span>()

def httpClient = createHttpClient()
.doOnRequest({ rq, con -> onRequestSpan.set(Span.current()) })
.doAfterRequest({ rq, con -> afterRequestSpan.set(Span.current()) })
.doOnResponse({ rs, con -> onResponseSpan.set(Span.current()) })
.doAfterResponse({ rs, con -> afterResponseSpan.set(Span.current()) })

when:
runUnderTrace("parent") {
httpClient.baseUrl(server.address.toString())
.get()
.uri("/success")
.response()
.block()
}

then:
assertTraces(1) {
trace(0, 3) {
def parentSpan = span(0)
def nettyClientSpan = span(1)

basicSpan(it, 0, "parent")
clientSpan(it, 1, parentSpan, "GET", server.address.resolve("/success"))
serverSpan(it, 2, nettyClientSpan)

assertSameSpan(parentSpan, onRequestSpan)
assertSameSpan(nettyClientSpan, afterRequestSpan)
assertSameSpan(nettyClientSpan, onResponseSpan)
assertSameSpan(nettyClientSpan, afterResponseSpan)
}
}
}

private static void assertSameSpan(SpanData expected, AtomicReference<Span> actual) {
def expectedSpanContext = expected.spanContext
def actualSpanContext = actual.get().spanContext
assert expectedSpanContext.traceId == actualSpanContext.traceId
assert expectedSpanContext.spanId == actualSpanContext.spanId
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,12 @@ public OnMessageDecorator(BiConsumer<? super M, ? super Connection> delegate) {
@Override
public void accept(M message, Connection connection) {
Context context = getChannelContext(message.currentContextView(), connection.channel());
try (Scope ignored = context.makeCurrent()) {
if (context == null) {
delegate.accept(message, connection);
} else {
try (Scope ignored = context.makeCurrent()) {
delegate.accept(message, connection);
}
}
}
}
Expand All @@ -50,23 +54,28 @@ public OnMessageErrorDecorator(BiConsumer<? super M, ? super Throwable> delegate
@Override
public void accept(M message, Throwable throwable) {
Context context = getChannelContext(message.currentContextView(), null);
try (Scope ignored = context.makeCurrent()) {
if (context == null) {
delegate.accept(message, throwable);
} else {
try (Scope ignored = context.makeCurrent()) {
delegate.accept(message, throwable);
}
}
}
}

@Nullable
private static Context getChannelContext(ContextView contextView, @Nullable Channel channel) {
Context context = null;
// try to get the client span context from the channel if it's available
if (channel != null) {
context = channel.attr(AttributeKeys.CLIENT_CONTEXT).get();
Context context = channel.attr(AttributeKeys.CLIENT_CONTEXT).get();
if (context != null) {
return context;
}
}
// otherwise use the parent span context
if (context == null) {
context = contextView.get(ReactorNettyInstrumentationModule.MapConnect.CONTEXT_ATTRIBUTE);
}
return context;
return contextView.getOrDefault(
ReactorNettyInstrumentationModule.MapConnect.CONTEXT_ATTRIBUTE, null);
}

private DecoratorFunctions() {}
Expand Down

0 comments on commit 444efb4

Please sign in to comment.