Don't clobber user decorators in Armeria client instrumentation (#3873)

* Don't clobber user decorators in Armeria client instrumentation * Spot
open-telemetry · Aug 19, 2021 · 3f6727d · 3f6727d
1 parent 73b1560
commit 3f6727d
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 28 deletions.
diff --git a/...emetry/javaagent/instrumentation/armeria/v1_3/ArmeriaWebClientBuilderInstrumentation.java b/...emetry/javaagent/instrumentation/armeria/v1_3/ArmeriaWebClientBuilderInstrumentation.java
@@ -8,12 +8,10 @@
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.named;
-import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
 import com.linecorp.armeria.client.WebClientBuilder;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
-import java.util.function.Function;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -27,37 +25,11 @@ public ElementMatcher<TypeDescription> typeMatcher() {
 
   @Override
   public void transform(TypeTransformer transformer) {
-    transformer.applyAdviceToMethod(
-        isMethod().and(isPublic()).and(named("decorator").and(takesArgument(0, Function.class))),
-        ArmeriaWebClientBuilderInstrumentation.class.getName() + "$SuppressDecoratorAdvice");
     transformer.applyAdviceToMethod(
         isMethod().and(isPublic()).and(named("build")),
         ArmeriaWebClientBuilderInstrumentation.class.getName() + "$BuildAdvice");
   }
 
-  // Intercept calls from app to register decorator and suppress them to avoid registering
-  // multiple decorators, one from user app and one from our auto instrumentation. Otherwise, we
-  // will end up with double telemetry.
-  // https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/903
-  @SuppressWarnings("unused")
-  public static class SuppressDecoratorAdvice {
-
-    @Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class)
-    public static boolean suppressDecorator(@Advice.Argument(0) Function<?, ?> decorator) {
-      return decorator != ArmeriaSingletons.CLIENT_DECORATOR;
-    }
-
-    @Advice.OnMethodExit
-    public static void handleSuppression(
-        @Advice.This WebClientBuilder builder,
-        @Advice.Enter boolean suppressed,
-        @Advice.Return(readOnly = false) WebClientBuilder returned) {
-      if (suppressed) {
-        returned = builder;
-      }
-    }
-  }
-
   @SuppressWarnings("unused")
   public static class BuildAdvice {
 

diff --git a/...ain/java/io/opentelemetry/instrumentation/armeria/v1_3/AbstractArmeriaHttpClientTest.java b/...ain/java/io/opentelemetry/instrumentation/armeria/v1_3/AbstractArmeriaHttpClientTest.java
@@ -5,6 +5,8 @@
 
 package io.opentelemetry.instrumentation.armeria.v1_3;
 
+import static org.assertj.core.api.Assertions.assertThat;
+
 import com.linecorp.armeria.client.ClientFactory;
 import com.linecorp.armeria.client.WebClient;
 import com.linecorp.armeria.client.WebClientBuilder;
@@ -21,19 +23,29 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.CompletionException;
+import java.util.concurrent.atomic.AtomicBoolean;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
 
 public abstract class AbstractArmeriaHttpClientTest extends AbstractHttpClientTest<HttpRequest> {
 
   protected abstract WebClientBuilder configureClient(WebClientBuilder clientBuilder);
 
+  private AtomicBoolean decoratorCalled;
+
   private WebClient client;
 
   @BeforeEach
   void setupClient() {
+    decoratorCalled = new AtomicBoolean();
     client =
         configureClient(
                 WebClient.builder()
+                    .decorator(
+                        (delegate, ctx, req) -> {
+                          decoratorCalled.set(true);
+                          return delegate.execute(ctx, req);
+                        })
                     .factory(ClientFactory.builder().connectTimeout(connectTimeout()).build()))
             .build();
   }
@@ -87,4 +99,10 @@ protected void configure(HttpClientTestOptions options) {
     extra.addAll(HttpClientTestOptions.DEFAULT_HTTP_ATTRIBUTES);
     options.setHttpAttributes(unused -> extra);
   }
+
+  @Test
+  void userDecoratorsNotClobbered() {
+    client.get(resolveAddress("/success").toString()).aggregate().join();
+    assertThat(decoratorCalled).isTrue();
+  }
 }