diff --git a/spring-cloud-netflix-core/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerClient.java b/spring-cloud-netflix-core/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerClient.java index 500892d6..84dbf271 100644 --- a/spring-cloud-netflix-core/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerClient.java +++ b/spring-cloud-netflix-core/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerClient.java @@ -12,4 +12,12 @@ public interface LoadBalancerClient { * @return */ public ServiceInstance choose(String serviceId); + + /** + * Choose a {@see ServiceInstance} from the LoadBalancer for the specified service + * @param serviceId The serviceId to use to look up the LoadBalancer + * @param request allows implementations to execute pre and post actions such as incrementing metrics + * @return + */ + public T choose(String serviceId, LoadBalancerRequest request); } diff --git a/spring-cloud-netflix-core/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerRequest.java b/spring-cloud-netflix-core/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerRequest.java new file mode 100644 index 00000000..396848d2 --- /dev/null +++ b/spring-cloud-netflix-core/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerRequest.java @@ -0,0 +1,10 @@ +package org.springframework.cloud.client.loadbalancer; + +import org.springframework.cloud.client.ServiceInstance; + +/** + * @author Spencer Gibb + */ +public interface LoadBalancerRequest { + public T apply(ServiceInstance instance); +} diff --git a/spring-cloud-netflix-core/src/main/java/org/springframework/cloud/netflix/ribbon/RibbonInterceptor.java b/spring-cloud-netflix-core/src/main/java/org/springframework/cloud/netflix/ribbon/RibbonInterceptor.java index 2cf7ea47..73fd6673 100644 --- a/spring-cloud-netflix-core/src/main/java/org/springframework/cloud/netflix/ribbon/RibbonInterceptor.java +++ b/spring-cloud-netflix-core/src/main/java/org/springframework/cloud/netflix/ribbon/RibbonInterceptor.java @@ -2,6 +2,7 @@ package org.springframework.cloud.netflix.ribbon; import org.springframework.cloud.client.ServiceInstance; import org.springframework.cloud.client.loadbalancer.LoadBalancerClient; +import org.springframework.cloud.client.loadbalancer.LoadBalancerRequest; import org.springframework.http.HttpRequest; import org.springframework.http.client.ClientHttpRequestExecution; import org.springframework.http.client.ClientHttpRequestInterceptor; @@ -28,14 +29,18 @@ public class RibbonInterceptor implements ClientHttpRequestInterceptor { HttpRequestWrapper wrapper = new HttpRequestWrapper(request) { @Override public URI getURI() { - URI originalUri = super.getURI(); + final URI originalUri = super.getURI(); String serviceName = originalUri.getHost(); - ServiceInstance instance = loadBalancer.choose(serviceName); - URI uri = UriComponentsBuilder.fromUri(originalUri) - .host(instance.getHost()) - .port(instance.getPort()) - .build() - .toUri(); + URI uri = loadBalancer.choose(serviceName, new LoadBalancerRequest() { + @Override + public URI apply(ServiceInstance instance) { + return UriComponentsBuilder.fromUri(originalUri) + .host(instance.getHost()) + .port(instance.getPort()) + .build() + .toUri(); + } + }); return uri; } }; diff --git a/spring-cloud-netflix-core/src/main/java/org/springframework/cloud/netflix/ribbon/RibbonLoadBalancerClient.java b/spring-cloud-netflix-core/src/main/java/org/springframework/cloud/netflix/ribbon/RibbonLoadBalancerClient.java index a92343a0..d177c194 100644 --- a/spring-cloud-netflix-core/src/main/java/org/springframework/cloud/netflix/ribbon/RibbonLoadBalancerClient.java +++ b/spring-cloud-netflix-core/src/main/java/org/springframework/cloud/netflix/ribbon/RibbonLoadBalancerClient.java @@ -4,6 +4,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import com.netflix.loadbalancer.AbstractLoadBalancer; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.cloud.client.ServiceInstance; import org.springframework.cloud.client.loadbalancer.LoadBalancerClient; @@ -11,6 +12,7 @@ import org.springframework.cloud.client.loadbalancer.LoadBalancerClient; import com.netflix.loadbalancer.BaseLoadBalancer; import com.netflix.loadbalancer.ILoadBalancer; import com.netflix.loadbalancer.Server; +import org.springframework.cloud.client.loadbalancer.LoadBalancerRequest; /** * @author Spencer Gibb @@ -34,20 +36,45 @@ public class RibbonLoadBalancerClient implements LoadBalancerClient { @Override public ServiceInstance choose(String serviceId) { - ribbonClientPreprocessor.preprocess(serviceId); - ILoadBalancer loadBalancer = this.balancers.get(serviceId); - if (loadBalancer == null) { - loadBalancer = clientFactory.getNamedLoadBalancer(serviceId); - } - Server server = loadBalancer.chooseServer("default"); - if (server == null) { - throw new IllegalStateException( - "Unable to locate ILoadBalancer for service: " + serviceId); - } - return new RibbonServer(serviceId, server); + return new RibbonServer(serviceId, getServer(serviceId)); } - private class RibbonServer implements ServiceInstance { + @Override + public T choose(String serviceId, LoadBalancerRequest request) { + ILoadBalancer loadBalancer = getLoadBalancer(serviceId); + Server server = getServer(serviceId, loadBalancer); + RibbonServer ribbonServer = new RibbonServer(serviceId, server); + try { + return request.apply(ribbonServer); + } finally { + if (loadBalancer instanceof AbstractLoadBalancer) { + AbstractLoadBalancer.class.cast(loadBalancer).getLoadBalancerStats().incrementNumRequests(server); + } + } + } + + protected Server getServer(String serviceId) { + return getServer(serviceId, getLoadBalancer(serviceId)); + } + protected Server getServer(String serviceId, ILoadBalancer loadBalancer) { + Server server = loadBalancer.chooseServer("default"); + if (server == null) { + throw new IllegalStateException( + "Unable to locate ILoadBalancer for service: " + serviceId); + } + return server; + } + + protected ILoadBalancer getLoadBalancer(String serviceId) { + ribbonClientPreprocessor.preprocess(serviceId); + ILoadBalancer loadBalancer = this.balancers.get(serviceId); + if (loadBalancer == null) { + loadBalancer = clientFactory.getNamedLoadBalancer(serviceId); + } + return loadBalancer; + } + + private class RibbonServer implements ServiceInstance { private String serviceId; private Server server;