Support to set loadBalancerKey (#50)

* Support to override CachingSpringLoadBalancerFactory to customize the FeignLoadBalancer instance.
This commit is contained in:
Kerwin Bryant
2018-09-13 02:38:15 +08:00
committed by Ryan Baxter
parent 15bda96e91
commit 3e3e486648
3 changed files with 44 additions and 4 deletions

View File

@@ -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<String, FeignLoadBalancer> cache = new ConcurrentReferenceHashMap<>();

View File

@@ -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;
}

View File

@@ -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<FeignLoadBalancer.RibbonResponse> 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));
}
}