Resolve end strategy after WithSpan method instead of before. (#5756)

* WIP * Resolve end strategy after WithSpan method instead of before.
open-telemetry · Apr 7, 2022 · cd725ec · cd725ec
1 parent 8b29f80
commit cd725ec
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 6 deletions.
diff --git a/...a/io/opentelemetry/javaagent/instrumentation/otelannotations/WithSpanInstrumentation.java b/...a/io/opentelemetry/javaagent/instrumentation/otelannotations/WithSpanInstrumentation.java
@@ -117,8 +117,6 @@ public static class WithSpanAdvice {
     public static void onEnter(
         @Advice.Origin Method originMethod,
         @Advice.Local("otelMethod") Method method,
-        @Advice.Local("otelOperationEndSupport")
-            AsyncOperationEndSupport<Method, Object> operationEndSupport,
         @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope) {
 
@@ -132,16 +130,12 @@ public static void onEnter(
       if (instrumenter.shouldStart(current, method)) {
         context = instrumenter.start(current, method);
         scope = context.makeCurrent();
-        operationEndSupport =
-            AsyncOperationEndSupport.create(instrumenter, Object.class, method.getReturnType());
       }
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void stopSpan(
         @Advice.Local("otelMethod") Method method,
-        @Advice.Local("otelOperationEndSupport")
-            AsyncOperationEndSupport<Method, Object> operationEndSupport,
         @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope,
         @Advice.Return(typing = Assigner.Typing.DYNAMIC, readOnly = false) Object returnValue,
@@ -150,6 +144,9 @@ public static void stopSpan(
         return;
       }
       scope.close();
+
+      AsyncOperationEndSupport<Method, Object> operationEndSupport =
+          AsyncOperationEndSupport.create(instrumenter(), Object.class, method.getReturnType());
       returnValue = operationEndSupport.asyncEnd(context, method, returnValue, throwable);
     }
   }

diff --git a/instrumentation/reactor/reactor-3.1/javaagent/build.gradle.kts b/instrumentation/reactor/reactor-3.1/javaagent/build.gradle.kts
@@ -34,3 +34,21 @@ dependencies {
   // Looks like later versions on reactor need this dependency for some reason even though it is marked as optional.
   latestDepTestLibrary("io.micrometer:micrometer-core:1.+")
 }
+
+testing {
+  suites {
+    val testInitialization by registering(JvmTestSuite::class) {
+      dependencies {
+        implementation(project(":instrumentation:reactor:reactor-3.1:library"))
+        implementation("io.opentelemetry:opentelemetry-extension-annotations")
+        implementation("io.projectreactor:reactor-test:3.1.0.RELEASE")
+      }
+    }
+  }
+}
+
+tasks {
+  check {
+    dependsOn(testing.suites)
+  }
+}
diff --git a/...alization/java/io/opentelemetry/javaagent/instrumentation/reactor/InitializationTest.java b/...alization/java/io/opentelemetry/javaagent/instrumentation/reactor/InitializationTest.java
@@ -0,0 +1,49 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.reactor;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.opentelemetry.extension.annotations.WithSpan;
+import java.util.stream.Collectors;
+import org.junit.jupiter.api.Test;
+import reactor.core.Scannable;
+import reactor.core.publisher.Mono;
+
+// Isolated test to use clean classloader because reactor instrumentation is applied on static
+// initialization.
+class InitializationTest {
+
+  @Test
+  void contextPropagated() {
+    Mono<String> mono = new Traced().traceMe();
+
+    // If reactor augmentation of WithSpan is working correctly, we will end up with these
+    // implementation details.
+    // TODO(anuraaga): This should just check actual context propagation instead of implementation
+    // but couldn't figure out how.
+    assertThat(((Scannable) mono).parents().collect(Collectors.toList()))
+        .anySatisfy(
+            op -> {
+              assertThat(op.getClass().getSimpleName()).isEqualTo("MonoFlatMap");
+              assertThat(op)
+                  .extracting("source")
+                  .satisfies(
+                      source ->
+                          assertThat(source.getClass().getSimpleName())
+                              .isEqualTo("ScalarPropagatingMono"));
+            });
+
+    assertThat(mono.block()).isEqualTo("foo");
+  }
+
+  static class Traced {
+    @WithSpan
+    Mono<String> traceMe() {
+      return Mono.just("foo");
+    }
+  }
+}