Add instrumentation for Guava ListenableFutures context propagation (DataDog/dd-trace-java#1665)
MODIFY
instrumentation/cassandra/cassandra-3.0/cassandra-3.0.gradle
instrumentation/cassandra/cassandra-3.0/cassandra-3.0.gradle
@@ -50,6 +50,7 @@ dependencies {
testLibrary group: 'com.datastax.cassandra', name: 'cassandra-driver-core', version: '3.2.0'
testImplementation group: 'org.cassandraunit', name: 'cassandra-unit', version: '3.1.3.2'
+ testImplementation project(':instrumentation:guava-10.0')
latestDepTestLibrary group: 'com.datastax.cassandra', name: 'cassandra-driver-core', version: '3.+'
}

MODIFY
instrumentation/cassandra/cassandra-3.0/src/test/groovy/CassandraClientTest.groovy
instrumentation/cassandra/cassandra-3.0/src/test/groovy/CassandraClientTest.groovy
@@ -24,6 +24,8 @@ import io.opentelemetry.auto.test.AgentTestRunner
import io.opentelemetry.auto.test.asserts.TraceAssert
import io.opentelemetry.sdk.trace.data.SpanData
import io.opentelemetry.trace.attributes.SemanticAttributes
+import java.util.concurrent.Executors
+import java.util.concurrent.atomic.AtomicBoolean
import org.cassandraunit.utils.EmbeddedCassandraServerHelper
import spock.lang.Shared
@@ -32,6 +34,9 @@ class CassandraClientTest extends AgentTestRunner {
@Shared
Cluster cluster
+ @Shared
+ def executor = Executors.newCachedThreadPool()
+
def setupSpec() {
/*
This timeout seems excessive but we've seen tests fail with timeout of 40s.
@@ -86,9 +91,15 @@ class CassandraClientTest extends AgentTestRunner {
def "test async"() {
setup:
+ def callbackExecuted = new AtomicBoolean()
Session session = cluster.connect(keyspace)
runUnderTrace("parent") {
- session.executeAsync(statement)
+ def future = session.executeAsync(statement)
+ future.addListener({ ->
+ runUnderTrace("callbackListener") {
+ callbackExecuted.set(true)
+ }
+ }, executor)
}
expect:
@@ -98,9 +109,10 @@ class CassandraClientTest extends AgentTestRunner {
cassandraSpan(it, 0, "USE $keyspace", null)
}
}
- trace(keyspace ? 1 : 0, 2) {
+ trace(keyspace ? 1 : 0, 3) {
basicSpan(it, 0, "parent")
cassandraSpan(it, 1, statement, keyspace, span(0))
+ basicSpan(it, 2, "callbackListener", span(0))
}
}

ADD
instrumentation/guava-10.0/guava-10.0.gradle
+apply from: "$rootDir/gradle/instrumentation.gradle"
+
+muzzle {
+ pass {
+ group = "com.google.guava"
+ module = "guava"
+ versions = "[10.0,]"
+ assertInverse = true
+ }
+}
+
+dependencies {
+ library group: 'com.google.guava', name: 'guava', version: '10.0'
+}

ADD
instrumentation/guava-10.0/src/main/java/io/opentelemetry/instrumentation/auto/guava/ListenableFutureInstrumentation.java
+/*
+ * Copyright The OpenTelemetry Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.opentelemetry.instrumentation.auto.guava;
+
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import com.google.auto.service.AutoService;
+import com.google.common.util.concurrent.AbstractFuture;
+import io.grpc.Context;
+import io.opentelemetry.instrumentation.auto.api.ContextStore;
+import io.opentelemetry.instrumentation.auto.api.InstrumentationContext;
+import io.opentelemetry.instrumentation.auto.api.concurrent.ExecutorInstrumentationUtils;
+import io.opentelemetry.instrumentation.auto.api.concurrent.RunnableWrapper;
+import io.opentelemetry.instrumentation.auto.api.concurrent.State;
+import io.opentelemetry.javaagent.tooling.Instrumenter;
+import java.util.Map;
+import java.util.concurrent.Executor;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import net.bytebuddy.matcher.ElementMatchers;
+
+@AutoService(Instrumenter.class)
+public class ListenableFutureInstrumentation extends Instrumenter.Default {
+
+ public ListenableFutureInstrumentation() {
+ super("guava");
+ }
+
+ @Override
+ public ElementMatcher<TypeDescription> typeMatcher() {
+ return named("com.google.common.util.concurrent.AbstractFuture");
+ }
+
+ @Override
+ public Map<String, String> contextStore() {
+ return singletonMap(Runnable.class.getName(), State.class.getName());
+ }
+
+ @Override
+ public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+ return singletonMap(
+ named("addListener").and(ElementMatchers.takesArguments(Runnable.class, Executor.class)),
+ ListenableFutureInstrumentation.class.getName() + "$AddListenerAdvice");
+ }
+
+ public static class AddListenerAdvice {
+ @Advice.OnMethodEnter(suppress = Throwable.class)
+ public static State addListenerEnter(
+ @Advice.Argument(value = 0, readOnly = false) Runnable task) {
+ final Context context = Context.current();
+ final Runnable newTask = RunnableWrapper.wrapIfNeeded(task);
+ if (ExecutorInstrumentationUtils.shouldAttachStateToTask(newTask)) {
+ task = newTask;
+ final ContextStore<Runnable, State> contextStore =
+ InstrumentationContext.get(Runnable.class, State.class);
+ return ExecutorInstrumentationUtils.setupState(contextStore, newTask, context);
+ }
+ return null;
+ }
+
+ @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+ public static void addListenerExit(
+ @Advice.Enter final State state, @Advice.Thrown final Throwable throwable) {
+ ExecutorInstrumentationUtils.cleanUpOnMethodExit(state, throwable);
+ }
+
+ private static void muzzleCheck(final AbstractFuture<?> future) {
+ future.addListener(null, null);
+ }
+ }
+}

ADD
instrumentation/guava-10.0/src/test/groovy/ListenableFutureTest.groovy
+/*
+ * Copyright The OpenTelemetry Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import com.google.common.util.concurrent.Futures
+import com.google.common.util.concurrent.ListenableFuture
+import com.google.common.util.concurrent.SettableFuture
+import io.opentelemetry.auto.test.base.AbstractPromiseTest
+import java.util.concurrent.Executors
+import spock.lang.Shared
+
+class ListenableFutureTest extends AbstractPromiseTest<SettableFuture<Boolean>, ListenableFuture<String>> {
+ @Shared
+ def executor = Executors.newFixedThreadPool(1)
+
+ @Override
+ SettableFuture<Boolean> newPromise() {
+ return SettableFuture.create()
+ }
+
+ @Override
+ ListenableFuture<String> map(SettableFuture<Boolean> promise, Closure<String> callback) {
+ return Futures.transform(promise, callback, executor)
+ }
+
+ @Override
+ void onComplete(ListenableFuture<String> promise, Closure callback) {
+ promise.addListener({ -> callback(promise.get()) }, executor)
+ }
+
+
+ @Override
+ void complete(SettableFuture<Boolean> promise, boolean value) {
+ promise.set(value)
+ }
+
+ @Override
+ Boolean get(SettableFuture<Boolean> promise) {
+ return promise.get()
+ }
+}

MODIFY
javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/AgentInstaller.java
javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/AgentInstaller.java
@@ -105,6 +105,7 @@ public class AgentInstaller {
new AgentBuilder.Default()
.disableClassFormatChanges()
.with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
+ .with(AgentBuilder.RedefinitionStrategy.DiscoveryStrategy.Reiterating.INSTANCE)
.with(AgentBuilder.DescriptionStrategy.Default.POOL_ONLY)
.with(AgentTooling.poolStrategy())
.with(new ClassLoadListener())

MODIFY
settings.gradle
settings.gradle
@@ -51,7 +51,6 @@ include ':utils:test-utils'
// smoke tests
include ':smoke-tests'
-// instrumentation:
include ':instrumentation:akka-context-propagation-2.5'
include ':instrumentation:akka-http-10.0'
include ':instrumentation:apache-httpasyncclient-4.0'
@@ -86,6 +85,7 @@ include ':instrumentation:google-http-client-1.19'
include ':instrumentation:grizzly-2.0'
include ':instrumentation:grizzly-client-1.9'
include ':instrumentation:grpc-1.5'
+include ':instrumentation:guava-10.0'
include ':instrumentation:hibernate:hibernate-3.3'
include ':instrumentation:hibernate:hibernate-4.0'
include ':instrumentation:hibernate:hibernate-4.3'

ADD
testing-common/src/main/groovy/io/opentelemetry/auto/test/base/AbstractPromiseTest.groovy
+/*
+ * Copyright The OpenTelemetry Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.opentelemetry.auto.test.base
+
+import static io.opentelemetry.auto.test.utils.TraceUtils.basicSpan
+import static io.opentelemetry.auto.test.utils.TraceUtils.runUnderTrace
+
+import io.opentelemetry.auto.test.AgentTestRunner
+
+// TODO: add a test for a longer chain of promises
+abstract class AbstractPromiseTest<P, M> extends AgentTestRunner {
+
+ abstract P newPromise()
+
+ abstract M map(P promise, Closure<String> callback)
+
+ abstract void onComplete(M promise, Closure callback)
+
+ abstract void complete(P promise, boolean value)
+
+ abstract Boolean get(P promise)
+
+ def "test call with parent"() {
+ setup:
+ def promise = newPromise()
+
+ when:
+ runUnderTrace("parent") {
+ def mapped = map(promise) { "$it" }
+ onComplete(mapped) {
+ assert it == "$value"
+ runUnderTrace("callback") {}
+ }
+ runUnderTrace("other") {
+ complete(promise, value)
+ }
+ }
+
+ then:
+ get(promise) == value
+ assertTraces(1) {
+ trace(0, 3) {
+ basicSpan(it, 0, "parent")
+ basicSpan(it, 1, "other", it.span(0))
+ basicSpan(it, 2, "callback", it.span(0))
+ }
+ }
+
+ where:
+ value << [true, false]
+ }
+
+ def "test call with parent delayed complete"() {
+ setup:
+ def promise = newPromise()
+
+ when:
+ runUnderTrace("parent") {
+ def mapped = map(promise) { "$it" }
+ onComplete(mapped) {
+ assert it == "$value"
+ runUnderTrace("callback") {}
+ }
+ }
+
+ runUnderTrace("other") {
+ complete(promise, value)
+ }
+
+ then:
+ get(promise) == value
+ assertTraces(2) {
+ trace(0, 2) {
+ basicSpan(it, 0, "parent")
+ basicSpan(it, 1, "callback", span(0))
+ }
+ trace(1, 1) {
+ basicSpan(it, 0, "other")
+ }
+ }
+
+ where:
+ value << [true, false]
+ }
+
+ def "test call with parent complete separate thread"() {
+ setup:
+ final promise = newPromise()
+
+ when:
+ runUnderTrace("parent") {
+ def mapped = map(promise) { "$it" }
+ onComplete(mapped) {
+ assert it == "$value"
+ runUnderTrace("callback") {}
+ }
+ Thread.start {
+ complete(promise, value)
+ }.join()
+ }
+
+ then:
+ get(promise) == value
+ assertTraces(1) {
+ trace(0, 2) {
+ basicSpan(it, 0, "parent")
+ basicSpan(it, 1, "callback", it.span(0))
+ }
+ }
+
+ where:
+ value << [true, false]
+ }
+
+ def "test call with no parent"() {
+ setup:
+ def promise = newPromise()
+
+ when:
+ def mapped = map(promise) { "$it" }
+ onComplete(mapped) {
+ assert it == "$value"
+ runUnderTrace("callback") {}
+ }
+
+ runUnderTrace("other") {
+ complete(promise, value)
+ }
+
+ then:
+ get(promise) == value
+ assertTraces(1) {
+ trace(0, 2) {
+ // TODO: is this really the behavior we want?
+ basicSpan(it, 0, "other")
+ basicSpan(it, 1, "callback", it.span(0))
+ }
+ }
+
+ where:
+ value << [true, false]
+ }
+}

Add instrumentation for Guava ListenableFutures context propagation (DataDog/dd-trace-java#1665)