Async @WithSpan Instrumentation for Guava ListenableFuture (#2811)

* Add Guava instrumentation library with AsyncSpanEndStrategy

* Enable span strategy in advice

* Spotless

* Nix attempt at typeInitializer advice, leave TODO comment to revisit

* Move async span strategy registration to helper class

* Remove use of sameThreadExecutor

* Make helper class final and add comment about relying on static initializer

instrumentation/guava-10.0/javaagent/guava-10.0-javaagent.gradle

@@ -11,4 +11,8 @@ muzzle {
 
 dependencies {
   library group: 'com.google.guava', name: 'guava', version: '10.0'
+
+  implementation project(':instrumentation:guava-10.0:library')
+
+  testImplementation deps.opentelemetryExtAnnotations
 }

instrumentation/guava-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/guava/GuavaInstrumentationModule.java

@@ -6,7 +6,7 @@
 package io.opentelemetry.javaagent.instrumentation.guava;
 
 import static java.util.Collections.singletonList;
-import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 
 import com.google.auto.service.AutoService;
@@ -20,6 +20,7 @@ import io.opentelemetry.javaagent.instrumentation.api.concurrent.RunnableWrapper
 import io.opentelemetry.javaagent.instrumentation.api.concurrent.State;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Executor;
@@ -49,9 +50,20 @@ public class GuavaInstrumentationModule extends InstrumentationModule {
 
     @Override
     public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
-      return singletonMap(
+      Map<ElementMatcher<? super MethodDescription>, String> map = new HashMap<>();
+      map.put(
+          isConstructor(), GuavaInstrumentationModule.class.getName() + "$AbstractFutureAdvice");
+      map.put(
           named("addListener").and(ElementMatchers.takesArguments(Runnable.class, Executor.class)),
           GuavaInstrumentationModule.class.getName() + "$AddListenerAdvice");
+      return map;
+    }
+  }
+
+  public static class AbstractFutureAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onConstruction() {
+      InstrumentationHelper.initialize();
     }
   }
 

instrumentation/guava-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/guava/InstrumentationHelper.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.guava;
+
+import io.opentelemetry.instrumentation.api.tracer.async.AsyncSpanEndStrategies;
+import io.opentelemetry.instrumentation.guava.GuavaAsyncSpanEndStrategy;
+
+public final class InstrumentationHelper {
+  static {
+    AsyncSpanEndStrategies.getInstance().registerStrategy(GuavaAsyncSpanEndStrategy.INSTANCE);
+  }
+
+  /**
+   * This method is invoked to trigger the runtime system to execute the static initializer block
+   * ensuring that the {@link GuavaAsyncSpanEndStrategy} is registered exactly once.
+   */
+  public static void initialize() {}
+
+  private InstrumentationHelper() {}
+}

instrumentation/guava-10.0/javaagent/src/test/groovy/GuavaWithSpanInstrumentationTest.groovy

@@ -0,0 +1,105 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import com.google.common.util.concurrent.Futures
+import com.google.common.util.concurrent.SettableFuture
+import io.opentelemetry.api.trace.SpanKind
+import io.opentelemetry.instrumentation.guava.TracedWithSpan
+import io.opentelemetry.instrumentation.test.AgentInstrumentationSpecification
+
+class GuavaWithSpanInstrumentationTest extends AgentInstrumentationSpecification {
+
+  def "should capture span for already done ListenableFuture"() {
+    setup:
+    new TracedWithSpan().listenableFuture(Futures.immediateFuture("Value"))
+
+    expect:
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          name "TracedWithSpan.listenableFuture"
+          kind SpanKind.INTERNAL
+          hasNoParent()
+          errored false
+          attributes {
+          }
+        }
+      }
+    }
+  }
+
+  def "should capture span for already failed ListenableFuture"() {
+    setup:
+    def error = new IllegalArgumentException("Boom")
+    new TracedWithSpan().listenableFuture(Futures.immediateFailedFuture(error))
+
+    expect:
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          name "TracedWithSpan.listenableFuture"
+          kind SpanKind.INTERNAL
+          hasNoParent()
+          errored true
+          errorEvent(IllegalArgumentException, "Boom")
+          attributes {
+          }
+        }
+      }
+    }
+  }
+
+  def "should capture span for eventually done ListenableFuture"() {
+    setup:
+    def future = SettableFuture.<String>create()
+    new TracedWithSpan().listenableFuture(future)
+
+    expect:
+    Thread.sleep(500) // sleep a bit just to make sure no span is captured
+    assertTraces(0) {}
+
+    future.set("Value")
+
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          name "TracedWithSpan.listenableFuture"
+          kind SpanKind.INTERNAL
+          hasNoParent()
+          errored false
+          attributes {
+          }
+        }
+      }
+    }
+  }
+
+  def "should capture span for eventually failed ListenableFuture"() {
+    setup:
+    def error = new IllegalArgumentException("Boom")
+    def future = SettableFuture.<String>create()
+    new TracedWithSpan().listenableFuture(future)
+
+    expect:
+    Thread.sleep(500) // sleep a bit just to make sure no span is captured
+    assertTraces(0) {}
+
+    future.setException(error)
+
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          name "TracedWithSpan.listenableFuture"
+          kind SpanKind.INTERNAL
+          hasNoParent()
+          errored true
+          errorEvent(IllegalArgumentException, "Boom")
+          attributes {
+          }
+        }
+      }
+    }
+  }
+}

