fix bug: get return type in a wrong way for method instrumentation (#6118)

* fix bug: get return type in a wrong way for method instrumentation

* add testcase for method instrumentation

* update code style

* update code style

* Use Span.isRecording() to test whether span has been ended

Co-authored-by: Lauri Tulmin <ltulmin@splunk.com>

instrumentation/methods/javaagent/build.gradle.kts

@@ -14,5 +14,5 @@ dependencies {
 }
 
 tasks.withType<Test>().configureEach {
-  jvmArgs("-Dotel.instrumentation.methods.include=io.opentelemetry.javaagent.instrumentation.methods.MethodTest\$ConfigTracedCallable[call]")
+  jvmArgs("-Dotel.instrumentation.methods.include=io.opentelemetry.javaagent.instrumentation.methods.MethodTest\$ConfigTracedCallable[call];io.opentelemetry.javaagent.instrumentation.methods.MethodTest\$ConfigTracedCompletableFuture[getResult]")
 }

instrumentation/methods/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/methods/MethodInstrumentation.java

@@ -18,6 +18,7 @@ import io.opentelemetry.instrumentation.api.annotation.support.async.AsyncOperat
 import io.opentelemetry.instrumentation.api.util.ClassAndMethod;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import java.lang.reflect.Method;
 import java.util.Set;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
@@ -72,7 +73,7 @@ public class MethodInstrumentation implements TypeInstrumentation {
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void stopSpan(
-        @Advice.Origin("#r") Class<?> returnType,
+        @Advice.Origin Method method,
         @Advice.Local("otelMethod") ClassAndMethod classAndMethod,
         @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope,
@@ -81,7 +82,7 @@ public class MethodInstrumentation implements TypeInstrumentation {
       scope.close();
 
       returnValue =
-          AsyncOperationEndSupport.create(instrumenter(), Void.class, returnType)
+          AsyncOperationEndSupport.create(instrumenter(), Void.class, method.getReturnType())
               .asyncEnd(context, classAndMethod, returnValue, throwable);
     }
   }

instrumentation/methods/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/methods/MethodTest.java

@@ -8,10 +8,14 @@ package io.opentelemetry.javaagent.instrumentation.methods;
 import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.assertThat;
 
 import io.opentelemetry.api.common.Attributes;
+import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
 import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
 import java.util.concurrent.Callable;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
@@ -23,7 +27,6 @@ class MethodTest {
   @Test
   void methodTraced() {
     assertThat(new ConfigTracedCallable().call()).isEqualTo("Hello!");
-
     testing.waitAndAssertTraces(
         trace ->
             trace.hasSpansSatisfyingExactly(
@@ -34,9 +37,52 @@ class MethodTest {
   }
 
   static class ConfigTracedCallable implements Callable<String> {
+
     @Override
     public String call() {
       return "Hello!";
     }
   }
+
+  @Test
+  void methodTracedWithAsyncStop() throws Exception {
+    ConfigTracedCompletableFuture traced = new ConfigTracedCompletableFuture();
+    CompletableFuture<String> future = traced.getResult();
+
+    // span is ended when CompletableFuture is completed
+    // verify that span has not been ended yet
+    assertThat(traced.span).isNotNull().satisfies(span -> assertThat(span.isRecording()).isTrue());
+
+    traced.countDownLatch.countDown();
+    assertThat(future.get(10, TimeUnit.SECONDS)).isEqualTo("Hello!");
+
+    testing.waitAndAssertTraces(
+        trace ->
+            trace.hasSpansSatisfyingExactly(
+                span ->
+                    span.hasName("ConfigTracedCompletableFuture.getResult")
+                        .hasKind(SpanKind.INTERNAL)
+                        .hasAttributes(Attributes.empty())));
+  }
+
+  static class ConfigTracedCompletableFuture {
+    final CountDownLatch countDownLatch = new CountDownLatch(1);
+    Span span;
+
+    CompletableFuture<String> getResult() {
+      CompletableFuture<String> completableFuture = new CompletableFuture<>();
+      span = Span.current();
+      new Thread(
+              () -> {
+                try {
+                  countDownLatch.await();
+                } catch (InterruptedException exception) {
+                  // ignore
+                }
+                completableFuture.complete("Hello!");
+              })
+          .start();
+      return completableFuture;
+    }
+  }
 }