From 3e3e486648fd836986dcfdd6b95e3e2d5f17f127 Mon Sep 17 00:00:00 2001 From: Kerwin Bryant Date: Thu, 13 Sep 2018 02:38:15 +0800 Subject: [PATCH] Support to set loadBalancerKey (#50) * Support to override CachingSpringLoadBalancerFactory to customize the FeignLoadBalancer instance. --- .../CachingSpringLoadBalancerFactory.java | 4 +-- .../openfeign/ribbon/FeignLoadBalancer.java | 11 +++++-- .../ribbon/FeignLoadBalancerTests.java | 33 +++++++++++++++++++ 3 files changed, 44 insertions(+), 4 deletions(-) diff --git a/spring-cloud-openfeign-core/src/main/java/org/springframework/cloud/openfeign/ribbon/CachingSpringLoadBalancerFactory.java b/spring-cloud-openfeign-core/src/main/java/org/springframework/cloud/openfeign/ribbon/CachingSpringLoadBalancerFactory.java index d64a3c9e..e47ea3ab 100644 --- a/spring-cloud-openfeign-core/src/main/java/org/springframework/cloud/openfeign/ribbon/CachingSpringLoadBalancerFactory.java +++ b/spring-cloud-openfeign-core/src/main/java/org/springframework/cloud/openfeign/ribbon/CachingSpringLoadBalancerFactory.java @@ -36,8 +36,8 @@ import com.netflix.loadbalancer.ILoadBalancer; */ public class CachingSpringLoadBalancerFactory { - private final SpringClientFactory factory; - private LoadBalancedRetryFactory loadBalancedRetryFactory = null; + protected final SpringClientFactory factory; + protected LoadBalancedRetryFactory loadBalancedRetryFactory = null; private volatile Map cache = new ConcurrentReferenceHashMap<>(); diff --git a/spring-cloud-openfeign-core/src/main/java/org/springframework/cloud/openfeign/ribbon/FeignLoadBalancer.java b/spring-cloud-openfeign-core/src/main/java/org/springframework/cloud/openfeign/ribbon/FeignLoadBalancer.java index e8aa2286..928b4157 100644 --- a/spring-cloud-openfeign-core/src/main/java/org/springframework/cloud/openfeign/ribbon/FeignLoadBalancer.java +++ b/spring-cloud-openfeign-core/src/main/java/org/springframework/cloud/openfeign/ribbon/FeignLoadBalancer.java @@ -118,7 +118,7 @@ public class FeignLoadBalancer extends private final Request request; private final Client client; - RibbonRequest(Client client, Request request, URI uri) { + protected RibbonRequest(Client client, Request request, URI uri) { this.client = client; setUri(uri); this.request = toRequest(request); @@ -170,6 +170,13 @@ public class FeignLoadBalancer extends }; } + public Request getRequest() { + return request; + } + + public Client getClient() { + return client; + } @Override public Object clone() { @@ -182,7 +189,7 @@ public class FeignLoadBalancer extends private final URI uri; private final Response response; - RibbonResponse(URI uri, Response response) { + protected RibbonResponse(URI uri, Response response) { this.uri = uri; this.response = response; } diff --git a/spring-cloud-openfeign-core/src/test/java/org/springframework/cloud/openfeign/ribbon/FeignLoadBalancerTests.java b/spring-cloud-openfeign-core/src/test/java/org/springframework/cloud/openfeign/ribbon/FeignLoadBalancerTests.java index 630647fa..7576e186 100644 --- a/spring-cloud-openfeign-core/src/test/java/org/springframework/cloud/openfeign/ribbon/FeignLoadBalancerTests.java +++ b/spring-cloud-openfeign-core/src/test/java/org/springframework/cloud/openfeign/ribbon/FeignLoadBalancerTests.java @@ -23,6 +23,9 @@ import java.util.Collections; import java.util.HashMap; import java.util.Map; +import com.netflix.loadbalancer.BaseLoadBalancer; +import com.netflix.loadbalancer.RoundRobinRule; +import com.netflix.loadbalancer.reactive.LoadBalancerCommand; import org.junit.Before; import org.junit.Test; import org.mockito.Mock; @@ -168,4 +171,34 @@ public class FeignLoadBalancerTests { } + @Test + public void testOverrideFeignLoadBalancer() throws Exception { + when(this.config.get(IsSecure)).thenReturn(false); + Server server1 = new Server("foo", 6666); + Server server2 = new Server("foo", 7777); + BaseLoadBalancer baseLoadBalancer = new BaseLoadBalancer(); + baseLoadBalancer.setRule(new RoundRobinRule() { + @Override + public Server choose(Object loadBalancerKey) { + return loadBalancerKey == null ? server2 : server1; + } + }); + + this.feignLoadBalancer = new FeignLoadBalancer(baseLoadBalancer, this.config, + this.inspector) { + protected void customizeLoadBalancerCommandBuilder(final FeignLoadBalancer.RibbonRequest request, final IClientConfig config, + final LoadBalancerCommand.Builder builder) { + builder.withServerLocator(request.getRequest().headers().get("c_ip")); + } + }; + Request request = new RequestTemplate().method("GET").request(); + RibbonResponse resp = this.feignLoadBalancer.executeWithLoadBalancer(new RibbonRequest(this.delegate, request, + new URI(request.url())), null); + assertThat(resp.getRequestedURI().getPort(), is(7777)); + request = new RequestTemplate().method("GET").header("c_ip", "666").request(); + resp = this.feignLoadBalancer.executeWithLoadBalancer(new RibbonRequest(this.delegate, request, + new URI(request.url())), null); + assertThat(resp.getRequestedURI().getPort(), is(6666)); + } + }