Pass HTTP server requests and hints to LoadBalancer. (#2062)

spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/config/GatewayLoadBalancerProperties.java

@@ -22,7 +22,7 @@ import org.springframework.boot.context.properties.ConfigurationProperties;
  * @author Ryan Baxter
  */
 @ConfigurationProperties("spring.cloud.gateway.loadbalancer")
-public class LoadBalancerProperties {
+public class GatewayLoadBalancerProperties {
 
 	private boolean use404;
 

spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/config/GatewayNoLoadBalancerClientAutoConfiguration.java

@@ -44,13 +44,13 @@ import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.G
 @Configuration(proxyBeanMethods = false)
 @ConditionalOnMissingClass("org.springframework.cloud.loadbalancer.core.ReactorLoadBalancer")
 @ConditionalOnMissingBean(ReactiveLoadBalancer.class)
-@EnableConfigurationProperties(LoadBalancerProperties.class)
+@EnableConfigurationProperties(GatewayLoadBalancerProperties.class)
 @AutoConfigureAfter(GatewayReactiveLoadBalancerClientAutoConfiguration.class)
 public class GatewayNoLoadBalancerClientAutoConfiguration {
 
 	@Bean
 	@ConditionalOnMissingBean(ReactiveLoadBalancerClientFilter.class)
-	public NoLoadBalancerClientFilter noLoadBalancerClientFilter(LoadBalancerProperties properties) {
+	public NoLoadBalancerClientFilter noLoadBalancerClientFilter(GatewayLoadBalancerProperties properties) {
 		return new NoLoadBalancerClientFilter(properties.isUse404());
 	}
 

spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/config/GatewayReactiveLoadBalancerClientAutoConfiguration.java

@@ -21,6 +21,7 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
+import org.springframework.cloud.client.loadbalancer.reactive.LoadBalancerProperties;
 import org.springframework.cloud.client.loadbalancer.reactive.ReactiveLoadBalancer;
 import org.springframework.cloud.gateway.config.conditional.ConditionalOnEnabledGlobalFilter;
 import org.springframework.cloud.gateway.filter.ReactiveLoadBalancerClientFilter;
@@ -39,7 +40,7 @@ import org.springframework.web.reactive.DispatcherHandler;
 @Configuration(proxyBeanMethods = false)
 @ConditionalOnClass({ ReactiveLoadBalancer.class, LoadBalancerAutoConfiguration.class, DispatcherHandler.class })
 @AutoConfigureAfter(LoadBalancerAutoConfiguration.class)
-@EnableConfigurationProperties(LoadBalancerProperties.class)
+@EnableConfigurationProperties(GatewayLoadBalancerProperties.class)
 public class GatewayReactiveLoadBalancerClientAutoConfiguration {
 
 	@Bean
@@ -47,8 +48,8 @@ public class GatewayReactiveLoadBalancerClientAutoConfiguration {
 	@ConditionalOnMissingBean(ReactiveLoadBalancerClientFilter.class)
 	@ConditionalOnEnabledGlobalFilter
 	public ReactiveLoadBalancerClientFilter gatewayLoadBalancerClientFilter(LoadBalancerClientFactory clientFactory,
-			LoadBalancerProperties properties) {
-		return new ReactiveLoadBalancerClientFilter(clientFactory, properties);
+			GatewayLoadBalancerProperties properties, LoadBalancerProperties loadBalancerProperties) {
+		return new ReactiveLoadBalancerClientFilter(clientFactory, properties, loadBalancerProperties);
 	}
 
 }

spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/filter/ReactiveLoadBalancerClientFilter.java

@@ -17,17 +17,19 @@
 package org.springframework.cloud.gateway.filter;
 
 import java.net.URI;
+import java.util.Map;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import reactor.core.publisher.Mono;
 
 import org.springframework.cloud.client.ServiceInstance;
+import org.springframework.cloud.client.loadbalancer.DefaultRequest;
 import org.springframework.cloud.client.loadbalancer.LoadBalancerUriTools;
-import org.springframework.cloud.client.loadbalancer.Request;
 import org.springframework.cloud.client.loadbalancer.Response;
-import org.springframework.cloud.client.loadbalancer.reactive.ReactiveLoadBalancer;
-import org.springframework.cloud.gateway.config.LoadBalancerProperties;
+import org.springframework.cloud.client.loadbalancer.ServerHttpRequestContext;
+import org.springframework.cloud.client.loadbalancer.reactive.LoadBalancerProperties;
+import org.springframework.cloud.gateway.config.GatewayLoadBalancerProperties;
 import org.springframework.cloud.gateway.support.DelegatingServiceInstance;
 import org.springframework.cloud.gateway.support.NotFoundException;
 import org.springframework.cloud.loadbalancer.core.ReactorLoadBalancer;
@@ -48,6 +50,7 @@ import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.a
  * @author Tim Ysewyn
  * @author Olga Maciaszek-Sharma
  */
+@SuppressWarnings("rawtypes")
 public class ReactiveLoadBalancerClientFilter implements GlobalFilter, Ordered {
 
 	private static final Log log = LogFactory.getLog(ReactiveLoadBalancerClientFilter.class);
@@ -59,12 +62,15 @@ public class ReactiveLoadBalancerClientFilter implements GlobalFilter, Ordered {
 
 	private final LoadBalancerClientFactory clientFactory;
 
-	private LoadBalancerProperties properties;
+	private final GatewayLoadBalancerProperties properties;
+
+	private final LoadBalancerProperties loadBalancerProperties;
 
 	public ReactiveLoadBalancerClientFilter(LoadBalancerClientFactory clientFactory,
-			LoadBalancerProperties properties) {
+			GatewayLoadBalancerProperties properties, LoadBalancerProperties loadBalancerProperties) {
 		this.clientFactory = clientFactory;
 		this.properties = properties;
+		this.loadBalancerProperties = loadBalancerProperties;
 	}
 
 	@Override
@@ -120,20 +126,23 @@ public class ReactiveLoadBalancerClientFilter implements GlobalFilter, Ordered {
 		return LoadBalancerUriTools.reconstructURI(serviceInstance, original);
 	}
 
-	@SuppressWarnings("deprecation")
 	private Mono<Response<ServiceInstance>> choose(ServerWebExchange exchange) {
 		URI uri = exchange.getAttribute(GATEWAY_REQUEST_URL_ATTR);
-		ReactorLoadBalancer<ServiceInstance> loadBalancer = this.clientFactory.getInstance(uri.getHost(),
+		String serviceId = uri.getHost();
+		ReactorLoadBalancer<ServiceInstance> loadBalancer = this.clientFactory.getInstance(serviceId,
 				ReactorServiceInstanceLoadBalancer.class);
 		if (loadBalancer == null) {
-			throw new NotFoundException("No loadbalancer available for " + uri.getHost());
+			throw new NotFoundException("No loadbalancer available for " + serviceId);
 		}
-		return loadBalancer.choose(createRequest());
+		DefaultRequest<ServerHttpRequestContext> lbRequest = new DefaultRequest<>(new ServerHttpRequestContext(
+				exchange.getRequest(), getHint(serviceId, loadBalancerProperties.getHint())));
+		return loadBalancer.choose(lbRequest);
 	}
 
-	@SuppressWarnings("deprecation")
-	private Request createRequest() {
-		return ReactiveLoadBalancer.REQUEST;
+	private String getHint(String serviceId, Map<String, String> hints) {
+		String defaultHint = hints.getOrDefault("default", "default");
+		String hintPropertyValue = hints.get(serviceId);
+		return hintPropertyValue != null ? hintPropertyValue : defaultHint;
 	}
 
 }

spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/discovery/ReactiveGatewayDiscoveryClientAutoConfigurationTests.java

@@ -26,7 +26,7 @@ import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.SpringBootConfiguration;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.test.context.SpringBootTest;
-import org.springframework.cloud.gateway.config.LoadBalancerProperties;
+import org.springframework.cloud.gateway.config.GatewayLoadBalancerProperties;
 import org.springframework.cloud.gateway.route.RouteDefinition;
 import org.springframework.test.context.junit4.SpringRunner;
 
@@ -46,7 +46,7 @@ public class ReactiveGatewayDiscoveryClientAutoConfigurationTests {
 		private DiscoveryClientRouteDefinitionLocator locator;
 
 		@Autowired(required = false)
-		private LoadBalancerProperties properties;
+		private GatewayLoadBalancerProperties properties;
 
 		@Test
 		public void routeLocatorBeanExists() {

spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/filter/ReactiveLoadBalancerClientFilterTests.java

@@ -17,8 +17,11 @@
 package org.springframework.cloud.gateway.filter;
 
 import java.net.URI;
+import java.util.HashMap;
 import java.util.LinkedHashSet;
+import java.util.Map;
 
+import org.jetbrains.annotations.NotNull;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -30,7 +33,11 @@ import reactor.core.publisher.Mono;
 
 import org.springframework.cloud.client.DefaultServiceInstance;
 import org.springframework.cloud.client.ServiceInstance;
-import org.springframework.cloud.gateway.config.LoadBalancerProperties;
+import org.springframework.cloud.client.loadbalancer.DefaultResponse;
+import org.springframework.cloud.client.loadbalancer.Request;
+import org.springframework.cloud.client.loadbalancer.ServerHttpRequestContext;
+import org.springframework.cloud.client.loadbalancer.reactive.LoadBalancerProperties;
+import org.springframework.cloud.gateway.config.GatewayLoadBalancerProperties;
 import org.springframework.cloud.gateway.support.NotFoundException;
 import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
 import org.springframework.cloud.loadbalancer.core.RoundRobinLoadBalancer;
@@ -44,6 +51,9 @@ import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.util.UriComponentsBuilder;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.argThat;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoInteractions;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
@@ -65,7 +75,7 @@ public class ReactiveLoadBalancerClientFilterTests {
 
 	private ServerWebExchange exchange;
 
-	private LoadBalancerProperties properties;
+	private GatewayLoadBalancerProperties properties;
 
 	@Mock
 	private GatewayFilterChain chain;
@@ -73,12 +83,15 @@ public class ReactiveLoadBalancerClientFilterTests {
 	@Mock
 	private LoadBalancerClientFactory clientFactory;
 
+	@Mock
+	private LoadBalancerProperties loadBalancerProperties;
+
 	@InjectMocks
 	private ReactiveLoadBalancerClientFilter filter;
 
 	@Before
 	public void setup() {
-		properties = new LoadBalancerProperties();
+		properties = new GatewayLoadBalancerProperties();
 		exchange = MockServerWebExchange.from(MockServerHttpRequest.get("/mypath").build());
 	}
 
@@ -237,7 +250,8 @@ public class ReactiveLoadBalancerClientFilterTests {
 				ServiceInstanceListSuppliers.toProvider("service1"), "service1", -1);
 		when(clientFactory.getInstance("service1", ReactorServiceInstanceLoadBalancer.class)).thenReturn(loadBalancer);
 		properties.setUse404(true);
-		ReactiveLoadBalancerClientFilter filter = new ReactiveLoadBalancerClientFilter(clientFactory, properties);
+		ReactiveLoadBalancerClientFilter filter = new ReactiveLoadBalancerClientFilter(clientFactory, properties,
+				loadBalancerProperties);
 		when(chain.filter(exchange)).thenReturn(Mono.empty());
 		try {
 			filter.filter(exchange, chain).block();
@@ -270,6 +284,40 @@ public class ReactiveLoadBalancerClientFilterTests {
 		verifyNoMoreInteractions(chain);
 	}
 
+	@SuppressWarnings({ "rawtypes" })
+	@Test
+	public void shouldPassRequestToLoadBalancer() {
+		String hint = "test";
+		when(loadBalancerProperties.getHint()).thenReturn(buildHints(hint));
+		MockServerHttpRequest request = MockServerHttpRequest.get("http://localhost/get?a=b").build();
+		URI lbUri = URI.create("lb://service1?a=b");
+		ServerWebExchange serverWebExchange = mock(ServerWebExchange.class);
+		when(serverWebExchange.getAttribute(GATEWAY_REQUEST_URL_ATTR)).thenReturn(lbUri);
+		when(serverWebExchange.getRequiredAttribute(GATEWAY_ORIGINAL_REQUEST_URL_ATTR))
+				.thenReturn(new LinkedHashSet<>());
+		when(serverWebExchange.getRequest()).thenReturn(request);
+		RoundRobinLoadBalancer loadBalancer = mock(RoundRobinLoadBalancer.class);
+		when(loadBalancer.choose(any(Request.class))).thenReturn(Mono.just(
+				new DefaultResponse(new DefaultServiceInstance("myservice1", "myservice", "localhost", 8080, false))));
+		when(clientFactory.getInstance("service1", ReactorServiceInstanceLoadBalancer.class)).thenReturn(loadBalancer);
+		when(chain.filter(any())).thenReturn(Mono.empty());
+
+		filter.filter(serverWebExchange, chain);
+
+		verify(loadBalancer)
+				.choose(argThat((Request passedRequest) -> ((ServerHttpRequestContext) passedRequest.getContext())
+						.getClientRequest().equals(request)
+						&& ((ServerHttpRequestContext) passedRequest.getContext()).getHint().equals(hint)));
+
+	}
+
+	@NotNull
+	private Map<String, String> buildHints(String hint) {
+		Map<String, String> hints = new HashMap<>();
+		hints.put("default", hint);
+		return hints;
+	}
+
 	private ServerWebExchange testFilter(MockServerHttpRequest request, URI uri) {
 		return testFilter(MockServerWebExchange.from(request), uri);
 	}
@@ -286,7 +334,8 @@ public class ReactiveLoadBalancerClientFilterTests {
 				"service1", -1);
 		when(clientFactory.getInstance("service1", ReactorServiceInstanceLoadBalancer.class)).thenReturn(loadBalancer);
 
-		ReactiveLoadBalancerClientFilter filter = new ReactiveLoadBalancerClientFilter(clientFactory, properties);
+		ReactiveLoadBalancerClientFilter filter = new ReactiveLoadBalancerClientFilter(clientFactory, properties,
+				loadBalancerProperties);
 		filter.filter(exchange, chain).block();
 
 		return captor.getValue();

spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterTests.java

@@ -61,12 +61,12 @@ public class RedisRateLimiterTests extends BaseWebClientTests {
 	private RedisRateLimiter rateLimiter;
 
 	@BeforeEach
-	public void setUp() throws Exception {
+	public void setUp() {
 		assumeThat("Ignore on Circle", System.getenv("CIRCLECI"), is(nullValue()));
 	}
 
 	@AfterEach
-	public void tearDown() throws Exception {
+	public void tearDown() {
 		rateLimiter.setIncludeHeaders(true);
 	}