Okhttp3: fix concurrency test with callback (#3669)

instrumentation/okhttp/okhttp-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/okhttp/v3_0/OkHttp3DispatcherInstrumentation.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.okhttp.v3_0;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
+import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
+import io.opentelemetry.javaagent.instrumentation.api.concurrent.ExecutorInstrumentationUtils;
+import io.opentelemetry.javaagent.instrumentation.api.concurrent.State;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class OkHttp3DispatcherInstrumentation implements TypeInstrumentation {
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("okhttp3.Dispatcher");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        named("enqueue").and(takesArgument(0, named("okhttp3.RealCall$AsyncCall"))),
+        OkHttp3DispatcherInstrumentation.class.getName() + "$AttachStateAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class AttachStateAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static State onEnter(@Advice.Argument(0) Runnable call) {
+      if (ExecutorInstrumentationUtils.shouldAttachStateToTask(call)) {
+        ContextStore<Runnable, State> contextStore =
+            InstrumentationContext.get(Runnable.class, State.class);
+        return ExecutorInstrumentationUtils.setupState(
+            contextStore, call, Java8BytecodeBridge.currentContext());
+      }
+      return null;
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(@Advice.Enter State state, @Advice.Thrown Throwable throwable) {
+      ExecutorInstrumentationUtils.cleanUpOnMethodExit(state, throwable);
+    }
+  }
+}

instrumentation/okhttp/okhttp-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/okhttp/v3_0/OkHttp3InstrumentationModule.java

@@ -5,7 +5,7 @@
 
 package io.opentelemetry.javaagent.instrumentation.okhttp.v3_0;
 
-import static java.util.Collections.singletonList;
+import static java.util.Arrays.asList;
 
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
@@ -21,6 +21,6 @@ public class OkHttp3InstrumentationModule extends InstrumentationModule {
 
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
-    return singletonList(new OkHttp3Instrumentation());
+    return asList(new OkHttp3Instrumentation(), new OkHttp3DispatcherInstrumentation());
   }
 }

instrumentation/okhttp/okhttp-3.0/library/src/test/groovy/io/opentelemetry/instrumentation/okhttp/v3_0/OkHttp3Test.groovy

@@ -24,4 +24,9 @@ class OkHttp3Test extends AbstractOkHttp3Test implements LibraryTestTrait {
   boolean testWithClientParent() {
     false
   }
+
+  @Override
+  boolean testCausalityWithCallback() {
+    false
+  }
 }

instrumentation/okhttp/okhttp-3.0/testing/src/main/groovy/io/opentelemetry/instrumentation/okhttp/v3_0/AbstractOkHttp3Test.groovy

@@ -63,12 +63,7 @@ abstract class AbstractOkHttp3Test extends HttpClientTest<Request> {
   }
 
   @Override
-  boolean testRedirects() {
-    false
-  }
-
-  @Override
-  boolean testCausality() {
+  boolean testCircularRedirects() {
     false
   }
 

testing-common/src/main/groovy/io/opentelemetry/instrumentation/test/base/HttpClientTest.groovy

@@ -865,6 +865,7 @@ abstract class HttpClientTest<REQUEST> extends InstrumentationSpecification {
   def "high concurrency test with callback"() {
     setup:
     assumeTrue(testCausality())
+    assumeTrue(testCausalityWithCallback())
     assumeTrue(testCallback())
     assumeTrue(testCallbackWithParent())
 
@@ -1120,6 +1121,10 @@ abstract class HttpClientTest<REQUEST> extends InstrumentationSpecification {
     true
   }
 
+  boolean testCausalityWithCallback() {
+    true
+  }
+
   boolean testCallback() {
     return true
   }