Add fix for WebClient

Signed-off-by: Lenin Jaganathan <[email protected]>

Author: lenin-jaganathan

instrumentation/spring/spring-boot-autoconfigure/src/main/java/io/opentelemetry/instrumentation/spring/autoconfigure/internal/instrumentation/webflux/SpringWebfluxInstrumentationAutoConfiguration.java

@@ -8,10 +8,14 @@
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.instrumentation.api.incubator.config.internal.InstrumentationConfig;
 import io.opentelemetry.instrumentation.spring.autoconfigure.internal.ConditionalOnEnabledInstrumentation;
+import io.opentelemetry.instrumentation.spring.webflux.v5_3.SpringWebfluxClientTelemetry;
 import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
+import org.springframework.boot.web.reactive.function.client.WebClientCustomizer;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.core.Ordered;
+import org.springframework.core.annotation.Order;
 import org.springframework.web.reactive.function.client.WebClient;
 import org.springframework.web.server.WebFilter;
 
@@ -24,7 +28,7 @@
  * at any time.
  */
 @ConditionalOnEnabledInstrumentation(module = "spring-webflux")
-@ConditionalOnClass(WebClient.class)
+@ConditionalOnClass({WebClient.class, WebClientCustomizer.class})
 @Configuration
 public class SpringWebfluxInstrumentationAutoConfiguration {
 
@@ -38,6 +42,15 @@ static WebClientBeanPostProcessor otelWebClientBeanPostProcessor(
     return new WebClientBeanPostProcessor(openTelemetryProvider, configProvider);
   }
 
+  @Bean
+  @Order(Ordered.HIGHEST_PRECEDENCE + 10)
+  WebClientCustomizer otelWebClientCustomizer(
+      OpenTelemetry openTelemetry, InstrumentationConfig config) {
+    SpringWebfluxClientTelemetry webfluxClientTelemetry =
+        WebClientBeanPostProcessor.getWebfluxClientTelemetry(openTelemetry, config);
+    return builder -> builder.filters(webfluxClientTelemetry::addFilter);
+  }
+
   @Bean
   WebFilter telemetryFilter(OpenTelemetry openTelemetry, InstrumentationConfig config) {
     return WebClientBeanPostProcessor.getWebfluxServerTelemetry(openTelemetry, config)

instrumentation/spring/spring-boot-autoconfigure/src/main/java/io/opentelemetry/instrumentation/spring/autoconfigure/internal/instrumentation/webflux/WebClientBeanPostProcessor.java

@@ -11,8 +11,12 @@
 import io.opentelemetry.instrumentation.spring.webflux.v5_3.SpringWebfluxClientTelemetry;
 import io.opentelemetry.instrumentation.spring.webflux.v5_3.SpringWebfluxServerTelemetry;
 import io.opentelemetry.instrumentation.spring.webflux.v5_3.internal.SpringWebfluxBuilderUtil;
+import io.opentelemetry.instrumentation.spring.webflux.v5_3.internal.WebClientTracingFilter;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
 import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.beans.factory.config.BeanPostProcessor;
+import org.springframework.web.reactive.function.client.ExchangeFilterFunction;
 import org.springframework.web.reactive.function.client.WebClient;
 
 /**
@@ -53,18 +57,30 @@ static SpringWebfluxServerTelemetry getWebfluxServerTelemetry(
   @Override
   public Object postProcessAfterInitialization(Object bean, String beanName) {
     if (bean instanceof WebClient) {
-      WebClient webClient = (WebClient) bean;
-      return wrapBuilder(webClient.mutate()).build();
-    } else if (bean instanceof WebClient.Builder) {
-      WebClient.Builder webClientBuilder = (WebClient.Builder) bean;
-      return wrapBuilder(webClientBuilder);
+      return addWebClientFilterIfNotPresent(
+          (WebClient) bean, openTelemetryProvider.getObject(), configProvider.getObject());
     }
     return bean;
   }
 
-  private WebClient.Builder wrapBuilder(WebClient.Builder webClientBuilder) {
-    SpringWebfluxClientTelemetry instrumentation =
-        getWebfluxClientTelemetry(openTelemetryProvider.getObject(), configProvider.getObject());
-    return webClientBuilder.filters(instrumentation::addFilter);
+  private static WebClient addWebClientFilterIfNotPresent(
+      WebClient webClient, OpenTelemetry openTelemetry, InstrumentationConfig config) {
+    AtomicBoolean filterAdded = new AtomicBoolean(false);
+    WebClient.Builder builder =
+        webClient
+            .mutate()
+            .filters(
+                filters -> {
+                  if (isFilterNotPresent(filters)) {
+                    getWebfluxClientTelemetry(openTelemetry, config).addFilter(filters);
+                    filterAdded.set(true);
+                  }
+                });
+
+    return filterAdded.get() ? builder.build() : webClient;
+  }
+
+  private static boolean isFilterNotPresent(List<ExchangeFilterFunction> filters) {
+    return filters.stream().noneMatch(WebClientTracingFilter.class::isInstance);
   }
 }

instrumentation/spring/spring-boot-autoconfigure/src/main/javaSpring3/io/opentelemetry/instrumentation/spring/autoconfigure/internal/instrumentation/web/RestClientBeanPostProcessor.java

@@ -10,6 +10,7 @@
 import io.opentelemetry.instrumentation.spring.autoconfigure.internal.properties.InstrumentationConfigUtil;
 import io.opentelemetry.instrumentation.spring.web.v3_1.SpringWebTelemetry;
 import io.opentelemetry.instrumentation.spring.web.v3_1.internal.WebTelemetryUtil;
+import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;
 import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.beans.factory.config.BeanPostProcessor;
@@ -57,7 +58,7 @@ private static RestClient addRestClientInterceptorIfNotPresent(
   }
 
   private static boolean isInterceptorNotPresent(
-      java.util.List<ClientHttpRequestInterceptor> interceptors,
+      List<ClientHttpRequestInterceptor> interceptors,
       ClientHttpRequestInterceptor instrumentationInterceptor) {
     return interceptors.stream()
         .noneMatch(interceptor -> interceptor.getClass() == instrumentationInterceptor.getClass());

instrumentation/spring/spring-boot-autoconfigure/src/main/javaSpring3/io/opentelemetry/instrumentation/spring/autoconfigure/internal/instrumentation/web/RestClientInstrumentationAutoConfiguration.java

@@ -15,6 +15,8 @@
 import org.springframework.boot.web.client.RestClientCustomizer;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.core.Ordered;
+import org.springframework.core.annotation.Order;
 import org.springframework.web.client.RestClient;
 
 /**
@@ -39,6 +41,7 @@ static RestClientBeanPostProcessor otelRestClientBeanPostProcessor(
   }
 
   @Bean
+  @Order(Ordered.HIGHEST_PRECEDENCE + 10)
   RestClientCustomizer otelRestClientCustomizer(
       ObjectProvider<OpenTelemetry> openTelemetryProvider,
       ObjectProvider<InstrumentationConfig> configProvider) {

instrumentation/spring/spring-boot-autoconfigure/src/main/javaSpring4/io/opentelemetry/instrumentation/spring/autoconfigure/internal/instrumentation/web/RestClientInstrumentationSpringBoot4AutoConfiguration.java

@@ -15,6 +15,8 @@
 import org.springframework.boot.restclient.autoconfigure.RestClientAutoConfiguration;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.core.Ordered;
+import org.springframework.core.annotation.Order;
 import org.springframework.web.client.RestClient;
 
 /**
@@ -39,6 +41,7 @@ static RestClientBeanPostProcessorSpring4 otelRestClientBeanPostProcessor(
   }
 
   @Bean
+  @Order(Ordered.HIGHEST_PRECEDENCE + 10)
   RestClientCustomizer otelRestClientCustomizer(
       ObjectProvider<OpenTelemetry> openTelemetryProvider,
       ObjectProvider<InstrumentationConfig> configProvider) {

instrumentation/spring/spring-boot-autoconfigure/src/test/java/io/opentelemetry/instrumentation/spring/autoconfigure/internal/instrumentation/webflux/SpringWebfluxInstrumentationAutoConfigurationTest.java

@@ -12,9 +12,12 @@
 import io.opentelemetry.instrumentation.spring.autoconfigure.internal.properties.ConfigPropertiesBridge;
 import io.opentelemetry.sdk.autoconfigure.spi.internal.DefaultConfigProperties;
 import java.util.Collections;
+import java.util.concurrent.atomic.AtomicLong;
 import org.junit.jupiter.api.Test;
 import org.springframework.boot.autoconfigure.AutoConfigurations;
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
+import org.springframework.boot.web.reactive.function.client.WebClientCustomizer;
+import org.springframework.web.reactive.function.client.WebClient;
 
 class SpringWebfluxInstrumentationAutoConfigurationTest {
 
@@ -35,10 +38,9 @@ void instrumentationEnabled() {
         .withPropertyValues("otel.instrumentation.spring-webflux.enabled=true")
         .run(
             context ->
-                assertThat(
-                        context.getBean(
-                            "otelWebClientBeanPostProcessor", WebClientBeanPostProcessor.class))
-                    .isNotNull());
+                assertThat(context)
+                    .hasBean("otelWebClientBeanPostProcessor")
+                    .hasBean("otelWebClientCustomizer"));
   }
 
   @Test
@@ -47,16 +49,45 @@ void instrumentationDisabled() {
         .withPropertyValues("otel.instrumentation.spring-webflux.enabled=false")
         .run(
             context ->
-                assertThat(context.containsBean("otelWebClientBeanPostProcessor")).isFalse());
+                assertThat(context)
+                    .doesNotHaveBean("otelWebClientBeanPostProcessor")
+                    .doesNotHaveBean("otelWebClientCustomizer"));
   }
 
   @Test
   void defaultConfiguration() {
     contextRunner.run(
-        context ->
-            assertThat(
-                    context.getBean(
-                        "otelWebClientBeanPostProcessor", WebClientBeanPostProcessor.class))
-                .isNotNull());
+        context -> {
+          assertThat(context)
+              .hasBean("otelWebClientBeanPostProcessor")
+              .hasBean("otelWebClientCustomizer");
+        });
+  }
+
+  @Test
+  void shouldAddTracingFilterWhenCustomizerApplied() {
+    contextRunner.run(
+        context -> {
+          WebClientCustomizer customizer =
+              context.getBean("otelWebClientCustomizer", WebClientCustomizer.class);
+          WebClient.Builder builder = WebClient.builder();
+          customizer.customize(builder);
+
+          AtomicLong count = new AtomicLong(0);
+          builder
+              .build()
+              .mutate()
+              .filters(
+                  filters ->
+                      count.set(
+                          filters.stream()
+                              .filter(
+                                  f ->
+                                      f.getClass()
+                                          .getName()
+                                          .startsWith("io.opentelemetry.instrumentation"))
+                              .count()));
+          assertThat(count.get()).isEqualTo(1);
+        });
   }
 }

instrumentation/spring/spring-boot-autoconfigure/src/test/java/io/opentelemetry/instrumentation/spring/autoconfigure/internal/instrumentation/webflux/WebClientBeanPostProcessorTest.java

@@ -12,6 +12,8 @@
 import io.opentelemetry.instrumentation.spring.autoconfigure.internal.properties.ConfigPropertiesBridge;
 import io.opentelemetry.sdk.autoconfigure.spi.internal.DefaultConfigProperties;
 import java.util.Collections;
+import java.util.concurrent.atomic.AtomicLong;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.DisplayName;
 import org.junit.jupiter.api.Test;
 import org.springframework.beans.factory.config.BeanPostProcessor;
@@ -29,92 +31,62 @@ class WebClientBeanPostProcessorTest {
         new ConfigPropertiesBridge(DefaultConfigProperties.createFromMap(Collections.emptyMap())));
   }
 
-  @Test
-  @DisplayName(
-      "when processed bean is NOT of type WebClient or WebClientBuilder should return Object")
-  void returnsObject() {
-    BeanPostProcessor underTest =
-        new WebClientBeanPostProcessor(
-            beanFactory.getBeanProvider(OpenTelemetry.class),
-            beanFactory.getBeanProvider(InstrumentationConfig.class));
-
-    assertThat(underTest.postProcessAfterInitialization(new Object(), "testObject"))
-        .isExactlyInstanceOf(Object.class);
-  }
+  private BeanPostProcessor underTest;
 
-  @Test
-  @DisplayName("when processed bean is of type WebClient should return WebClient")
-  void returnsWebClient() {
-    BeanPostProcessor underTest =
+  @BeforeEach
+  void setUp() {
+    underTest =
         new WebClientBeanPostProcessor(
             beanFactory.getBeanProvider(OpenTelemetry.class),
             beanFactory.getBeanProvider(InstrumentationConfig.class));
-
-    assertThat(underTest.postProcessAfterInitialization(WebClient.create(), "testWebClient"))
-        .isInstanceOf(WebClient.class);
   }
 
   @Test
-  @DisplayName("when processed bean is of type WebClientBuilder should return WebClientBuilder")
-  void returnsWebClientBuilder() {
-    BeanPostProcessor underTest =
-        new WebClientBeanPostProcessor(
-            beanFactory.getBeanProvider(OpenTelemetry.class),
-            beanFactory.getBeanProvider(InstrumentationConfig.class));
+  @DisplayName("when processed bean is NOT of type WebClient should return same Object")
+  void returnsObject() {
+    Object original = new Object();
 
-    assertThat(
-            underTest.postProcessAfterInitialization(WebClient.builder(), "testWebClientBuilder"))
-        .isInstanceOf(WebClient.Builder.class);
+    assertThat(underTest.postProcessAfterInitialization(original, "testObject")).isSameAs(original);
   }
 
   @Test
-  @DisplayName("when processed bean is of type WebClient should add exchange filter to WebClient")
-  void addsExchangeFilterWebClient() {
-    BeanPostProcessor underTest =
-        new WebClientBeanPostProcessor(
-            beanFactory.getBeanProvider(OpenTelemetry.class),
-            beanFactory.getBeanProvider(InstrumentationConfig.class));
-
+  @DisplayName("when processed bean is of type WebClient should return WebClient with filter")
+  void returnsWebClientWithFilter() {
     WebClient webClient = WebClient.create();
     Object processedWebClient =
         underTest.postProcessAfterInitialization(webClient, "testWebClient");
 
-    assertThat(processedWebClient).isInstanceOf(WebClient.class);
-    ((WebClient) processedWebClient)
-        .mutate()
-        .filters(
-            functions ->
-                assertThat(
-                        functions.stream()
-                            .filter(WebClientBeanPostProcessorTest::isOtelExchangeFilter)
-                            .count())
-                    .isEqualTo(1));
+    assertThat(processedWebClient).isInstanceOf(WebClient.class).isNotSameAs(webClient);
+    assertFilterCount((WebClient) processedWebClient, 1);
   }
 
   @Test
-  @DisplayName(
-      "when processed bean is of type WebClientBuilder should add ONE exchange filter to WebClientBuilder")
-  void addsExchangeFilterWebClientBuilder() {
-    BeanPostProcessor underTest =
-        new WebClientBeanPostProcessor(
-            beanFactory.getBeanProvider(OpenTelemetry.class),
-            beanFactory.getBeanProvider(InstrumentationConfig.class));
+  @DisplayName("when WebClient already has filter should return same instance")
+  void doesNotAddDuplicateFilter() {
+    WebClient webClient = WebClient.create();
+    WebClient firstProcessed =
+        (WebClient) underTest.postProcessAfterInitialization(webClient, "testWebClient");
+    WebClient secondProcessed =
+        (WebClient) underTest.postProcessAfterInitialization(firstProcessed, "testWebClient");
 
-    WebClient.Builder webClientBuilder = WebClient.builder();
-    underTest.postProcessAfterInitialization(webClientBuilder, "testWebClientBuilder");
-    underTest.postProcessAfterInitialization(webClientBuilder, "testWebClientBuilder");
-    underTest.postProcessAfterInitialization(webClientBuilder, "testWebClientBuilder");
+    assertThat(secondProcessed).isSameAs(firstProcessed);
+    assertFilterCount(secondProcessed, 1);
+  }
 
-    webClientBuilder.filters(
-        functions ->
-            assertThat(
-                    functions.stream()
+  private static void assertFilterCount(WebClient webClient, long expectedCount) {
+    AtomicLong count = new AtomicLong(0);
+    webClient
+        .mutate()
+        .filters(
+            filters ->
+                count.set(
+                    filters.stream()
                         .filter(WebClientBeanPostProcessorTest::isOtelExchangeFilter)
-                        .count())
-                .isEqualTo(1));
+                        .count()));
+    assertThat(count.get()).isEqualTo(expectedCount);
   }
 
-  private static boolean isOtelExchangeFilter(ExchangeFilterFunction wctf) {
-    return wctf.getClass().getName().startsWith("io.opentelemetry.instrumentation");
+  private static boolean isOtelExchangeFilter(ExchangeFilterFunction filter) {
+    return filter.getClass().getName().startsWith("io.opentelemetry.instrumentation");
   }
 }

instrumentation/spring/spring-boot-autoconfigure/src/testSpring3/java/io/opentelemetry/instrumentation/spring/autoconfigure/internal/instrumentation/web/RestClientInstrumentationAutoConfigurationTest.java

@@ -5,8 +5,14 @@
 
 package io.opentelemetry.instrumentation.spring.autoconfigure.internal.instrumentation.web;
 
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.instrumentation.api.incubator.config.internal.InstrumentationConfig;
 import io.opentelemetry.instrumentation.spring.autoconfigure.internal.AbstractRestClientInstrumentationAutoConfigurationTest;
+import org.junit.jupiter.api.Test;
 import org.springframework.boot.autoconfigure.AutoConfigurations;
+import org.springframework.web.client.RestClient;
 
 class RestClientInstrumentationAutoConfigurationTest
     extends AbstractRestClientInstrumentationAutoConfigurationTest {