Add asynchronous tracing for Java 8 CompletableFuture in WithSpanAdvice

instrumentation/opentelemetry-annotations-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/otelannotations/WithSpanAdvice.java

@@ -13,6 +13,7 @@
 import io.opentelemetry.context.Scope;
 import java.lang.reflect.Method;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.implementation.bytecode.assign.Assigner;
 
 /**
  * Instrumentation for methods annotated with {@link WithSpan} annotation.
@@ -40,8 +41,10 @@ public static void onEnter(
 
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
   public static void stopSpan(
+      @Advice.Origin Method method,
       @Advice.Local("otelContext") Context context,
       @Advice.Local("otelScope") Scope scope,
+      @Advice.Return(typing = Assigner.Typing.DYNAMIC) Object returnValue,
       @Advice.Thrown Throwable throwable) {
     if (scope == null) {
       return;
@@ -51,7 +54,7 @@ public static void stopSpan(
     if (throwable != null) {
       tracer().endExceptionally(context, throwable);
     } else {
-      tracer().end(context);
+      tracer().end(context, method, returnValue);
     }
   }
 }

instrumentation/opentelemetry-annotations-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/otelannotations/WithSpanTracer.java

@@ -10,6 +10,7 @@
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.tracer.BaseTracer;
+import io.opentelemetry.javaagent.instrumentation.otelannotations.async.MethodSpanStrategies;
 import java.lang.reflect.Method;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -23,6 +24,8 @@ public static WithSpanTracer tracer() {
 
   private static final Logger log = LoggerFactory.getLogger(WithSpanTracer.class);
 
+  private final MethodSpanStrategies methodSpanStrategies = MethodSpanStrategies.getInstance();
+
   public Context startSpan(
       Context parentContext, WithSpan applicationAnnotation, Method method, SpanKind kind) {
     Span span =
@@ -69,6 +72,21 @@ public static SpanKind toAgentOrNull(
     }
   }
 
+  /**
+   * Denotes the end of the invocation of the traced method with a successful result which will end
+   * the span stored in the passed {@code context}. If the method returned a value representing an
+   * asynchronous operation then the span will remain open until the asynchronous operation has
+   * completed.
+   *
+   * @param result Return value from the traced method.
+   * @return Either {@code result} or a value composing over {@code result} for notification of
+   *     completion.
+   */
+  public Object end(Context context, Method method, Object result) {
+    Class<?> returnType = method.getReturnType();
+    return methodSpanStrategies.resolveStrategy(returnType).end(this, context, returnType, result);
+  }
+
   @Override
   protected String getInstrumentationName() {
     return "io.opentelemetry.javaagent.opentelemetry-annotations-1.0";

instrumentation/opentelemetry-annotations-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/otelannotations/async/Jdk8MethodStrategy.java

@@ -0,0 +1,59 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.otelannotations.async;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.tracer.BaseTracer;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+
+enum Jdk8MethodStrategy implements MethodSpanStrategy {
+  INSTANCE;
+
+  @Override
+  public boolean supports(Class<?> returnType) {
+    return returnType == CompletionStage.class || returnType == CompletableFuture.class;
+  }
+
+  @Override
+  public Object end(BaseTracer tracer, Context context, Class<?> returnType, Object result) {
+    if (result instanceof CompletableFuture) {
+      CompletableFuture<?> future = (CompletableFuture<?>) result;
+      if (future.isDone()) {
+        return endSynchronously(future, tracer, context);
+      }
+      return endWhenComplete(future, tracer, context);
+    } else if (result instanceof CompletionStage) {
+      CompletionStage<?> stage = (CompletionStage<?>) result;
+      return endWhenComplete(stage, tracer, context);
+    }
+    tracer.end(context);
+    return result;
+  }
+
+  private CompletableFuture<?> endSynchronously(
+      CompletableFuture<?> future, BaseTracer tracer, Context context) {
+    try {
+      future.join();
+      tracer.end(context);
+    } catch (Exception exception) {
+      tracer.endExceptionally(context, exception);
+    }
+    return future;
+  }
+
+  private CompletionStage<?> endWhenComplete(
+      CompletionStage<?> stage, BaseTracer tracer, Context context) {
+    return stage.whenComplete(
+        (result, exception) -> {
+          if (exception != null) {
+            tracer.endExceptionally(context, exception);
+          } else {
+            tracer.end(context);
+          }
+        });
+  }
+}

instrumentation/opentelemetry-annotations-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/otelannotations/async/MethodSpanStrategies.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.otelannotations.async;
+
+import java.util.List;
+import java.util.Objects;
+import java.util.concurrent.CopyOnWriteArrayList;
+
+/**
+ * Registry of {@link MethodSpanStrategy} implementations for tracing the asynchronous operations
+ * represented by the return type of a traced method.
+ */
+public class MethodSpanStrategies {
+  private static final MethodSpanStrategies instance = new MethodSpanStrategies();
+
+  public static MethodSpanStrategies getInstance() {
+    return instance;
+  }
+
+  private final List<MethodSpanStrategy> strategies = new CopyOnWriteArrayList<>();
+
+  private MethodSpanStrategies() {
+    strategies.add(Jdk8MethodStrategy.INSTANCE);
+  }
+
+  public void registerStrategy(MethodSpanStrategy strategy) {
+    Objects.requireNonNull(strategy);
+    strategies.add(strategy);
+  }
+
+  public MethodSpanStrategy resolveStrategy(Class<?> returnType) {
+    for (MethodSpanStrategy strategy : strategies) {
+      if (strategy.supports(returnType)) {
+        return strategy;
+      }
+    }
+    return MethodSpanStrategy.synchronous();
+  }
+}

instrumentation/opentelemetry-annotations-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/otelannotations/async/MethodSpanStrategy.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.otelannotations.async;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.tracer.BaseTracer;
+
+/**
+ * Represents an implementation of a strategy for composing over the return value of a traced
+ * method. If the return value represents the result of an asynchronous operation the implementation
+ * can compose or register for notification of completion at which point the span representing the
+ * invocation of the method will be ended.
+ */
+public interface MethodSpanStrategy {
+  boolean supports(Class<?> returnType);
+
+  /**
+   * Denotes the end of the invocation of the traced method with a successful result which will end
+   * the span stored in the passed {@code context}. If the method returned a value representing an
+   * asynchronous operation then the span will remain open until the asynchronous operation has
+   * completed.
+   *
+   * @param tracer {@link BaseTracer} tracer to be used to end the span stored in the {@code
+   *     context}.
+   * @param result Return value of the traced method.
+   * @return Either {@code result} or a value composing over {@code result} for notification of
+   *     completion.
+   */
+  Object end(BaseTracer tracer, Context context, Class<?> returnType, Object result);
+
+  /**
+   * Returns a {@link MethodSpanStrategy} for tracing synchronous methods where the return value
+   * does not represent the completion of an asynchronous operation.
+   */
+  static MethodSpanStrategy synchronous() {
+    return SynchronousMethodSpanStrategy.INSTANCE;
+  }
+}

instrumentation/opentelemetry-annotations-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/otelannotations/async/SynchronousMethodSpanStrategy.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.otelannotations.async;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.tracer.BaseTracer;
+
+enum SynchronousMethodSpanStrategy implements MethodSpanStrategy {
+  INSTANCE;
+
+  @Override
+  public boolean supports(Class<?> returnType) {
+    return true;
+  }
+
+  @Override
+  public Object end(BaseTracer tracer, Context context, Class<?> returnType, Object result) {
+    tracer.end(context);
+    return result;
+  }
+}

instrumentation/opentelemetry-annotations-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/otelannotations/async/package-info.java

@@ -0,0 +1,5 @@
+/**
+ * Provides implementations of strategies for tracing methods that return asynchronous and reactive
+ * values so that the span can be ended when the asynchronous operation completes.
+ */
+package io.opentelemetry.javaagent.instrumentation.otelannotations.async;

instrumentation/opentelemetry-annotations-1.0/javaagent/src/test/groovy/WithSpanInstrumentationTest.groovy

@@ -3,6 +3,8 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import java.util.concurrent.CompletableFuture
+
 import static io.opentelemetry.api.trace.SpanKind.CLIENT
 import static io.opentelemetry.api.trace.SpanKind.PRODUCER
 import static io.opentelemetry.api.trace.SpanKind.SERVER
@@ -144,4 +146,102 @@ class WithSpanInstrumentationTest extends AgentInstrumentationSpecification {
     Thread.sleep(500) // sleep a bit just to make sure no span is captured
     assertTraces(0) {}
   }
+
+  def "should capture span for CompletionStage"() {
+    setup:
+    def future = new CompletableFuture<String>()
+    new TracedWithSpan().completionStage(future)
+
+    expect:
+    Thread.sleep(500) // sleep a bit just to make sure no span is captured
+    assertTraces(0) {}
+
+    future.complete("Done")
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          name "TracedWithSpan.completionStage"
+          kind SpanKind.INTERNAL
+          hasNoParent()
+          errored false
+          attributes {
+          }
+        }
+      }
+    }
+  }
+
+  def "should capture span for CompletionStage on completed exceptionally"() {
+    setup:
+    def future = new CompletableFuture<String>()
+    new TracedWithSpan().completionStage(future)
+
+    expect:
+    Thread.sleep(500) // sleep a bit just to make sure no span is captured
+    assertTraces(0) {}
+
+    future.completeExceptionally(new IllegalArgumentException("Boom"))
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          name "TracedWithSpan.completionStage"
+          kind SpanKind.INTERNAL
+          hasNoParent()
+          errored true
+          errorEvent(IllegalArgumentException, "Boom")
+          attributes {
+          }
+        }
+      }
+    }
+  }
+
+  def "should capture span for CompletableFuture"() {
+    setup:
+    def future = new CompletableFuture<String>()
+    new TracedWithSpan().completableFuture(future)
+
+    expect:
+    Thread.sleep(500) // sleep a bit just to make sure no span is captured
+    assertTraces(0) {}
+
+    future.complete("Done")
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          name "TracedWithSpan.completableFuture"
+          kind SpanKind.INTERNAL
+          hasNoParent()
+          errored false
+          attributes {
+          }
+        }
+      }
+    }
+  }
+
+  def "should capture span for CompletableFuture on completed exceptionally"() {
+    setup:
+    def future = new CompletableFuture<String>()
+    new TracedWithSpan().completableFuture(future)
+
+    expect:
+    Thread.sleep(500) // sleep a bit just to make sure no span is captured
+    assertTraces(0) {}
+
+    future.completeExceptionally(new IllegalArgumentException("Boom"))
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          name "TracedWithSpan.completableFuture"
+          kind SpanKind.INTERNAL
+          hasNoParent()
+          errored true
+          errorEvent(IllegalArgumentException, "Boom")
+          attributes {
+          }
+        }
+      }
+    }
+  }
 }

instrumentation/opentelemetry-annotations-1.0/javaagent/src/test/java/io/opentelemetry/test/annotation/TracedWithSpan.java

@@ -7,6 +7,8 @@
 
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.extension.annotations.WithSpan;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
 
 public class TracedWithSpan {
 
@@ -54,4 +56,14 @@ public String nestedClients() {
   public String innerClient() {
     return "hello!";
   }
+
+  @WithSpan
+  public CompletionStage<String> completionStage(CompletableFuture<String> future) {
+    return future;
+  }
+
+  @WithSpan
+  public CompletableFuture<String> completableFuture(CompletableFuture<String> future) {
+    return future;
+  }
 }