Support @LoadBalanced RestTemplateBuilder (#1403)

docs/modules/ROOT/pages/spring-cloud-commons/common-abstractions.adoc

@@ -218,6 +218,89 @@ IMPORTANT: Notice the use of the `@Primary` annotation on the plain `RestTemplat
 
 TIP: If you see errors such as `java.lang.IllegalArgumentException: Can not set org.springframework.web.client.RestTemplate field com.my.app.Foo.restTemplate to com.sun.proxy.$Proxy89`, try injecting `RestOperations` or setting `spring.aop.proxyTargetClass=true`.
 
+[[rest-template-builder-loadbalancer-client]]
+== Using `@LoadBalanced RestTemplateBuilder` to create a LoadBalancer Client
+
+You can also configure a `RestTemplate` to use a Load-Balancer client by annotating a
+`RestTemplateBuilder` bean with `@LoadBalanced`:
+
+[source,java,indent=0]
+----
+import org.springframework.boot.web.client.RestTemplateBuilder;@Configuration
+public class MyConfiguration {
+
+    @Bean
+	@LoadBalanced
+	RestTemplateBuilder loadBalancedRestTemplateBuilder() {
+		return new RestTemplateBuilder();
+	}
+}
+
+public class MyClass {
+
+    private final RestTemplate restTemplate;
+
+	MyClass(@LoadBalanced RestTemplateBuilder restTemplateBuilder) {
+		this.restTemplate = restTemplateBuilder.build();
+	}
+
+    public String getStores() {
+        return restTemplate.getForObject("http://stores/stores", String.class);
+    }
+}
+----
+
+The URI needs to use a virtual host name (that is, a service name, not a host name).
+The `BlockingLoadBalancerClient` is used to create a full physical address.
+
+IMPORTANT: To use it, add xref:spring-cloud-commons/loadbalancer.adoc#spring-cloud-loadbalancer-starter[Spring Cloud LoadBalancer starter] to your project.
+
+[[multiple-resttemplate-builder-beans]]
+=== Multiple `RestTemplateBuilder` beans
+
+If you want a `RestTemplateBuilder` that is not load-balanced, create a `RestTemplateBuilder` bean and inject it.
+To access the load-balanced `RestTemplateBuilder`, use the `@LoadBalanced` qualifier when you create your `@Bean`, as the following example shows:
+
+[source,java,indent=0]
+----
+@Configuration
+public class MyConfiguration {
+
+    @LoadBalanced
+    @Bean
+    RestTemplateBuilder loadBalancedRestTemplateBuilder() {
+		return new RestTemplateBuilder();
+	}
+
+    @Primary
+    @Bean
+    RestTemplateBuilder restTemplateBuilder() {
+		return new RestTemplateBuilder();
+	}
+}
+
+public class MyClass {
+
+    @Autowired
+    private RestTemplateBuilder restTemplateBuilder;
+
+    @Autowired
+    @LoadBalanced
+    private RestTemplateBuilder loadBalanced;
+
+    public String doOtherStuff() {
+        return loadBalanced.getForObject("http://stores/stores", String.class);
+    }
+
+    public String doStuff() {
+		return restTemplateBuilder.build().getForObject("http://example.com", String.class);
+    }
+}
+----
+
+IMPORTANT: Notice the use of the `@Primary` annotation on the plain `RestTemplateBuilder` declaration in the preceding example to disambiguate the unqualified `@Autowired` injection.
+
+
 [[rest-client-loadbalancer-client]]
 == Spring RestClient as a LoadBalancer Client
 
@@ -256,7 +339,7 @@ IMPORTANT: To use it, add xref:spring-cloud-commons/loadbalancer.adoc#spring-clo
 === Multiple `RestClient.Builder` Objects
 
 If you want a `RestClient.Builder` that is not load-balanced, create a `RestClient.Builder` bean and inject it.
-To access the load-balanced `RestClient`, use the `@LoadBalanced` qualifier when you create your `@Bean`, as the following example shows:
+To access the load-balanced `RestClient.Builder`, use the `@LoadBalanced` qualifier when you create your `@Bean`, as the following example shows:
 
 [source,java,indent=0]
 ----
@@ -296,7 +379,7 @@ public class MyClass {
 }
 ----
 
-IMPORTANT: Notice the use of the `@Primary` annotation on the plain `RestTemplate` declaration in the preceding example to disambiguate the unqualified `@Autowired` injection.
+IMPORTANT: Notice the use of the `@Primary` annotation on the plain `RestClient.Builder` declaration in the preceding example to disambiguate the unqualified `@Autowired` injection.
 
 [[webclinet-loadbalancer-client]]
 == Spring WebClient as a LoadBalancer Client

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

@@ -0,0 +1,66 @@
+/*
+ * 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.beans.BeansException;
+import org.springframework.beans.factory.ObjectProvider;
+import org.springframework.beans.factory.config.BeanPostProcessor;
+import org.springframework.context.ApplicationContext;
+import org.springframework.http.client.ClientHttpRequestInterceptor;
+
+/**
+ * A {@link BeanPostProcessor} that adds the provided {@link ClientHttpRequestInterceptor}
+ * to bean instances annotated with {@link LoadBalanced}.
+ *
+ * @author Olga Maciaszek-Sharma
+ * @since 4.2.0
+ */
+public abstract class AbstractLoadBalancerBlockingBuilderBeanPostProcessor<T extends ClientHttpRequestInterceptor>
+		implements BeanPostProcessor {
+
+	protected final ObjectProvider<T> loadBalancerInterceptorProvider;
+
+	protected final ApplicationContext context;
+
+	AbstractLoadBalancerBlockingBuilderBeanPostProcessor(ObjectProvider<T> loadBalancerInterceptorProvider,
+			ApplicationContext context) {
+		this.loadBalancerInterceptorProvider = loadBalancerInterceptorProvider;
+		this.context = context;
+	}
+
+	@Override
+	public Object postProcessBeforeInitialization(Object bean, String beanName) throws BeansException {
+		// a separate call to verify supported type before searching for annotation for
+		// performance reasons
+		if (isSupported(bean)) {
+			if (context.findAnnotationOnBean(beanName, LoadBalanced.class) == null) {
+				return bean;
+			}
+			ClientHttpRequestInterceptor interceptor = loadBalancerInterceptorProvider.getIfAvailable();
+			if (interceptor == null) {
+				throw new IllegalStateException(ClientHttpRequestInterceptor.class.getSimpleName() + " not available.");
+			}
+			bean = apply(bean, interceptor);
+		}
+		return bean;
+	}
+
+	protected abstract Object apply(Object bean, ClientHttpRequestInterceptor interceptor);
+
+	protected abstract boolean isSupported(Object bean);
+
+}

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

@@ -101,6 +101,15 @@ public class LoadBalancerAutoConfiguration {
 			return new LoadBalancerRestClientBuilderBeanPostProcessor<>(loadBalancerInterceptorProvider, context);
 		}
 
+		@Bean
+		@ConditionalOnBean(DeferringLoadBalancerInterceptor.class)
+		@ConditionalOnMissingBean(LoadBalancerRestTemplateBuilderBeanPostProcessor.class)
+		static LoadBalancerRestTemplateBuilderBeanPostProcessor<DeferringLoadBalancerInterceptor> lbRestTemplateBuilderPostProcessor(
+				ObjectProvider<DeferringLoadBalancerInterceptor> loadBalancerInterceptorProvider,
+				ApplicationContext context) {
+			return new LoadBalancerRestTemplateBuilderBeanPostProcessor<>(loadBalancerInterceptorProvider, context);
+		}
+
 	}
 
 	@AutoConfiguration

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

@@ -16,26 +16,22 @@
 
 package org.springframework.cloud.client.loadbalancer;
 
-import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.ObjectProvider;
-import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.context.ApplicationContext;
 import org.springframework.http.client.ClientHttpRequestInterceptor;
 import org.springframework.web.client.RestClient;
 
 /**
- * A {@link BeanPostProcessor} that adds the provided {@link ClientHttpRequestInterceptor}
- * to all {@link RestClient.Builder} instances annotated with {@link LoadBalanced}.
+ * {@link RestClient.Builder}-specific
+ * {@link AbstractLoadBalancerBlockingBuilderBeanPostProcessor} implementation. Adds the
+ * provided {@link ClientHttpRequestInterceptor} to all {@link RestClient.Builder}
+ * instances annotated with {@link LoadBalanced}.
  *
  * @author Olga Maciaszek-Sharma
  * @since 4.1.0
  */
 public class LoadBalancerRestClientBuilderBeanPostProcessor<T extends ClientHttpRequestInterceptor>
-		implements BeanPostProcessor {
-
-	private final ObjectProvider<T> loadBalancerInterceptorProvider;
-
-	private final ApplicationContext context;
+		extends AbstractLoadBalancerBlockingBuilderBeanPostProcessor<T> {
 
 	/**
 	 * Creates a {@link LoadBalancerRestClientBuilderBeanPostProcessor} instance using a
@@ -48,8 +44,7 @@ public class LoadBalancerRestClientBuilderBeanPostProcessor<T extends ClientHttp
 	 */
 	@Deprecated(forRemoval = true)
 	public LoadBalancerRestClientBuilderBeanPostProcessor(T loadBalancerInterceptor, ApplicationContext context) {
-		this.loadBalancerInterceptorProvider = new SimpleObjectProvider<>(loadBalancerInterceptor);
-		this.context = context;
+		this(new SimpleObjectProvider<>(loadBalancerInterceptor), context);
 	}
 
 	/**
@@ -61,23 +56,17 @@ public class LoadBalancerRestClientBuilderBeanPostProcessor<T extends ClientHttp
 	 */
 	public LoadBalancerRestClientBuilderBeanPostProcessor(ObjectProvider<T> loadBalancerInterceptorProvider,
 			ApplicationContext context) {
-		this.loadBalancerInterceptorProvider = loadBalancerInterceptorProvider;
-		this.context = context;
+		super(loadBalancerInterceptorProvider, context);
 	}
 
 	@Override
-	public Object postProcessBeforeInitialization(Object bean, String beanName) throws BeansException {
-		if (bean instanceof RestClient.Builder) {
-			if (context.findAnnotationOnBean(beanName, LoadBalanced.class) == null) {
-				return bean;
-			}
-			ClientHttpRequestInterceptor interceptor = loadBalancerInterceptorProvider.getIfAvailable();
-			if (interceptor == null) {
-				throw new IllegalStateException(ClientHttpRequestInterceptor.class.getSimpleName() + " not available.");
-			}
-			((RestClient.Builder) bean).requestInterceptor(interceptor);
-		}
-		return bean;
+	protected Object apply(Object bean, ClientHttpRequestInterceptor interceptor) {
+		return ((RestClient.Builder) bean).requestInterceptor(interceptor);
+	}
+
+	@Override
+	protected boolean isSupported(Object bean) {
+		return bean instanceof RestClient.Builder;
 	}
 
 }

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

@@ -0,0 +1,51 @@
+/*
+ * 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.beans.factory.ObjectProvider;
+import org.springframework.boot.web.client.RestTemplateBuilder;
+import org.springframework.context.ApplicationContext;
+import org.springframework.http.client.ClientHttpRequestInterceptor;
+
+/**
+ * {@link RestTemplateBuilder}-specific
+ * {@link AbstractLoadBalancerBlockingBuilderBeanPostProcessor} implementation. Adds the
+ * provided {@link ClientHttpRequestInterceptor} to all {@link RestTemplateBuilder}
+ * instances annotated with {@link LoadBalanced}.
+ *
+ * @author Olga Maciaszek-Sharma
+ * @since 4.2.0
+ */
+public class LoadBalancerRestTemplateBuilderBeanPostProcessor<T extends ClientHttpRequestInterceptor>
+		extends AbstractLoadBalancerBlockingBuilderBeanPostProcessor<T> {
+
+	public LoadBalancerRestTemplateBuilderBeanPostProcessor(ObjectProvider<T> loadBalancerInterceptorProvider,
+			ApplicationContext context) {
+		super(loadBalancerInterceptorProvider, context);
+	}
+
+	@Override
+	protected boolean isSupported(Object bean) {
+		return bean instanceof RestTemplateBuilder;
+	}
+
+	@Override
+	protected Object apply(Object bean, ClientHttpRequestInterceptor interceptor) {
+		return ((RestTemplateBuilder) bean).interceptors(interceptor);
+	}
+
+}

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

@@ -16,7 +16,6 @@
 
 package org.springframework.cloud.client.loadbalancer;
 
-import java.io.IOException;
 import java.net.URI;
 
 import org.junit.jupiter.api.Test;
@@ -25,7 +24,6 @@ import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.SpringBootConfiguration;
 import org.springframework.boot.test.context.SpringBootTest;
 import org.springframework.cloud.client.ServiceInstance;
-import org.springframework.context.ApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.web.client.RestClient;
 
@@ -41,17 +39,16 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy;
 		properties = "spring.cloud.loadbalancer.retry.enabled=false")
 public class LoadBalancedRestClientIntegrationTests {
 
-	private final RestClient client;
+	private final RestClient.Builder restClientBuilder;
 
-	@Autowired
-	ApplicationContext context;
-
-	public LoadBalancedRestClientIntegrationTests(@Autowired RestClient.Builder clientBuilder) {
-		this.client = clientBuilder.build();
+	public LoadBalancedRestClientIntegrationTests(@Autowired RestClient.Builder restClientBuilder) {
+		this.restClientBuilder = restClientBuilder;
 	}
 
 	@Test
 	void shouldBuildLoadBalancedRestClientInConstructor() {
+		RestClient client = restClientBuilder.build();
+
 		// Interceptors are not visible in RestClient
 		assertThatThrownBy(() -> client.get().uri("http://test-service").retrieve())
 			.hasMessage("LoadBalancerInterceptor invoked.");
@@ -70,13 +67,13 @@ public class LoadBalancedRestClientIntegrationTests {
 		LoadBalancerClient testLoadBalancerClient() {
 			return new LoadBalancerClient() {
 				@Override
-				public <T> T execute(String serviceId, LoadBalancerRequest<T> request) throws IOException {
+				public <T> T execute(String serviceId, LoadBalancerRequest<T> request) {
 					throw new UnsupportedOperationException("LoadBalancerInterceptor invoked.");
 				}
 
 				@Override
-				public <T> T execute(String serviceId, ServiceInstance serviceInstance, LoadBalancerRequest<T> request)
-						throws IOException {
+				public <T> T execute(String serviceId, ServiceInstance serviceInstance,
+						LoadBalancerRequest<T> request) {
 					throw new UnsupportedOperationException("LoadBalancerInterceptor invoked.");
 				}
 

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

@@ -0,0 +1,101 @@
+/*
+ * 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 java.net.URI;
+
+import org.junit.jupiter.api.Test;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.SpringBootConfiguration;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.boot.web.client.RestTemplateBuilder;
+import org.springframework.cloud.client.ServiceInstance;
+import org.springframework.context.annotation.Bean;
+import org.springframework.web.client.RestTemplate;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Integration tests for load-balanced {@link RestTemplateBuilder}.
+ *
+ * @author Olga Maciaszek-Sharma
+ */
+@SpringBootTest(classes = { LoadBalancedRestTemplateBuilderIntegrationTests.TestConfig.class,
+		LoadBalancerAutoConfiguration.class }, properties = "spring.cloud.loadbalancer.retry.enabled=false")
+public class LoadBalancedRestTemplateBuilderIntegrationTests {
+
+	private final RestTemplateBuilder restTemplateBuilder;
+
+	public LoadBalancedRestTemplateBuilderIntegrationTests(@Autowired RestTemplateBuilder restTemplateBuilder) {
+		this.restTemplateBuilder = restTemplateBuilder;
+	}
+
+	@Test
+	void shouldBuildLoadBalancedRestTemplate() {
+		RestTemplate restTemplate = restTemplateBuilder.build();
+
+		assertThat(restTemplate.getInterceptors()).hasSize(1);
+		assertThat(restTemplate.getInterceptors().get(0)).isInstanceOf(DeferringLoadBalancerInterceptor.class);
+		assertThat(((DeferringLoadBalancerInterceptor) restTemplate.getInterceptors().get(0))
+			.getLoadBalancerInterceptorProvider()
+			.getObject()).isInstanceOf(BlockingLoadBalancerInterceptor.class);
+	}
+
+	@SpringBootConfiguration
+	static class TestConfig {
+
+		@LoadBalanced
+		@Bean
+		RestTemplateBuilder restTemplateBuilder() {
+			return new RestTemplateBuilder();
+		}
+
+		@Bean
+		LoadBalancerClient testLoadBalancerClient() {
+			return new LoadBalancerClient() {
+				@Override
+				public <T> T execute(String serviceId, LoadBalancerRequest<T> request) {
+					throw new UnsupportedOperationException("LoadBalancerInterceptor invoked.");
+				}
+
+				@Override
+				public <T> T execute(String serviceId, ServiceInstance serviceInstance,
+						LoadBalancerRequest<T> request) {
+					throw new UnsupportedOperationException("LoadBalancerInterceptor invoked.");
+				}
+
+				@Override
+				public URI reconstructURI(ServiceInstance instance, URI original) {
+					throw new UnsupportedOperationException("LoadBalancerInterceptor invoked.");
+				}
+
+				@Override
+				public ServiceInstance choose(String serviceId) {
+					throw new UnsupportedOperationException("LoadBalancerInterceptor invoked.");
+				}
+
+				@Override
+				public <T> ServiceInstance choose(String serviceId, Request<T> request) {
+					throw new UnsupportedOperationException("LoadBalancerInterceptor invoked.");
+				}
+			};
+		}
+
+	}
+
+}