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.
@@ -42,6 +43,7 @@ public static void onEnter(
  public static void stopSpan(
  @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 +53,7 @@ public static void stopSpan(
  if (throwable != null) {
  tracer().endExceptionally(context, throwable);
  } else {
- tracer().end(context);
+ tracer().end(context, returnValue);
  }
  }
 }

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

@@ -10,6 +10,8 @@
 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 io.opentelemetry.javaagent.instrumentation.otelannotations.async.MethodSpanStrategy;
 import java.lang.reflect.Method;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -23,19 +25,33 @@ 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) {
+
+ Context spanStrategyContext = withMethodSpanStrategy(parentContext, method);
  Span span =
  spanBuilder(
  parentContext, spanNameForMethodWithAnnotation(applicationAnnotation, method), kind)
  .startSpan();
  if (kind == SpanKind.SERVER) {
- return withServerSpan(parentContext, span);
+ return withServerSpan(spanStrategyContext, span);
  }
  if (kind == SpanKind.CLIENT) {
- return withClientSpan(parentContext, span);
+ return withClientSpan(spanStrategyContext, span);
  }
- return parentContext.with(span);
+ return spanStrategyContext.with(span);
+ }
+
+ /**
+ * Resolves the {@link MethodSpanStrategy} for tracing the specified {@code method} and stores
+ * that strategy in the returned {@code Context}.
+ */
+ protected Context withMethodSpanStrategy(Context context, Method method) {
+ Class<?> returnType = method.getReturnType();
+ MethodSpanStrategy methodSpanStrategy = methodSpanStrategies.resolveStrategy(returnType);
+ return context.with(methodSpanStrategy);
  }
 
  /**
@@ -69,6 +85,20 @@ 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, Object result) {
+ return MethodSpanStrategy.fromContext(context).end(this, context, 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/CompletableFutureMethodSpanStrategy.java

@@ -0,0 +1,51 @@
+/*
+ * 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;
+
+enum CompletableFutureMethodSpanStrategy implements MethodSpanStrategy {
+ INSTANCE;
+
+ @Override
+ public Object end(BaseTracer tracer, Context context, Object result) {
+ if (result instanceof CompletableFuture) {
+ CompletableFuture<?> future = (CompletableFuture<?>) result;
+ if (future.isDone()) {
+ return endSynchronously(tracer, context, future);
+ } else {
+ return endOnCompletion(tracer, context, future);
+ }
+ }
+ tracer.end(context);
+ return result;
+ }
+
+ private CompletableFuture<?> endSynchronously(
+ BaseTracer tracer, Context context, CompletableFuture<?> future) {
+ try {
+ future.join();
+ tracer.end(context);
+ } catch (Exception exception) {
+ tracer.endExceptionally(context, exception);
+ }
+ return future;
+ }
+
+ private CompletableFuture<?> endOnCompletion(
+ BaseTracer tracer, Context context, CompletableFuture<?> future) {
+ return future.whenComplete(
+ (value, error) -> {
+ if (error != null) {
+ tracer.endExceptionally(context, error);
+ } else {
+ tracer.end(context);
+ }
+ });
+ }
+}

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

@@ -0,0 +1,31 @@
+/*
+ * 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.CompletionStage;
+
+enum CompletionStageMethodSpanStrategy implements MethodSpanStrategy {
+ INSTANCE;
+
+ @Override
+ public Object end(BaseTracer tracer, Context context, Object result) {
+ if (result instanceof CompletionStage) {
+ CompletionStage<?> stage = (CompletionStage<?>) result;
+ return stage.whenComplete(
+ (value, error) -> {
+ if (error != null) {
+ tracer.endExceptionally(context, error);
+ } else {
+ tracer.end(context);
+ }
+ });
+ }
+ tracer.end(context);
+ return result;
+ }
+}

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

@@ -0,0 +1,44 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.otelannotations.async;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+
+/**
+ * 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;
+
+ static {
+ Map<Class<?>, MethodSpanStrategy> strategies = new HashMap<>();
+ strategies.put(CompletionStage.class, MethodSpanStrategy.forCompletionStage());
+ strategies.put(CompletableFuture.class, MethodSpanStrategy.forCompletableFuture());
+ instance = new MethodSpanStrategies(strategies);
+ }
+
+ public static MethodSpanStrategies getInstance() {
+ return instance;
+ }
+
+ private final Map<Class<?>, MethodSpanStrategy> strategies;
+
+ private MethodSpanStrategies(Map<Class<?>, MethodSpanStrategy> strategies) {
+ this.strategies = strategies;
+ }
+
+ public void registerStrategy(Class<?> returnType, MethodSpanStrategy strategy) {
+ strategies.put(returnType, strategy);
+ }
+
+ public MethodSpanStrategy resolveStrategy(Class<?> returnType) {
+ return strategies.getOrDefault(returnType, MethodSpanStrategy.synchronous());
+ }
+}

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

@@ -0,0 +1,71 @@
+/*
+ * 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.context.ContextKey;
+import io.opentelemetry.context.ImplicitContextKeyed;
+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 extends ImplicitContextKeyed {
+ ContextKey<MethodSpanStrategy> CONTEXT_KEY =
+ ContextKey.named("opentelemetry-spring-autoconfigure-aspects-method-span-strategy");
+
+ /**
+ * 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, Object result);
+
+ @Override
+ default Context storeInContext(Context context) {
+ return context.with(CONTEXT_KEY, this);
+ }
+
+ static MethodSpanStrategy fromContext(Context context) {
+ MethodSpanStrategy methodSpanStrategy = context.get(CONTEXT_KEY);
+ return methodSpanStrategy != null ? methodSpanStrategy : synchronous();
+ }
+
+ /**
+ * 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;
+ }
+
+ /**
+ * Returns a {@link MethodSpanStrategy} for tracing a method that returns a {@link
+ * java.util.concurrent.CompletionStage} representing the completion of an asynchronous operation.
+ */
+ static MethodSpanStrategy forCompletionStage() {
+ return CompletionStageMethodSpanStrategy.INSTANCE;
+ }
+
+ /**
+ * Returns a {@link MethodSpanStrategy} for tracing a method that returns a {@link
+ * java.util.concurrent.CompletableFuture} representing the completion of an asynchronous
+ * operation.
+ */
+ static MethodSpanStrategy forCompletableFuture() {
+ return CompletableFutureMethodSpanStrategy.INSTANCE;
+ }
+}

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

@@ -0,0 +1,19 @@
+/*
+ * 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 Object end(BaseTracer tracer, Context context, 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;