From 3f6727d6b734e469bf86ded0475b7dfdd4b0474a Mon Sep 17 00:00:00 2001
From: Anuraag Agrawal <aanuraag@amazon.co.jp>
Date: Thu, 19 Aug 2021 18:55:58 +0900
Subject: [PATCH] Don't clobber user decorators in Armeria client
 instrumentation (#3873)

* Don't clobber user decorators in Armeria client instrumentation

* Spot
---
 ...rmeriaWebClientBuilderInstrumentation.java | 28 -------------------
 .../v1_3/AbstractArmeriaHttpClientTest.java   | 18 ++++++++++++
 2 files changed, 18 insertions(+), 28 deletions(-)

diff --git a/instrumentation/armeria-1.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/armeria/v1_3/ArmeriaWebClientBuilderInstrumentation.java b/instrumentation/armeria-1.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/armeria/v1_3/ArmeriaWebClientBuilderInstrumentation.java
index 522a26d81353..f72d14dd8341 100644
--- a/instrumentation/armeria-1.3/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/armeria/v1_3/ArmeriaWebClientBuilderInstrumentation.java
+++ b/instrumentation/armeria-1.3/javaagent/src/main/java/io/opentelemetry/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/instrumentation/armeria-1.3/testing/src/main/java/io/opentelemetry/instrumentation/armeria/v1_3/AbstractArmeriaHttpClientTest.java b/instrumentation/armeria-1.3/testing/src/main/java/io/opentelemetry/instrumentation/armeria/v1_3/AbstractArmeriaHttpClientTest.java
index 0f599757511a..589d9a32c076 100644
--- a/instrumentation/armeria-1.3/testing/src/main/java/io/opentelemetry/instrumentation/armeria/v1_3/AbstractArmeriaHttpClientTest.java
+++ b/instrumentation/armeria-1.3/testing/src/main/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();
+  }
 }