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.getReturnType(), 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,26 @@ 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 not be finished until the asynchronous operation has
+   * completed.
+   *
+   * @param returnType Return type of the traced method.
+   * @param returnValue Return value from the traced method.
+   * @return Either {@code returnValue} or a value composing over {@code returnValue} for
+   *     notification of completion.
+   * @throws ClassCastException if returnValue is not an instance of returnType
+   */
+  public Object end(Context context, Class<?> returnType, Object returnValue) {
+    if (!returnType.isInstance(returnValue)) {
+      end(context);
+      return returnValue;
+    }
+    return methodSpanStrategies.resolveStrategy(returnType).end(this, context, returnValue);
+  }
+
   @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,75 @@
+/*
+ * 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, Object returnValue) {
+    if (returnValue instanceof CompletableFuture) {
+      CompletableFuture<?> future = (CompletableFuture<?>) returnValue;
+      if (endSynchronously(future, tracer, context)) {
+        return future;
+      }
+      return endWhenComplete(future, tracer, context);
+    }
+    CompletionStage<?> stage = (CompletionStage<?>) returnValue;
+    return endWhenComplete(stage, tracer, context);
+  }
+
+  /**
+   * Checks to see if the {@link CompletableFuture} has already been completed and if so
+   * synchronously ends the span to avoid additional allocations and overhead registering for
+   * notification of completion.
+   */
+  private boolean endSynchronously(
+      CompletableFuture<?> future, BaseTracer tracer, Context context) {
+
+    if (!future.isDone()) {
+      return false;
+    }
+
+    if (future.isCompletedExceptionally()) {
+      // If the future completed exceptionally then join to catch the exception
+      // so that it can be recorded to the span
+      try {
+        future.join();
+      } catch (Exception exception) {
+        tracer.endExceptionally(context, exception);
+        return true;
+      }
+    }
+    tracer.end(context);
+    return true;
+  }
+
+  /**
+   * Registers for notification of the completion of the {@link CompletionStage} at which time the
+   * span will be ended.
+   */
+  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,43 @@
+/*
+ * 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 returnValue Return value from the traced method. Must be an instance of a {@code
+   *     returnType} for which {@link #supports(Class)} returned true (in particular it must not be
+   *     {@code null}).
+   * @return Either {@code returnValue} or a value composing over {@code returnValue} for
+   *     notification of completion.
+   */
+  Object end(BaseTracer tracer, Context context, Object returnValue);
+
+  /**
+   * 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, Object returnValue) {
+    tracer.end(context);
+    return returnValue;
+  }
+}

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;