Use SmartInitializingSingleton instead of BeanPostProcessor to set up @LoadBalanced WebClient.Builder. (#1319)

Author: OlgaMaciaszek

spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerAutoConfiguration.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2023 the original author or authors.
+ * Copyright 2012-2024 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -23,6 +23,7 @@
 import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.beans.factory.SmartInitializingSingleton;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.AnyNestedCondition;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -32,12 +33,12 @@
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.cloud.client.ServiceInstance;
 import org.springframework.cloud.client.loadbalancer.reactive.ReactiveLoadBalancer;
-import org.springframework.context.ApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Conditional;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.http.client.ClientHttpRequestInterceptor;
 import org.springframework.retry.support.RetryTemplate;
+import org.springframework.web.client.RestClient;
 import org.springframework.web.client.RestTemplate;
 
 /**
@@ -49,7 +50,7 @@
  * @author Gang Li
  * @author Olga Maciaszek-Sharma
  */
-@Configuration(proxyBeanMethods = false)
+@AutoConfiguration
 @Conditional(BlockingRestClassesPresentCondition.class)
 @ConditionalOnBean(LoadBalancerClient.class)
 @EnableConfigurationProperties(LoadBalancerClientsProperties.class)
@@ -59,6 +60,10 @@ public class LoadBalancerAutoConfiguration {
 	@Autowired(required = false)
 	private List<RestTemplate> restTemplates = Collections.emptyList();
 
+	@LoadBalanced
+	@Autowired(required = false)
+	private List<RestClient.Builder> restClientBuilders = Collections.emptyList();
+
 	@Autowired(required = false)
 	private List<LoadBalancerRequestTransformer> transformers = Collections.emptyList();
 
@@ -74,6 +79,19 @@ public SmartInitializingSingleton loadBalancedRestTemplateInitializerDeprecated(
 		});
 	}
 
+	@Bean
+	public SmartInitializingSingleton loadBalancedRestClientBuilderInitializer(
+			final ObjectProvider<List<RestClientBuilderCustomizer>> restClientBuilderCustomizers) {
+		return () -> restClientBuilderCustomizers.ifAvailable(customizers -> {
+			for (RestClient.Builder restClientBuilder : LoadBalancerAutoConfiguration.this.restClientBuilders) {
+				for (RestClientBuilderCustomizer customizer : customizers) {
+					customizer.customize(restClientBuilder);
+				}
+
+			}
+		});
+	}
+
 	@Bean
 	@ConditionalOnMissingBean
 	public LoadBalancerRequestFactory loadBalancerRequestFactory(LoadBalancerClient loadBalancerClient) {
@@ -101,11 +119,10 @@ public RestTemplateCustomizer restTemplateCustomizer(final LoadBalancerIntercept
 		}
 
 		@Bean
-		@ConditionalOnBean(LoadBalancerInterceptor.class)
 		@ConditionalOnMissingBean
-		LoadBalancerRestClientBuilderBeanPostProcessor lbRestClientPostProcessor(
-				final LoadBalancerInterceptor loadBalancerInterceptor, ApplicationContext context) {
-			return new LoadBalancerRestClientBuilderBeanPostProcessor(loadBalancerInterceptor, context);
+		public RestClientBuilderCustomizer restClientBuilderCustomizer(
+				final LoadBalancerInterceptor loadBalancerInterceptor) {
+			return restClientBuilder -> restClientBuilder.requestInterceptor(loadBalancerInterceptor);
 		}
 
 	}
@@ -174,11 +191,10 @@ public RestTemplateCustomizer restTemplateCustomizer(
 		}
 
 		@Bean
-		@ConditionalOnBean(RetryLoadBalancerInterceptor.class)
 		@ConditionalOnMissingBean
-		LoadBalancerRestClientBuilderBeanPostProcessor lbRestClientPostProcessor(
-				final RetryLoadBalancerInterceptor loadBalancerInterceptor, ApplicationContext context) {
-			return new LoadBalancerRestClientBuilderBeanPostProcessor(loadBalancerInterceptor, context);
+		public RestClientBuilderCustomizer restClientBuilderCustomizer(
+				final RetryLoadBalancerInterceptor loadBalancerInterceptor) {
+			return restClientBuilder -> restClientBuilder.requestInterceptor(loadBalancerInterceptor);
 		}
 
 	}

spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerRestClientBuilderBeanPostProcessor.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2023 the original author or authors.
+ * Copyright 2012-2024 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -28,7 +28,9 @@
  *
  * @author Olga Maciaszek-Sharma
  * @since 4.1.0
+ * @deprecated to be removed in the next release.
  */
+@Deprecated
 public class LoadBalancerRestClientBuilderBeanPostProcessor implements BeanPostProcessor {
 
 	private final ClientHttpRequestInterceptor loadBalancerInterceptor;

spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/RestClientBuilderCustomizer.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2012-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.client.loadbalancer;
+
+import org.springframework.web.client.RestClient;
+
+/**
+ * A customizer interface for {@link RestClient.Builder}. Used to set
+ * {@link LoadBalancerInterceptor} on the builder at the end of the singleton
+ * pre-instantiation phase.
+ *
+ * @author Olga Maciaszek-Sharma
+ * @since 4.1.1
+ */
+public interface RestClientBuilderCustomizer {
+
+	void customize(RestClient.Builder restClientBuilder);
+
+}

spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/AbstractLoadBalancerAutoConfigurationTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2020 the original author or authors.
+ * Copyright 2012-2024 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -117,6 +117,43 @@ void multipleRestClientBuilders() {
 		});
 	}
 
