Allows LoadBalancerClientFilter to be extended (#448)

Makes `LoadBalancerClient` protected and creates protected `choose(ServerWebExchange)` method.

spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/config/GatewayLoadBalancerClientAutoConfiguration.java

@@ -20,6 +20,7 @@ package org.springframework.cloud.gateway.config;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
 import org.springframework.cloud.gateway.filter.LoadBalancerClientFilter;
 import org.springframework.cloud.netflix.ribbon.RibbonAutoConfiguration;
@@ -39,6 +40,7 @@ public class GatewayLoadBalancerClientAutoConfiguration {
 
 	@Bean
 	@ConditionalOnBean(LoadBalancerClient.class)
+	@ConditionalOnMissingBean(LoadBalancerClientFilter.class)
 	public LoadBalancerClientFilter loadBalancerClientFilter(LoadBalancerClient client) {
 		return new LoadBalancerClientFilter(client);
 	}

spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/LoadBalancerClientFilter.java

@@ -43,7 +43,7 @@ public class LoadBalancerClientFilter implements GlobalFilter, Ordered {
 	private static final Log log = LogFactory.getLog(LoadBalancerClientFilter.class);
 	public static final int LOAD_BALANCER_CLIENT_FILTER_ORDER = 10100;
 
-	private final LoadBalancerClient loadBalancer;
+	protected final LoadBalancerClient loadBalancer;
 
 	public LoadBalancerClientFilter(LoadBalancerClient loadBalancer) {
 		this.loadBalancer = loadBalancer;
@@ -66,7 +66,7 @@ public class LoadBalancerClientFilter implements GlobalFilter, Ordered {
 
 		log.trace("LoadBalancerClientFilter url before: " + url);
 
-		final ServiceInstance instance = loadBalancer.choose(url.getHost());
+		final ServiceInstance instance = choose(exchange);
 
 		if (instance == null) {
 			throw new NotFoundException("Unable to find instance for " + url.getHost());
@@ -88,6 +88,10 @@ public class LoadBalancerClientFilter implements GlobalFilter, Ordered {
 		return chain.filter(exchange);
 	}
 
+	protected ServiceInstance choose(ServerWebExchange exchange) {
+		return loadBalancer.choose(((URI) exchange.getAttribute(GATEWAY_REQUEST_URL_ATTR)).getHost());
+	}
+
 	class DelegatingServiceInstance implements ServiceInstance {
 		final ServiceInstance delegate;
 		private String overrideScheme;

spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/LoadBalancerClientFilterTests.java

@@ -1,7 +1,6 @@
 package org.springframework.cloud.gateway.filter;
 
 import java.net.URI;
-import java.util.Collections;
 import java.util.LinkedHashSet;
 
 import com.netflix.loadbalancer.ILoadBalancer;
@@ -23,6 +22,7 @@ import org.springframework.cloud.netflix.ribbon.SpringClientFactory;
 import org.springframework.http.HttpMethod;
 import org.springframework.mock.http.server.reactive.MockServerHttpRequest;
 import org.springframework.mock.web.server.MockServerWebExchange;
+import org.springframework.util.StringUtils;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.util.UriComponentsBuilder;
 
@@ -227,6 +227,52 @@ public class LoadBalancerClientFilterTests {
 		verifyZeroInteractions(loadBalancerClient);
 	}
 
+	@Test
+	public void shouldSelectSpecifiedServer() {
+		URI uri1 = UriComponentsBuilder.fromUriString("lb://myservice").port(11111).build().toUri();
+		URI uri2 = UriComponentsBuilder.fromUriString("lb://myservice").port(22222).build().toUri();
+
+		SpringClientFactory clientFactory = mock(SpringClientFactory.class);
+		ILoadBalancer loadBalancer = mock(ILoadBalancer.class);
+		when(clientFactory.getLoadBalancerContext("myservice")).thenReturn(new RibbonLoadBalancerContext(loadBalancer));
+		when(clientFactory.getLoadBalancer("myservice")).thenReturn(loadBalancer);
+
+		when(loadBalancer.chooseServer("11111")).thenReturn(new Server("myservice-host1", 8081));
+		when(loadBalancer.chooseServer("22222")).thenReturn(new Server("myservice-host2", 8081));
+
+		LoadBalancerClient loadBalancerClient = new RibbonLoadBalancerClient(clientFactory) {
+			private String loadBalancerKey;
+			public ServiceInstance choose(String serviceId) {
+				String[] strings = serviceId.split("<<>>");
+				loadBalancerKey = strings[1];
+				return super.choose(strings[0]);
+			}
+			protected Server getServer(ILoadBalancer loadBalancer) {
+				return loadBalancer == null ? null : loadBalancer.chooseServer(StringUtils.isEmpty(loadBalancerKey) ? "default" : loadBalancerKey);
+			}
+		};
+
+		LoadBalancerClientFilter loadBalancerClientFilter = new LoadBalancerClientFilter(loadBalancerClient) {
+			protected ServiceInstance choose(ServerWebExchange exchange) {
+				URI attribute = (URI) exchange.getAttribute(GATEWAY_REQUEST_URL_ATTR);
+				return loadBalancer.choose(attribute.getHost() + "<<>>" + attribute.getPort());
+			}
+		};
+
+		MockServerHttpRequest request = MockServerHttpRequest
+				.get("http://localhost/get")
+				.build();
+		ServerWebExchange exchange = MockServerWebExchange.from(request);
+
+		exchange.getAttributes().put(GATEWAY_REQUEST_URL_ATTR, uri1);
+		loadBalancerClientFilter.filter(exchange, chain);
+		assertThat(((URI)exchange.getAttributes().get(GATEWAY_REQUEST_URL_ATTR)).getHost()).isEqualTo("myservice-host1");
+
+		exchange.getAttributes().put(GATEWAY_REQUEST_URL_ATTR, uri2);
+		loadBalancerClientFilter.filter(exchange, chain);
+		assertThat(((URI)exchange.getAttributes().get(GATEWAY_REQUEST_URL_ATTR)).getHost()).isEqualTo("myservice-host2");
+	}
+
 	private ServerWebExchange testFilter(MockServerHttpRequest request, URI uri) {
 		return testFilter(MockServerWebExchange.from(request), uri);
 	}