Make OkHttp3 instrumentation more reliable

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

@@ -6,20 +6,24 @@
 package io.opentelemetry.javaagent.instrumentation.okhttp.v3_0;
 
 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.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
+import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
+import okhttp3.Interceptor;
 import okhttp3.OkHttpClient;
 
 @AutoService(InstrumentationModule.class)
@@ -42,16 +46,65 @@ public ElementMatcher<TypeDescription> typeMatcher() {
 
     @Override
     public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
-      return singletonMap(
+      Map<ElementMatcher.Junction<MethodDescription>, String> transformers = new HashMap<>();
+      transformers.put(
           isConstructor().and(takesArgument(0, named("okhttp3.OkHttpClient$Builder"))),
-          OkHttp3InstrumentationModule.class.getName() + "$OkHttp3Advice");
+          OkHttp3InstrumentationModule.class.getName() + "$OkHttp3ClientConstructorAdvice");
+      transformers.put(
+          isMethod().and(named("newBuilder").and(returns(named("okhttp3.OkHttpClient$Builder")))),
+          OkHttp3InstrumentationModule.class.getName() + "$OkHttp3ClientNewBuilderAdvice");
+      return transformers;
     }
   }
 
-  public static class OkHttp3Advice {
+  // This advice makes two guarantees:
+  // 1) The state of the builder (specifically interceptor list) is the same before and after the
+  //    OkHttpClient constructor invocation
+  // 2) The interceptor list of the created OkHttpClient has exactly one instance of the tracing
+  //    interceptor and it is in the end (assuming no other instrumentations)
+  public static class OkHttp3ClientConstructorAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void addTracingInterceptor(@Advice.Argument(0) OkHttpClient.Builder builder) {
+    public static void addTracingInterceptor(
+        @Advice.Argument(0) OkHttpClient.Builder builder,
+        @Advice.Local("otelOriginalInterceptors") List<Interceptor> originalInterceptors) {
+
+      if (builder.interceptors().contains(OkHttp3Interceptors.TRACING_INTERCEPTOR)) {
+        // Potential corner case - the tracing interceptor may be in the builder due to the builder
+        // being manually constructed by adding interceptors from an existing client. In this case,
+        // save the original interceptors so we can restore them after the constructor call, and
+        // then remove all tracing interceptors before adding one as the last.
+        originalInterceptors = new ArrayList<>(builder.interceptors());
+
+        while (builder.interceptors().remove(OkHttp3Interceptors.TRACING_INTERCEPTOR)) {}
+      }
+
       builder.addInterceptor(OkHttp3Interceptors.TRACING_INTERCEPTOR);
     }
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void removeTracingInterceptor(
+        @Advice.Argument(0) OkHttpClient.Builder builder,
+        @Advice.Local("otelOriginalInterceptors") List<Interceptor> originalInterceptors) {
+
+      // Restore the interceptor list to what it was before the constructor call. In the common
+      // case, a single instance of the tracing interceptor was appended to the end, so it will be
+      // removed. For the corner case where builder already contained the tracing interceptor, the
+      // original interceptor list was saved to a local and will be restored from there.
+      if (originalInterceptors != null) {
+        builder.interceptors().clear();
+        builder.interceptors().addAll(originalInterceptors);
+      } else {
+        builder.interceptors().remove(OkHttp3Interceptors.TRACING_INTERCEPTOR);
+      }
+    }
+  }
+
+  public static class OkHttp3ClientNewBuilderAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void removeTracingInterceptor(@Advice.Return OkHttpClient.Builder builder) {
+      // Remove the interceptor from the builder returned by newBuilder, as it should be added only
+      // by the constructor instrumentation to guarantee that it is the last one in the chain.
+      builder.interceptors().remove(OkHttp3Interceptors.TRACING_INTERCEPTOR);
+    }
   }
 }

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

