From eddaaac5cb5883a8866d5f6721efb97bb20b6ce7 Mon Sep 17 00:00:00 2001 From: Toshiaki Maki Date: Fri, 12 Mar 2021 21:30:32 +0900 Subject: [PATCH] Add LoadBalancerClientRequestTransformer to 2.2.x (#925) --- .../main/asciidoc/spring-cloud-commons.adoc | 47 +++++++ .../reactive/ExchangeFilterFunctionUtils.java | 15 ++- .../LoadBalancerClientRequestTransformer.java | 39 ++++++ ...orLoadBalancerClientAutoConfiguration.java | 10 +- ...torLoadBalancerExchangeFilterFunction.java | 19 ++- ...bleLoadBalancerExchangeFilterFunction.java | 37 +++++- ...dBalancerClientRequestTransformerTest.java | 125 ++++++++++++++++++ ...xchangeFilterFunctionIntegrationTests.java | 2 +- ...adBalancerExchangeFilterFunctionTests.java | 5 +- 9 files changed, 285 insertions(+), 14 deletions(-) create mode 100644 spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/LoadBalancerClientRequestTransformer.java create mode 100644 spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/reactive/LoadBalancerClientRequestTransformerTest.java diff --git a/docs/src/main/asciidoc/spring-cloud-commons.adoc b/docs/src/main/asciidoc/spring-cloud-commons.adoc index 622c04c7..aae3bd93 100644 --- a/docs/src/main/asciidoc/spring-cloud-commons.adoc +++ b/docs/src/main/asciidoc/spring-cloud-commons.adoc @@ -1080,6 +1080,53 @@ public class CustomLoadBalancerConfiguration { TIP: This is also a replacement for Zookeeper `StickyRule`. +=== Transform the load-balanced HTTP request + +You can use the selected `ServiceInstance` to transform the load-balanced HTTP Request. + +For `RestTemplate`, you need to implement and define `LoadBalancerRequestTransformer` as follows: + +[source,java,indent=0] +---- + @Bean + public LoadBalancerRequestTransformer transformer() { + return new LoadBalancerRequestTransformer() { + @Override + public HttpRequest transformRequest(HttpRequest request, ServiceInstance instance) { + return new HttpRequestWrapper(request) { + @Override + public HttpHeaders getHeaders() { + HttpHeaders headers = new HttpHeaders(); + headers.putAll(super.getHeaders()); + headers.add("X-InstanceId", instance.getInstanceId()); + return headers; + } + }; + } + }; + } +---- + +For `WebClient`, you need to implement and define `LoadBalancerClientRequestTransformer` as follows: + +[source,java,indent=0] +---- + @Bean + public LoadBalancerClientRequestTransformer transformer() { + return new LoadBalancerClientRequestTransformer() { + @Override + public ClientRequest transformRequest(ClientRequest request, ServiceInstance instance) { + return ClientRequest.from(request) + .header("X-InstanceId", instance.getInstanceId()) + .build(); + } + }; + } +---- + +If multiple transformers are defined, they are applied in the order in which Beans are defined. +Alternatively, you can use `LoadBalancerRequestTransformer.DEFAULT_ORDER` or `LoadBalancerClientRequestTransformer.DEFAULT_ORDER` to specify the order. + [[spring-cloud-loadbalancer-starter]] === Spring Cloud LoadBalancer Starter diff --git a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ExchangeFilterFunctionUtils.java b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ExchangeFilterFunctionUtils.java index 8e4ce151..e13f4ceb 100644 --- a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ExchangeFilterFunctionUtils.java +++ b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ExchangeFilterFunctionUtils.java @@ -17,7 +17,9 @@ package org.springframework.cloud.client.loadbalancer.reactive; import java.net.URI; +import java.util.List; +import org.springframework.cloud.client.ServiceInstance; import org.springframework.web.reactive.function.client.ClientRequest; import org.springframework.web.reactive.function.client.ExchangeFilterFunction; @@ -33,12 +35,21 @@ public final class ExchangeFilterFunctionUtils { throw new IllegalStateException("Can't instantiate a utility class."); } - static ClientRequest buildClientRequest(ClientRequest request, URI uri) { - return ClientRequest.create(request.method(), uri) + static ClientRequest buildClientRequest(ClientRequest request, URI uri, + ServiceInstance serviceInstance, + List transformers) { + ClientRequest clientRequest = ClientRequest.create(request.method(), uri) .headers(headers -> headers.addAll(request.headers())) .cookies(cookies -> cookies.addAll(request.cookies())) .attributes(attributes -> attributes.putAll(request.attributes())) .body(request.body()).build(); + if (transformers != null) { + for (LoadBalancerClientRequestTransformer transformer : transformers) { + clientRequest = transformer.transformRequest(clientRequest, + serviceInstance); + } + } + return clientRequest; } } diff --git a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/LoadBalancerClientRequestTransformer.java b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/LoadBalancerClientRequestTransformer.java new file mode 100644 index 00000000..2701cf53 --- /dev/null +++ b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/LoadBalancerClientRequestTransformer.java @@ -0,0 +1,39 @@ +/* + * Copyright 2012-2020 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.reactive; + +import org.springframework.cloud.client.ServiceInstance; +import org.springframework.core.annotation.Order; +import org.springframework.web.reactive.function.client.ClientRequest; + +/** + * Allows applications to transform the load-balanced {@link ClientRequest} given the + * chosen {@link ServiceInstance}. + * + * @author Toshiaki Maki + */ +@Order(LoadBalancerClientRequestTransformer.DEFAULT_ORDER) +public interface LoadBalancerClientRequestTransformer { + + /** + * Order for the load balancer request tranformer. + */ + int DEFAULT_ORDER = 0; + + ClientRequest transformRequest(ClientRequest request, ServiceInstance instance); + +} diff --git a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ReactorLoadBalancerClientAutoConfiguration.java b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ReactorLoadBalancerClientAutoConfiguration.java index f4f6a2f1..fd586838 100644 --- a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ReactorLoadBalancerClientAutoConfiguration.java +++ b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ReactorLoadBalancerClientAutoConfiguration.java @@ -16,11 +16,15 @@ package org.springframework.cloud.client.loadbalancer.reactive; +import java.util.Collections; +import java.util.List; + import javax.annotation.PostConstruct; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; +import org.springframework.beans.factory.ObjectProvider; import org.springframework.boot.autoconfigure.condition.ConditionalOnBean; import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; @@ -60,8 +64,10 @@ public class ReactorLoadBalancerClientAutoConfiguration { @ConditionalOnMissingBean @Bean public ReactorLoadBalancerExchangeFilterFunction loadBalancerExchangeFilterFunction( - ReactiveLoadBalancer.Factory loadBalancerFactory) { - return new ReactorLoadBalancerExchangeFilterFunction(loadBalancerFactory); + ReactiveLoadBalancer.Factory loadBalancerFactory, + ObjectProvider> transformers) { + return new ReactorLoadBalancerExchangeFilterFunction(loadBalancerFactory, + transformers.getIfAvailable(Collections::emptyList)); } } diff --git a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ReactorLoadBalancerExchangeFilterFunction.java b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ReactorLoadBalancerExchangeFilterFunction.java index 88fdd673..c0c2fdc7 100644 --- a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ReactorLoadBalancerExchangeFilterFunction.java +++ b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ReactorLoadBalancerExchangeFilterFunction.java @@ -17,6 +17,8 @@ package org.springframework.cloud.client.loadbalancer.reactive; import java.net.URI; +import java.util.Collections; +import java.util.List; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; @@ -46,9 +48,24 @@ public class ReactorLoadBalancerExchangeFilterFunction implements ExchangeFilter private final ReactiveLoadBalancer.Factory loadBalancerFactory; + private final List transformers; + + /** + * @deprecated Deprecated in favor of + * {@link #ReactorLoadBalancerExchangeFilterFunction(ReactiveLoadBalancer.Factory, List)}. + * @param loadBalancerFactory the loadbalancer factory + */ + @Deprecated public ReactorLoadBalancerExchangeFilterFunction( ReactiveLoadBalancer.Factory loadBalancerFactory) { + this(loadBalancerFactory, Collections.emptyList()); + } + + public ReactorLoadBalancerExchangeFilterFunction( + ReactiveLoadBalancer.Factory loadBalancerFactory, + List transformers) { this.loadBalancerFactory = loadBalancerFactory; + this.transformers = transformers; } @Override @@ -82,7 +99,7 @@ public class ReactorLoadBalancerExchangeFilterFunction implements ExchangeFilter serviceId, instance.getUri())); } ClientRequest newRequest = buildClientRequest(request, - reconstructURI(instance, originalUrl)); + reconstructURI(instance, originalUrl), instance, transformers); return next.exchange(newRequest); }); } diff --git a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/RetryableLoadBalancerExchangeFilterFunction.java b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/RetryableLoadBalancerExchangeFilterFunction.java index c213c4ed..4ace9c45 100644 --- a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/RetryableLoadBalancerExchangeFilterFunction.java +++ b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/RetryableLoadBalancerExchangeFilterFunction.java @@ -19,6 +19,7 @@ package org.springframework.cloud.client.loadbalancer.reactive; import java.io.IOException; import java.net.URI; import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.concurrent.TimeoutException; @@ -66,22 +67,46 @@ public class RetryableLoadBalancerExchangeFilterFunction private final ReactiveLoadBalancer.Factory loadBalancerFactory; + private final List transformers; + + /** + * @deprecated Deprecated in favor of + * {@link #RetryableLoadBalancerExchangeFilterFunction(LoadBalancerRetryPolicy, ReactiveLoadBalancer.Factory, LoadBalancerRetryProperties, List)}. + * @param retryPolicy the retry policy + * @param loadBalancerFactory the loadbalancer factory + * @param retryProperties the retry properties + */ + @Deprecated public RetryableLoadBalancerExchangeFilterFunction( LoadBalancerRetryPolicy retryPolicy, ReactiveLoadBalancer.Factory loadBalancerFactory, LoadBalancerRetryProperties retryProperties) { - this.retryPolicy = retryPolicy; - this.loadBalancerFactory = loadBalancerFactory; - this.retryProperties = retryProperties; + this(retryPolicy, loadBalancerFactory, retryProperties, Collections.emptyList()); } + /** + * @deprecated Deprecated in favor of + * {@link #RetryableLoadBalancerExchangeFilterFunction(LoadBalancerRetryPolicy, ReactiveLoadBalancer.Factory, LoadBalancerRetryProperties, List)}. + * @param loadBalancerFactory the loadbalancer factory + * @param retryProperties the retry properties + */ + @Deprecated public RetryableLoadBalancerExchangeFilterFunction( ReactiveLoadBalancer.Factory loadBalancerFactory, LoadBalancerRetryProperties retryProperties) { - this.retryPolicy = new RetryableExchangeFilterFunctionLoadBalancerRetryPolicy( - retryProperties); + this(new RetryableExchangeFilterFunctionLoadBalancerRetryPolicy(retryProperties), + loadBalancerFactory, retryProperties); + } + + public RetryableLoadBalancerExchangeFilterFunction( + LoadBalancerRetryPolicy retryPolicy, + ReactiveLoadBalancer.Factory loadBalancerFactory, + LoadBalancerRetryProperties retryProperties, + List transformers) { + this.retryPolicy = retryPolicy; this.loadBalancerFactory = loadBalancerFactory; this.retryProperties = retryProperties; + this.transformers = transformers; } @SuppressWarnings({ "rawtypes", "unchecked" }) @@ -129,7 +154,7 @@ public class RetryableLoadBalancerExchangeFilterFunction serviceId, instance.getUri())); } ClientRequest newRequest = buildClientRequest(clientRequest, - reconstructURI(instance, originalUrl)); + reconstructURI(instance, originalUrl), instance, transformers); return next.exchange(newRequest).map(clientResponse -> { loadBalancerRetryContext.setClientResponse(clientResponse); if (shouldRetrySameServiceInstance(loadBalancerRetryContext)) { diff --git a/spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/reactive/LoadBalancerClientRequestTransformerTest.java b/spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/reactive/LoadBalancerClientRequestTransformerTest.java new file mode 100644 index 00000000..73a86875 --- /dev/null +++ b/spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/reactive/LoadBalancerClientRequestTransformerTest.java @@ -0,0 +1,125 @@ +/* + * Copyright 2012-2020 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.reactive; + +import java.net.URI; +import java.util.Arrays; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import reactor.core.publisher.Mono; + +import org.springframework.cloud.client.ServiceInstance; +import org.springframework.cloud.client.loadbalancer.LoadBalancerRetryProperties; +import org.springframework.http.HttpHeaders; +import org.springframework.http.HttpMethod; +import org.springframework.http.HttpStatus; +import org.springframework.util.LinkedMultiValueMap; +import org.springframework.web.reactive.function.client.ClientRequest; +import org.springframework.web.reactive.function.client.ClientResponse; +import org.springframework.web.reactive.function.client.ExchangeFunction; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Tests for {@link LoadBalancerClientRequestTransformer}. + * + * @author Toshiaki Maki + */ +class LoadBalancerClientRequestTransformerTest { + + private final LoadBalancerRetryProperties properties = new LoadBalancerRetryProperties(); + + private final LoadBalancerRetryPolicy policy = new RetryableExchangeFilterFunctionLoadBalancerRetryPolicy( + properties); + + private final ReactiveLoadBalancer.Factory factory = mock( + ReactiveLoadBalancer.Factory.class); + + private final ClientRequest clientRequest = mock(ClientRequest.class); + + private final ClientResponse clientResponse = mock(ClientResponse.class); + + private final ExchangeFunction next = mock(ExchangeFunction.class); + + @BeforeEach + void setUp() { + when(factory.getInstance("testServiceId")) + .thenReturn(new TestReactiveLoadBalancer()); + when(clientRequest.method()).thenReturn(HttpMethod.GET); + when(clientRequest.url()).thenReturn(URI.create("http://testServiceId")); + when(clientRequest.headers()).thenReturn(new HttpHeaders()); + when(clientRequest.cookies()).thenReturn(new LinkedMultiValueMap<>()); + when(next.exchange(any())).thenReturn(Mono.just(clientResponse)); + when(clientResponse.statusCode()).thenReturn(HttpStatus.OK); + } + + @Test + void transformReactorLoadBalancerExchangeFilterFunction() { + ArgumentCaptor captor = ArgumentCaptor + .forClass(ClientRequest.class); + ReactorLoadBalancerExchangeFilterFunction filterFunction = new ReactorLoadBalancerExchangeFilterFunction( + factory, Arrays.asList(new Transformer1(), new Transformer2())); + filterFunction.filter(clientRequest, next).subscribe(); + verify(next).exchange(captor.capture()); + HttpHeaders headers = captor.getValue().headers(); + assertThat(headers.getFirst("X-ServiceId")).isEqualTo("testServiceId"); + assertThat(headers.getFirst("X-InstanceId")).isEqualTo("testServiceId"); + } + + @Test + void transformRetryableLoadBalancerExchangeFilterFunction() { + ArgumentCaptor captor = ArgumentCaptor + .forClass(ClientRequest.class); + RetryableLoadBalancerExchangeFilterFunction filterFunction = new RetryableLoadBalancerExchangeFilterFunction( + policy, factory, properties, + Arrays.asList(new Transformer1(), new Transformer2())); + filterFunction.filter(clientRequest, next).subscribe(); + verify(next).exchange(captor.capture()); + HttpHeaders headers = captor.getValue().headers(); + assertThat(headers.getFirst("X-ServiceId")).isEqualTo("testServiceId"); + assertThat(headers.getFirst("X-InstanceId")).isEqualTo("testServiceId"); + } + + class Transformer1 implements LoadBalancerClientRequestTransformer { + + @Override + public ClientRequest transformRequest(ClientRequest request, + ServiceInstance instance) { + return ClientRequest.from(request) + .header("X-ServiceId", instance.getServiceId()).build(); + } + + } + + class Transformer2 implements LoadBalancerClientRequestTransformer { + + @Override + public ClientRequest transformRequest(ClientRequest request, + ServiceInstance instance) { + return ClientRequest.from(request) + .header("X-InstanceId", instance.getInstanceId()).build(); + } + + } + +} diff --git a/spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/reactive/RetryableLoadBalancerExchangeFilterFunctionIntegrationTests.java b/spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/reactive/RetryableLoadBalancerExchangeFilterFunctionIntegrationTests.java index 9f57080c..64f64446 100644 --- a/spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/reactive/RetryableLoadBalancerExchangeFilterFunctionIntegrationTests.java +++ b/spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/reactive/RetryableLoadBalancerExchangeFilterFunctionIntegrationTests.java @@ -210,7 +210,7 @@ class RetryableLoadBalancerExchangeFilterFunctionIntegrationTests { return new RetryableLoadBalancerExchangeFilterFunction( new RetryableExchangeFilterFunctionLoadBalancerRetryPolicy( properties), - factory, properties); + factory, properties, Collections.emptyList()); } } diff --git a/spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/reactive/RetryableLoadBalancerExchangeFilterFunctionTests.java b/spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/reactive/RetryableLoadBalancerExchangeFilterFunctionTests.java index c0ed3a6d..a39ac064 100644 --- a/spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/reactive/RetryableLoadBalancerExchangeFilterFunctionTests.java +++ b/spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/reactive/RetryableLoadBalancerExchangeFilterFunctionTests.java @@ -18,6 +18,7 @@ package org.springframework.cloud.client.loadbalancer.reactive; import java.io.IOException; import java.net.URI; +import java.util.Collections; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -58,7 +59,7 @@ class RetryableLoadBalancerExchangeFilterFunctionTests { ReactiveLoadBalancer.Factory.class); private final RetryableLoadBalancerExchangeFilterFunction filterFunction = new RetryableLoadBalancerExchangeFilterFunction( - policy, factory, properties); + policy, factory, properties, Collections.emptyList()); private final ClientRequest clientRequest = mock(ClientRequest.class); @@ -145,7 +146,7 @@ class RetryableLoadBalancerExchangeFilterFunctionTests { LoadBalancerRetryPolicy policy = new RetryableExchangeFilterFunctionLoadBalancerRetryPolicy( properties); RetryableLoadBalancerExchangeFilterFunction filterFunction = new RetryableLoadBalancerExchangeFilterFunction( - policy, factory, properties); + policy, factory, properties, Collections.emptyList()); when(clientRequest.method()).thenReturn(HttpMethod.POST); when(clientResponse.statusCode()).thenReturn(HttpStatus.NOT_FOUND); when(next.exchange(any())).thenReturn(Mono.just(clientResponse));