make reactor 0.9 behave the same way 1.0 does, add request error test

open-telemetry · Apr 28, 2021 · 1f267aa · 1f267aa
1 parent 2b80e1f
commit 1f267aa
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 8 deletions.
diff --git a/...java/io/opentelemetry/javaagent/instrumentation/reactornetty/v0_9/DecoratorFunctions.java b/...java/io/opentelemetry/javaagent/instrumentation/reactornetty/v0_9/DecoratorFunctions.java
@@ -24,14 +24,23 @@ public static boolean shouldDecorate(Class<?> callbackClass) {
 
   private abstract static class OnMessageDecorator<M> implements BiConsumer<M, Connection> {
     private final BiConsumer<? super M, ? super Connection> delegate;
+    private final boolean forceParentContext;
 
-    public OnMessageDecorator(BiConsumer<? super M, ? super Connection> delegate) {
+    public OnMessageDecorator(BiConsumer<? super M, ? super Connection> delegate,
+        boolean forceParentContext) {
       this.delegate = delegate;
+      this.forceParentContext = forceParentContext;
     }
 
     @Override
     public final void accept(M message, Connection connection) {
-      Context context = getChannelContext(currentContext(message), connection.channel());
+      Channel channel = connection.channel();
+      // don't try to get the client span from the netty channel when forceParentSpan is true
+      // this way the parent context will always be propagated
+      if (forceParentContext) {
+        channel = null;
+      }
+      Context context = getChannelContext(currentContext(message), channel);
       if (context == null) {
         delegate.accept(message, connection);
       } else {
@@ -46,7 +55,7 @@ public final void accept(M message, Connection connection) {
 
   public static final class OnRequestDecorator extends OnMessageDecorator<HttpClientRequest> {
     public OnRequestDecorator(BiConsumer<? super HttpClientRequest, ? super Connection> delegate) {
-      super(delegate);
+      super(delegate, false);
     }
 
     @Override
@@ -57,8 +66,9 @@ reactor.util.context.Context currentContext(HttpClientRequest message) {
 
   public static final class OnResponseDecorator extends OnMessageDecorator<HttpClientResponse> {
     public OnResponseDecorator(
-        BiConsumer<? super HttpClientResponse, ? super Connection> delegate) {
-      super(delegate);
+        BiConsumer<? super HttpClientResponse, ? super Connection> delegate,
+        boolean forceParentContext) {
+      super(delegate, forceParentContext);
     }
 
     @Override

diff --git a/...emetry/javaagent/instrumentation/reactornetty/v0_9/ReactorNettyInstrumentationModule.java b/...emetry/javaagent/instrumentation/reactornetty/v0_9/ReactorNettyInstrumentationModule.java
@@ -171,9 +171,12 @@ 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) {
+            BiConsumer<? super HttpClientResponse, ? super Connection> callback,
+        @Advice.Origin("#m") String methodName) {
       if (DecoratorFunctions.shouldDecorate(callback.getClass())) {
-        callback = new DecoratorFunctions.OnResponseDecorator(callback);
+        boolean forceParentContext = methodName.equals("doAfterResponse");
+        callback = new DecoratorFunctions.OnResponseDecorator(callback,
+            forceParentContext);
       }
     }
   }

diff --git a/...tty/reactor-netty-0.9/javaagent/src/test/groovy/AbstractReactorNettyHttpClientTest.groovy b/...tty/reactor-netty-0.9/javaagent/src/test/groovy/AbstractReactorNettyHttpClientTest.groovy
@@ -3,10 +3,14 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+
+import static io.opentelemetry.instrumentation.test.utils.PortUtils.UNUSABLE_PORT
 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.api.trace.SpanKind
+import io.opentelemetry.api.trace.StatusCode
 import io.opentelemetry.instrumentation.test.AgentTestTrait
 import io.opentelemetry.instrumentation.test.base.HttpClientTest
 import io.opentelemetry.sdk.trace.data.SpanData
@@ -94,11 +98,48 @@ abstract class AbstractReactorNettyHttpClientTest extends HttpClientTest<HttpCli
         assertSameSpan(parentSpan, onRequestSpan)
         assertSameSpan(nettyClientSpan, afterRequestSpan)
         assertSameSpan(nettyClientSpan, onResponseSpan)
-        assertSameSpan(nettyClientSpan, afterResponseSpan)
+        assertSameSpan(parentSpan, afterResponseSpan)
+      }
+    }
+  }
+
+  def "should expose context to http request error callback"() {
+    given:
+    def onRequestErrorSpan = new AtomicReference<Span>()
+
+    def httpClient = createHttpClient()
+      .doOnRequestError({ rq, err -> onRequestErrorSpan.set(Span.current()) })
+
+    when:
+    runUnderTrace("parent") {
+      httpClient.get()
+        .uri("http://localhost:$UNUSABLE_PORT/")
+        .response()
+        .block()
+    }
+
+    then:
+    def ex = thrown(Exception)
+
+    assertTraces(1) {
+      trace(0, 2) {
+        def parentSpan = span(0)
+
+        basicSpan(it, 0, "parent", null, ex)
+        span(1) {
+          def actualException = ex.cause
+          kind SpanKind.CLIENT
+          childOf parentSpan
+          status StatusCode.ERROR
+          errorEvent(actualException.class, actualException.message)
+        }
+
+        assertSameSpan(parentSpan, onRequestErrorSpan)
       }
     }
   }
 
+
   private static void assertSameSpan(SpanData expected, AtomicReference<Span> actual) {
     def expectedSpanContext = expected.spanContext
     def actualSpanContext = actual.get().spanContext