Port Spring Boot WithSpanAspect to Instrumenter API (#3607)

* Start porting WithSpanAspect to use Instrumenter API

* Some cleanup and refactoring

* Switch caching dependency to compile only

* Minor refactors, javadocs

* Fix instrumentation name

* Rename builder methods

* spotless

* Add request type to extract Method and WithSpan annotation, use method references

* Add comment about IntelliJ dependency workaround

* Make cache non-configurable, use AsyncOperationEndSupport directly

* Address PR comments

* Move to static factory method, method-keyed cache

instrumentation-api-annotation-support/build.gradle.kts

@@ -9,6 +9,10 @@ group = "io.opentelemetry.instrumentation"
 dependencies {
   implementation(project(":instrumentation-api"))
 
+  // this only exists to make Intellij happy since it doesn't (currently at least) understand our
+  // inclusion of this artifact inside of :instrumentation-api
+  compileOnly(project(":instrumentation-api-caching"))
+
   api("io.opentelemetry:opentelemetry-api")
   api("io.opentelemetry:opentelemetry-semconv")
 

instrumentation-api-annotation-support/src/main/java/io/opentelemetry/instrumentation/api/annotation/support/MethodArgumentsExtractor.java

@@ -0,0 +1,13 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.annotation.support;
+
+/** Extractor for the actual arguments passed to the parameters of the traced method. */
+@FunctionalInterface
+public interface MethodArgumentsExtractor<REQUEST> {
+  /** Extracts an array of the actual arguments from the {@link REQUEST}. */
+  Object[] extract(REQUEST request);
+}

instrumentation-api-annotation-support/src/main/java/io/opentelemetry/instrumentation/api/annotation/support/MethodCache.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.annotation.support;
+
+import io.opentelemetry.instrumentation.api.caching.Cache;
+import java.lang.reflect.Method;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
+
+/**
+ * Implementation of {@link Cache} that uses {@link ClassValue} to store values keyed by {@link
+ * Method} compared by value equality while allowing the declaring class to be unloaded.
+ */
+final class MethodCache<V> extends ClassValue<Map<Method, V>> implements Cache<Method, V> {
+  @Override
+  public V computeIfAbsent(Method key, Function<? super Method, ? extends V> mappingFunction) {
+    return this.get(key.getDeclaringClass()).computeIfAbsent(key, mappingFunction);
+  }
+
+  @Override
+  public V get(Method key) {
+    return this.get(key.getDeclaringClass()).get(key);
+  }
+
+  @Override
+  public void put(Method key, V value) {
+    this.get(key.getDeclaringClass()).put(key, value);
+  }
+
+  @Override
+  public void remove(Method key) {
+    this.get(key.getDeclaringClass()).remove(key);
+  }
+
+  @Override
+  protected Map<Method, V> computeValue(Class<?> type) {
+    return new ConcurrentHashMap<>();
+  }
+}

instrumentation-api-annotation-support/src/main/java/io/opentelemetry/instrumentation/api/annotation/support/MethodExtractor.java

@@ -0,0 +1,16 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.annotation.support;
+
+import java.lang.reflect.Method;
+
+/** Extractor for the traced {@link Method}. */
+@FunctionalInterface
+public interface MethodExtractor<REQUEST> {
+
+  /** Extracts the {@link Method} corresponding to the {@link REQUEST}. */
+  Method extract(REQUEST request);
+}

instrumentation-api-annotation-support/src/main/java/io/opentelemetry/instrumentation/api/annotation/support/MethodSpanAttributesExtractor.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.annotation.support;
+
+import io.opentelemetry.api.common.AttributesBuilder;
+import io.opentelemetry.instrumentation.api.caching.Cache;
+import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
+import java.lang.reflect.Method;
+import java.lang.reflect.Parameter;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Extractor of {@link io.opentelemetry.api.common.Attributes} for a traced method. */
+public final class MethodSpanAttributesExtractor<REQUEST, RESPONSE>
+    extends AttributesExtractor<REQUEST, RESPONSE> {
+
+  private final BaseAttributeBinder binder;
+  private final MethodExtractor<REQUEST> methodExtractor;
+  private final MethodArgumentsExtractor<REQUEST> methodArgumentsExtractor;
+  private final Cache<Method, AttributeBindings> cache;
+
+  public static <REQUEST, RESPONSE> MethodSpanAttributesExtractor<REQUEST, RESPONSE> newInstance(
+      MethodExtractor<REQUEST> methodExtractor,
+      ParameterAttributeNamesExtractor parameterAttributeNamesExtractor,
+      MethodArgumentsExtractor<REQUEST> methodArgumentsExtractor) {
+
+    return new MethodSpanAttributesExtractor<>(
+        methodExtractor, parameterAttributeNamesExtractor, methodArgumentsExtractor);
+  }
+
+  MethodSpanAttributesExtractor(
+      MethodExtractor<REQUEST> methodExtractor,
+      ParameterAttributeNamesExtractor parameterAttributeNamesExtractor,
+      MethodArgumentsExtractor<REQUEST> methodArgumentsExtractor) {
+    this.methodExtractor = methodExtractor;
+    this.methodArgumentsExtractor = methodArgumentsExtractor;
+    this.binder = new MethodSpanAttributeBinder(parameterAttributeNamesExtractor);
+    this.cache = new MethodCache<>();
+  }
+
+  @Override
+  protected void onStart(AttributesBuilder attributes, REQUEST request) {
+    Method method = methodExtractor.extract(request);
+    AttributeBindings bindings = cache.computeIfAbsent(method, binder::bind);
+    if (!bindings.isEmpty()) {
+      Object[] args = methodArgumentsExtractor.extract(request);
+      bindings.apply(attributes::put, args);
+    }
+  }
+
+  @Override
+  protected void onEnd(
+      AttributesBuilder attributes, REQUEST request, @Nullable RESPONSE response) {}
+
+  private static class MethodSpanAttributeBinder extends BaseAttributeBinder {
+    private final ParameterAttributeNamesExtractor parameterAttributeNamesExtractor;
+
+    public MethodSpanAttributeBinder(
+        ParameterAttributeNamesExtractor parameterAttributeNamesExtractor) {
+      this.parameterAttributeNamesExtractor = parameterAttributeNamesExtractor;
+    }
+
+    @Override
+    protected @Nullable String[] attributeNamesForParameters(
+        Method method, Parameter[] parameters) {
+      return parameterAttributeNamesExtractor.extract(method, parameters);
+    }
+  }
+}

instrumentation-api-annotation-support/src/main/java/io/opentelemetry/instrumentation/api/annotation/support/ParameterAttributeNamesExtractor.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.api.annotation.support;
+
+import java.lang.reflect.Method;
+import java.lang.reflect.Parameter;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Extractor for the attribute names for the parameters of a traced method. */
+@FunctionalInterface
+public interface ParameterAttributeNamesExtractor {
+  /**
+   * Returns an array of the names of the attributes for the parameters of the traced method. The
+   * array should be the same length as the array of the method parameters. An element may be {@code
+   * null} to indicate that the parameter should not be bound to an attribute. The array may also be
+   * {@code null} to indicate that the method has no parameters to bind to attributes.
+   *
+   * @param method the traced method
+   * @param parameters the method parameters
+   * @return an array of the attribute names
+   */
+  @Nullable
+  String[] extract(Method method, Parameter[] parameters);
+}

instrumentation-api-annotation-support/src/main/java/io/opentelemetry/instrumentation/api/annotation/support/async/AsyncOperationEndSupport.java

@@ -7,6 +7,7 @@ package io.opentelemetry.instrumentation.api.annotation.support.async;
 
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
  * A wrapper over {@link Instrumenter} that is able to defer {@link Instrumenter#end(Context,
@@ -35,13 +36,13 @@ public final class AsyncOperationEndSupport<REQUEST, RESPONSE> {
   private final Instrumenter<REQUEST, RESPONSE> instrumenter;
   private final Class<RESPONSE> responseType;
   private final Class<?> asyncType;
-  private final AsyncOperationEndStrategy asyncOperationEndStrategy;
+  private final @Nullable AsyncOperationEndStrategy asyncOperationEndStrategy;
 
   private AsyncOperationEndSupport(
       Instrumenter<REQUEST, RESPONSE> instrumenter,
       Class<RESPONSE> responseType,
       Class<?> asyncType,
-      AsyncOperationEndStrategy asyncOperationEndStrategy) {
+      @Nullable AsyncOperationEndStrategy asyncOperationEndStrategy) {
     this.instrumenter = instrumenter;
     this.responseType = responseType;
     this.asyncType = asyncType;
@@ -61,8 +62,10 @@ public final class AsyncOperationEndSupport<REQUEST, RESPONSE> {
    * won't be {@link Instrumenter#end(Context, Object, Object, Throwable) ended} until {@code
    * asyncValue} completes.
    */
+  @SuppressWarnings("unchecked")
+  @Nullable
   public <ASYNC> ASYNC asyncEnd(
-      Context context, REQUEST request, ASYNC asyncValue, Throwable throwable) {
+      Context context, REQUEST request, @Nullable ASYNC asyncValue, @Nullable Throwable throwable) {
     // we can end early if an exception was thrown
     if (throwable != null) {
       instrumenter.end(context, request, null, throwable);

instrumentation/spring/spring-boot-autoconfigure/src/main/java/io/opentelemetry/instrumentation/spring/autoconfigure/aspects/JoinPointRequest.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.spring.autoconfigure.aspects;
+
+import io.opentelemetry.extension.annotations.WithSpan;
+import java.lang.reflect.Method;
+import org.aspectj.lang.JoinPoint;
+import org.aspectj.lang.reflect.MethodSignature;
+
+final class JoinPointRequest {
+  private final JoinPoint joinPoint;
+  private final Method method;
+  private final WithSpan annotation;
+
+  JoinPointRequest(JoinPoint joinPoint) {
+    this.joinPoint = joinPoint;
+    MethodSignature methodSignature = (MethodSignature) joinPoint.getSignature();
+    this.method = methodSignature.getMethod();
+    this.annotation = this.method.getDeclaredAnnotation(WithSpan.class);
+  }
+
+  Method method() {
+    return method;
+  }
+
+  WithSpan annotation() {
+    return annotation;
+  }
+
+  Object[] args() {
+    return joinPoint.getArgs();
+  }
+}

instrumentation/spring/spring-boot-autoconfigure/src/main/java/io/opentelemetry/instrumentation/spring/autoconfigure/aspects/TraceAspectAutoConfiguration.java

@@ -30,15 +30,9 @@ public class TraceAspectAutoConfiguration {
     return new DefaultParameterNameDiscoverer();
   }
 
-  @Bean
-  public WithSpanAspectAttributeBinder withSpanAspectAttributeBinder(
-      ParameterNameDiscoverer parameterNameDiscoverer) {
-    return new WithSpanAspectAttributeBinder(parameterNameDiscoverer);
-  }
-
   @Bean
   public WithSpanAspect withSpanAspect(
-      OpenTelemetry openTelemetry, WithSpanAspectAttributeBinder withSpanAspectAttributeBinder) {
+      OpenTelemetry openTelemetry, ParameterNameDiscoverer parameterNameDiscoverer) {
-    return new WithSpanAspect(openTelemetry, withSpanAspectAttributeBinder);
+    return new WithSpanAspect(openTelemetry, parameterNameDiscoverer);
   }
 }

instrumentation/spring/spring-boot-autoconfigure/src/main/java/io/opentelemetry/instrumentation/spring/autoconfigure/aspects/WithSpanAspect.java

@@ -6,14 +6,19 @@
 package io.opentelemetry.instrumentation.spring.autoconfigure.aspects;
 
 import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.extension.annotations.WithSpan;
-import java.lang.reflect.Method;
+import io.opentelemetry.instrumentation.api.annotation.support.MethodSpanAttributesExtractor;
+import io.opentelemetry.instrumentation.api.annotation.support.ParameterAttributeNamesExtractor;
+import io.opentelemetry.instrumentation.api.annotation.support.async.AsyncOperationEndSupport;
+import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+import io.opentelemetry.instrumentation.api.tracer.SpanNames;
 import org.aspectj.lang.ProceedingJoinPoint;
 import org.aspectj.lang.annotation.Around;
 import org.aspectj.lang.annotation.Aspect;
-import org.aspectj.lang.reflect.MethodSignature;
+import org.springframework.core.ParameterNameDiscoverer;
 
 /**
  * Uses Spring-AOP to wrap methods marked by {@link WithSpan} in a {@link
@@ -27,30 +32,57 @@ import org.aspectj.lang.reflect.MethodSignature;
  */
 @Aspect
 public class WithSpanAspect {
-  private final WithSpanAspectTracer tracer;
+  private static final String INSTRUMENTATION_NAME = "io.opentelemetry.spring-boot-autoconfigure";
+
+  private final Instrumenter<JoinPointRequest, Object> instrumenter;
 
   public WithSpanAspect(
-      OpenTelemetry openTelemetry, WithSpanAspectAttributeBinder withSpanAspectAttributeBinder) {
+      OpenTelemetry openTelemetry, ParameterNameDiscoverer parameterNameDiscoverer) {
-    tracer = new WithSpanAspectTracer(openTelemetry, withSpanAspectAttributeBinder);
+
+    ParameterAttributeNamesExtractor parameterAttributeNamesExtractor =
+        new WithSpanAspectParameterAttributeNamesExtractor(parameterNameDiscoverer);
+
+    instrumenter =
+        Instrumenter.newBuilder(openTelemetry, INSTRUMENTATION_NAME, WithSpanAspect::spanName)
+            .addAttributesExtractor(
+                MethodSpanAttributesExtractor.newInstance(
+                    JoinPointRequest::method,
+                    parameterAttributeNamesExtractor,
+                    JoinPointRequest::args))
+            .newInstrumenter(WithSpanAspect::spanKind);
+  }
+
+  private static String spanName(JoinPointRequest request) {
+    WithSpan annotation = request.annotation();
+    String spanName = annotation.value();
+    if (spanName.isEmpty()) {
+      return SpanNames.fromMethod(request.method());
+    }
+    return spanName;
+  }
+
+  private static SpanKind spanKind(JoinPointRequest request) {
+    return request.annotation().kind();
   }
 
   @Around("@annotation(io.opentelemetry.extension.annotations.WithSpan)")
   public Object traceMethod(ProceedingJoinPoint pjp) throws Throwable {
-    MethodSignature signature = (MethodSignature) pjp.getSignature();
-    Method method = signature.getMethod();
-    WithSpan withSpan = method.getAnnotation(WithSpan.class);
 
+    JoinPointRequest request = new JoinPointRequest(pjp);
     Context parentContext = Context.current();
-    if (!tracer.shouldStartSpan(parentContext, withSpan.kind())) {
+    if (!instrumenter.shouldStart(parentContext, request)) {
       return pjp.proceed();
     }
 
-    Context context = tracer.startSpan(parentContext, withSpan, method, pjp);
+    Context context = instrumenter.start(parentContext, request);
+    AsyncOperationEndSupport<JoinPointRequest, Object> asyncOperationEndSupport =
+        AsyncOperationEndSupport.create(
+            instrumenter, Object.class, request.method().getReturnType());
     try (Scope ignored = context.makeCurrent()) {
-      Object result = pjp.proceed();
+      Object response = pjp.proceed();
-      return tracer.end(context, method.getReturnType(), result);
+      return asyncOperationEndSupport.asyncEnd(context, request, response, null);
     } catch (Throwable t) {
-      tracer.endExceptionally(context, t);
+      asyncOperationEndSupport.asyncEnd(context, request, null, t);
       throw t;
     }
   }

instrumentation/spring/spring-boot-autoconfigure/src/main/java/io/opentelemetry/instrumentation/spring/autoconfigure/aspects/WithSpanAspectParameterAttributeNamesExtractor.java

@@ -6,32 +6,22 @@
 package io.opentelemetry.instrumentation.spring.autoconfigure.aspects;
 
 import io.opentelemetry.extension.annotations.SpanAttribute;
-import io.opentelemetry.instrumentation.api.annotation.support.AttributeBindings;
+import io.opentelemetry.instrumentation.api.annotation.support.ParameterAttributeNamesExtractor;
-import io.opentelemetry.instrumentation.api.annotation.support.BaseAttributeBinder;
-import io.opentelemetry.instrumentation.api.caching.Cache;
 import java.lang.reflect.Method;
 import java.lang.reflect.Parameter;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.springframework.core.ParameterNameDiscoverer;
 
-public class WithSpanAspectAttributeBinder extends BaseAttributeBinder {
+class WithSpanAspectParameterAttributeNamesExtractor implements ParameterAttributeNamesExtractor {
-
-  private static final Cache<Method, AttributeBindings> bindings =
-      Cache.newBuilder().setWeakKeys().build();
-
   private final ParameterNameDiscoverer parameterNameDiscoverer;
 
-  public WithSpanAspectAttributeBinder(ParameterNameDiscoverer parameterNameDiscoverer) {
+  public WithSpanAspectParameterAttributeNamesExtractor(
+      ParameterNameDiscoverer parameterNameDiscoverer) {
     this.parameterNameDiscoverer = parameterNameDiscoverer;
   }
 
   @Override
-  public AttributeBindings bind(Method method) {
+  public @Nullable String[] extract(Method method, Parameter[] parameters) {
-    return bindings.computeIfAbsent(method, super::bind);
-  }
-
-  @Override
-  protected @Nullable String[] attributeNamesForParameters(Method method, Parameter[] parameters) {
     String[] parameterNames = parameterNameDiscoverer.getParameterNames(method);
     String[] attributeNames = new String[parameters.length];
 

instrumentation/spring/spring-boot-autoconfigure/src/main/java/io/opentelemetry/instrumentation/spring/autoconfigure/aspects/WithSpanAspectTracer.java

@@ -1,92 +0,0 @@
-/*
- * Copyright The OpenTelemetry Authors
- * SPDX-License-Identifier: Apache-2.0
- */
-
-package io.opentelemetry.instrumentation.spring.autoconfigure.aspects;
-
-import io.opentelemetry.api.OpenTelemetry;
-import io.opentelemetry.api.trace.Span;
-import io.opentelemetry.api.trace.SpanBuilder;
-import io.opentelemetry.context.Context;
-import io.opentelemetry.extension.annotations.WithSpan;
-import io.opentelemetry.instrumentation.api.annotation.support.AttributeBindings;
-import io.opentelemetry.instrumentation.api.tracer.BaseTracer;
-import io.opentelemetry.instrumentation.api.tracer.SpanNames;
-import io.opentelemetry.instrumentation.api.tracer.async.AsyncSpanEndStrategies;
-import io.opentelemetry.instrumentation.api.tracer.async.AsyncSpanEndStrategy;
-import java.lang.reflect.Method;
-import org.aspectj.lang.JoinPoint;
-
-class WithSpanAspectTracer extends BaseTracer {
-
-  private final WithSpanAspectAttributeBinder withSpanAspectAttributeBinder;
-  private final AsyncSpanEndStrategies asyncSpanEndStrategies =
-      AsyncSpanEndStrategies.getInstance();
-
-  WithSpanAspectTracer(
-      OpenTelemetry openTelemetry, WithSpanAspectAttributeBinder withSpanAspectAttributeBinder) {
-    super(openTelemetry);
-    this.withSpanAspectAttributeBinder = withSpanAspectAttributeBinder;
-  }
-
-  @Override
-  protected String getInstrumentationName() {
-    return "io.opentelemetry.spring-boot-autoconfigure-aspect";
-  }
-
-  Context startSpan(
-      Context parentContext, WithSpan annotation, Method method, JoinPoint joinPoint) {
-    SpanBuilder spanBuilder =
-        spanBuilder(parentContext, spanName(annotation, method), annotation.kind());
-    Span span = withSpanAttributes(spanBuilder, method, joinPoint).startSpan();
-    switch (annotation.kind()) {
-      case SERVER:
-        return withServerSpan(parentContext, span);
-      case CLIENT:
-        return withClientSpan(parentContext, span);
-      default:
-        return parentContext.with(span);
-    }
-  }
-
-  private static String spanName(WithSpan annotation, Method method) {
-    String spanName = annotation.value();
-    if (spanName.isEmpty()) {
-      return SpanNames.fromMethod(method);
-    }
-    return spanName;
-  }
-
-  public SpanBuilder withSpanAttributes(
-      SpanBuilder spanBuilder, Method method, JoinPoint joinPoint) {
-    AttributeBindings bindings = withSpanAspectAttributeBinder.bind(method);
-    if (!bindings.isEmpty()) {
-      bindings.apply(spanBuilder::setAttribute, joinPoint.getArgs());
-    }
-    return spanBuilder;
-  }
-
-  /**
-   * 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.
-   */
-  public Object end(Context context, Class<?> returnType, Object returnValue) {
-    if (returnType.isInstance(returnValue)) {
-      AsyncSpanEndStrategy asyncSpanEndStrategy =
-          asyncSpanEndStrategies.resolveStrategy(returnType);
-      if (asyncSpanEndStrategy != null) {
-        return asyncSpanEndStrategy.end(this, context, returnValue);
-      }
-    }
-    end(context);
-    return returnValue;
-  }
-}

instrumentation/spring/spring-boot-autoconfigure/src/test/java/io/opentelemetry/instrumentation/spring/autoconfigure/aspects/WithSpanAspectTest.java

@@ -104,9 +104,8 @@ public class WithSpanAspectTest {
             return null;
           }
         };
-    WithSpanAspectAttributeBinder attributeBinder =
+
-        new WithSpanAspectAttributeBinder(parameterNameDiscoverer);
+    WithSpanAspect aspect = new WithSpanAspect(testing.getOpenTelemetry(), parameterNameDiscoverer);
-    WithSpanAspect aspect = new WithSpanAspect(testing.getOpenTelemetry(), attributeBinder);
     factory.addAspect(aspect);
 
     withSpanTester = factory.getProxy();