+	@Test
+	void restTemplatesAndRestClientsFromUsersAutoConfiguration() {
+		applicationContextRunner
+				.withConfiguration(AutoConfigurations.of(TwoRestTemplatesAndTwoRestClientBuilders.class))
+				.run(context -> {
+					final Map<String, RestClient.Builder> restClientBuilders = context
+							.getBeansOfType(RestClient.Builder.class);
+					final Map<String, RestTemplate> restTemplates = context.getBeansOfType(RestTemplate.class);
+
+					assertThat(restClientBuilders).isNotNull();
+					assertThat(restClientBuilders.values()).hasSize(2);
+
+					TwoRestTemplatesAndTwoRestClientBuilders.Two two = context
+							.getBean(TwoRestTemplatesAndTwoRestClientBuilders.Two.class);
+
+					assertThat(two.loadBalancedRestClientBuilder).isNotNull();
+					assertLoadBalanced(two.loadBalancedRestClientBuilder);
+
+					assertThat(two.nonLoadBalancedRestClientBuilder).isNotNull();
+					two.nonLoadBalancedRestClientBuilder
+							.requestInterceptors(interceptors -> assertThat(interceptors).isEmpty());
+
+					assertThat(restTemplates).isNotNull();
+					Collection<RestTemplate> templates = restTemplates.values();
+					assertThat(templates).hasSize(2);
+
+					TwoRestTemplatesAndTwoRestClientBuilders.Two twoRestTemplate = context
+							.getBean(TwoRestTemplatesAndTwoRestClientBuilders.Two.class);
+
+					assertThat(twoRestTemplate.loadBalanced).isNotNull();
+					assertLoadBalanced(twoRestTemplate.loadBalanced);
+
+					assertThat(twoRestTemplate.nonLoadBalanced).isNotNull();
+					assertThat(twoRestTemplate.nonLoadBalanced.getInterceptors()).isEmpty();
+				});
+	}
+
 	protected abstract void assertLoadBalanced(RestClient.Builder restClientBuilder);
 
 	protected abstract void assertLoadBalanced(RestTemplate restTemplate);

spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/support/LoadBalancerEnvironmentPropertyUtils.java

@@ -24,7 +24,7 @@
 public final class LoadBalancerEnvironmentPropertyUtils {
 
 	private LoadBalancerEnvironmentPropertyUtils() {
-		throw new IllegalStateException("Should not instantiate a utility class");
+		throw new UnsupportedOperationException("Cannot instantiate a utility class");
 	}
 
 	public static boolean trueForClientOrDefault(Environment environment, String propertySuffix) {

src/checkstyle/checkstyle-suppressions.xml

@@ -18,5 +18,7 @@
 	<suppress files=".*Tests.*" checks="JavadocVariable"/>
 	<suppress files=".*Tests.*" checks="JavadocMethod"/>
 	<suppress files=".*Tests.*" checks="HideUtilityClassConstructor"/>
+	<suppress files=".*AutoConfiguration.*" checks="HideUtilityClassConstructor"/>
+	<suppress files=".*AutoConfiguration.*" checks="FinalClass"/>
 	<suppress files=".*ReactiveDiscoveryClient.*" checks="JavadocVariable"/>
 </suppressions>