instrumentation/guava-10.0/javaagent/src/test/java/io/opentelemetry/instrumentation/guava/TracedWithSpan.java

@@ -0,0 +1,16 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.guava;
+
+import com.google.common.util.concurrent.ListenableFuture;
+import io.opentelemetry.extension.annotations.WithSpan;
+
+public class TracedWithSpan {
+  @WithSpan
+  public ListenableFuture<String> listenableFuture(ListenableFuture<String> future) {
+    return future;
+  }
+}

instrumentation/guava-10.0/library/guava-10.0-library.gradle

@@ -0,0 +1,5 @@
+apply from: "$rootDir/gradle/instrumentation-library.gradle"
+
+dependencies {
+  library group: 'com.google.guava', name: 'guava', version: '10.0'
+}

instrumentation/guava-10.0/library/src/main/java/io/opentelemetry/instrumentation/guava/GuavaAsyncSpanEndStrategy.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.guava;
+
+import com.google.common.util.concurrent.ListenableFuture;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.tracer.BaseTracer;
+import io.opentelemetry.instrumentation.api.tracer.async.AsyncSpanEndStrategy;
+
+public enum GuavaAsyncSpanEndStrategy implements AsyncSpanEndStrategy {
+  INSTANCE;
+
+  @Override
+  public boolean supports(Class<?> returnType) {
+    return ListenableFuture.class.isAssignableFrom(returnType);
+  }
+
+  @Override
+  public Object end(BaseTracer tracer, Context context, Object returnValue) {
+    ListenableFuture<?> future = (ListenableFuture<?>) returnValue;
+    if (future.isDone()) {
+      endSpan(tracer, context, future);
+    } else {
+      future.addListener(() -> endSpan(tracer, context, future), Runnable::run);
+    }
+    return future;
+  }
+
+  private void endSpan(BaseTracer tracer, Context context, ListenableFuture<?> future) {
+    try {
+      future.get();
+      tracer.end(context);
+    } catch (Throwable exception) {
+      tracer.endExceptionally(context, exception);
+    }
+  }
+}

instrumentation/guava-10.0/library/src/test/groovy/GuavaAsyncSpanEndStrategyTest.groovy

@@ -0,0 +1,89 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import com.google.common.util.concurrent.Futures
+import com.google.common.util.concurrent.ListenableFuture
+import com.google.common.util.concurrent.SettableFuture
+import io.opentelemetry.context.Context
+import io.opentelemetry.instrumentation.api.tracer.BaseTracer
+import io.opentelemetry.instrumentation.guava.GuavaAsyncSpanEndStrategy
+import spock.lang.Specification
+
+class GuavaAsyncSpanEndStrategyTest extends Specification  {
+  BaseTracer tracer
+
+  Context context
+
+  def underTest = GuavaAsyncSpanEndStrategy.INSTANCE
+
+  void setup() {
+    tracer = Mock()
+    context = Mock()
+  }
+
+  def "ListenableFuture is supported"() {
+    expect:
+    underTest.supports(ListenableFuture)
+  }
+
+  def "SettableFuture is also supported"() {
+    expect:
+    underTest.supports(SettableFuture)
+  }
+
+  def "ends span on already done future"() {
+    when:
+    underTest.end(tracer, context, Futures.immediateFuture("Value"))
+
+    then:
+    1 * tracer.end(context)
+  }
+
+  def "ends span on already failed future"() {
+    given:
+    def exception = new IllegalStateException()
+
+    when:
+    underTest.end(tracer, context, Futures.immediateFailedFuture(exception))
+
+    then:
+    1 * tracer.endExceptionally(context, { it.getCause() == exception })
+  }
+
+  def "ends span on eventually done future"() {
+    given:
+    def future = SettableFuture.<String>create()
+
+    when:
+    underTest.end(tracer, context, future)
+
+    then:
+    0 * tracer._
+
+    when:
+    future.set("Value")
+
+    then:
+    1 * tracer.end(context)
+  }
+
+  def "ends span on eventually failed future"() {
+    given:
+    def future = SettableFuture.<String>create()
+    def exception = new IllegalStateException()
+
+    when:
+    underTest.end(tracer, context, future)
+
+    then:
+    0 * tracer._
+
+    when:
+    future.setException(exception)
+
+    then:
+    1 * tracer.endExceptionally(context, { it.getCause() == exception })
+  }
+}

settings.gradle

@@ -118,6 +118,7 @@ include ':instrumentation:grpc-1.5:javaagent'
 include ':instrumentation:grpc-1.5:library'
 include ':instrumentation:grpc-1.5:testing'
 include ':instrumentation:guava-10.0:javaagent'
+include ':instrumentation:guava-10.0:library'
 include ':instrumentation:gwt-2.0:javaagent'
 include ':instrumentation:hibernate:hibernate-3.3:javaagent'
 include ':instrumentation:hibernate:hibernate-4.0:javaagent'