@@ -7,11 +7,97 @@ package io.opentelemetry.javaagent.instrumentation.okhttp.v3_0
 
 import io.opentelemetry.instrumentation.okhttp.v3_0.AbstractOkHttp3Test
 import io.opentelemetry.instrumentation.test.AgentTestTrait
+import okhttp3.Interceptor
 import okhttp3.OkHttpClient
+import okhttp3.Response
+import org.jetbrains.annotations.NotNull
 
 class OkHttp3Test extends AbstractOkHttp3Test implements AgentTestTrait {
   @Override
   OkHttpClient.Builder configureClient(OkHttpClient.Builder clientBuilder) {
     return clientBuilder
   }
+
+  /**
+   * Makes sure that the builder interceptors are the same before and after invoking
+   * {@link OkHttpClient.Builder#build}. This guarantees that there will be no duplicate tracing
+   * interceptors and that the tracing interceptor is not in the middle of the chain.
+   */
+  def "builder only contains the original interceptors after build"() {
+    setup:
+    def customInterceptor = new TestInterceptor()
+    def clientBuilder = new OkHttpClient.Builder().addInterceptor(customInterceptor)
+    def originalInterceptors = new ArrayList<>(clientBuilder.interceptors())
+
+    when:
+    clientBuilder.build()
+
+    then:
+    clientBuilder.interceptors() == originalInterceptors
+  }
+
+  /**
+   * Makes sure that the tracing interceptor is not present in the builder returned by
+   * {@link OkHttpClient#newBuilder()}. This reduces the chance of the corner case where when
+   * building a client, the builder already contains the tracing interceptor.
+   */
+  def "builder created from client does not contain tracing interceptor"() {
+    setup:
+    def customInterceptor = new TestInterceptor()
+    def clientBuilder = new OkHttpClient.Builder().addInterceptor(customInterceptor)
+    def originalInterceptors = new ArrayList<>(clientBuilder.interceptors())
+
+    when:
+    def builderFromClient = clientBuilder.build().newBuilder()
+
+    then:
+    builderFromClient.interceptors() == originalInterceptors
+  }
+
+  /**
+   * Tests the corner case where the tracing interceptor has been manually added to the builder by
+   * accessing it from the interceptors list of an existing client. In this case, a client created
+   * from that builder should only have one tracing interceptor and it must be in the end, but the
+   * builder state should still be the same before and after building.
+   */
+  def "tracing interceptor in the end for client even if it is in the middle for builder"() {
+    setup:
+    def customInterceptor = new TestInterceptor()
+    def clientBuilder = new OkHttpClient.Builder().addInterceptor(customInterceptor)
+
+    when:
+    def client = clientBuilder.build()
+
+    then:
+    clientBuilder.interceptors().size() == 1
+    clientBuilder.interceptors().get(0).is(customInterceptor)
+    client.interceptors().size() == 2
+    client.interceptors().get(0).is(customInterceptor)
+    client.interceptors().get(1).is(OkHttp3Interceptors.TRACING_INTERCEPTOR)
+
+    when:
+    def otherInterceptor = new TestInterceptor()
+    def newClientBuilder = new OkHttpClient.Builder()
+      .addInterceptor(client.interceptors().get(0))
+      .addInterceptor(client.interceptors().get(1))
+      .addInterceptor(client.interceptors().get(1))
+      .addInterceptor(otherInterceptor)
+    def originalNewInterceptors = newClientBuilder.interceptors()
+    def newClient = newClientBuilder.build()
+
+    then:
+    newClientBuilder.interceptors() == originalNewInterceptors
+    newClient.interceptors().size() == 3
+    newClient.interceptors().get(0).is(customInterceptor)
+    newClient.interceptors().get(1).is(otherInterceptor)
+    newClient.interceptors().get(2).is(OkHttp3Interceptors.TRACING_INTERCEPTOR)
+  }
+
+  private static class TestInterceptor implements Interceptor {
+
+    @Override
+    Response intercept(@NotNull Chain chain) throws IOException {
+      return chain.proceed(chain.request())
+    }
+  }
 }