diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/config/GatewayLoadBalancerClientAutoConfiguration.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/config/GatewayLoadBalancerClientAutoConfiguration.java index b6848a34..119c3adc 100644 --- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/config/GatewayLoadBalancerClientAutoConfiguration.java +++ b/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); } diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/LoadBalancerClientFilter.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/LoadBalancerClientFilter.java index d997611b..ff12c753 100644 --- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/LoadBalancerClientFilter.java +++ b/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; diff --git a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/LoadBalancerClientFilterTests.java b/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/LoadBalancerClientFilterTests.java index acaefd52..d3f45095 100644 --- a/spring-cloud-gateway-core/src/test/java/org/springframework/cloud/gateway/filter/LoadBalancerClientFilterTests.java +++ b/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); }