updated to use RestTemplateCustomizer

spring-cloud-netflix-core/src/main/java/org/springframework/cloud/netflix/ribbon/RibbonAutoConfiguration.java

@@ -19,23 +19,21 @@ package org.springframework.cloud.netflix.ribbon;
 import java.util.ArrayList;
 import java.util.List;
 
-import com.netflix.client.http.HttpRequest;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.AutoConfigureBefore;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
-import org.springframework.cloud.client.loadbalancer.LoadBalanced;
 import org.springframework.cloud.client.loadbalancer.LoadBalancerAutoConfiguration;
 import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
+import org.springframework.cloud.client.loadbalancer.RestTemplateCustomizer;
 import org.springframework.cloud.netflix.eureka.EurekaClientAutoConfiguration;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.web.client.RestTemplate;
 
 import com.netflix.client.IClient;
-
-import javax.annotation.PostConstruct;
+import com.netflix.client.http.HttpRequest;
 
 /**
  * Auto configuration for Ribbon (client side load balancing).
@@ -70,23 +68,26 @@ public class RibbonAutoConfiguration {
     @Configuration
     @ConditionalOnClass(HttpRequest.class)
     protected static class RibbonClientConfig {
-        @Autowired(required = false)
-        @LoadBalanced
-        private RestTemplate restTemplate;
 
         @Autowired
         private SpringClientFactory springClientFactory;
 
-        @PostConstruct
-        public void init() {
-            if (restTemplate != null) {
-                restTemplate.setRequestFactory(ribbonClientHttpRequestFactory());
-            }
+        @Autowired
+        private LoadBalancerClient loadBalancerClient;
+
+        @Bean
+        public RestTemplateCustomizer restTemplateCustomizer() {
+            return new RestTemplateCustomizer() {
+                @Override
+                public void customize(RestTemplate restTemplate) {
+                    restTemplate.setRequestFactory(ribbonClientHttpRequestFactory());
+                }
+            };
         }
 
         @Bean
         public RibbonClientHttpRequestFactory ribbonClientHttpRequestFactory() {
-            return new RibbonClientHttpRequestFactory(springClientFactory);
+            return new RibbonClientHttpRequestFactory(springClientFactory, loadBalancerClient);
         }
     }
 

spring-cloud-netflix-core/src/main/java/org/springframework/cloud/netflix/ribbon/RibbonClientHttpRequestFactory.java

@@ -24,7 +24,8 @@ import java.util.List;
 import java.util.Map;
 
 import org.springframework.cloud.client.ServiceInstance;
-import org.springframework.cloud.client.loadbalancer.LoadBalancerInterceptor;
+import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
+import org.springframework.cloud.client.loadbalancer.LoadBalancerRequest;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.HttpStatus;
@@ -46,16 +47,21 @@ public class RibbonClientHttpRequestFactory implements ClientHttpRequestFactory
 
 	private final SpringClientFactory clientFactory;
 
-	public RibbonClientHttpRequestFactory(SpringClientFactory clientFactory) {
+    private LoadBalancerClient loadBalancer;
+
+	public RibbonClientHttpRequestFactory(SpringClientFactory clientFactory, LoadBalancerClient loadBalancer) {
 		this.clientFactory = clientFactory;
+        this.loadBalancer = loadBalancer;
 	}
 
 	@Override
 	@SuppressWarnings("deprecation")
-	public ClientHttpRequest createRequest(URI uri, HttpMethod httpMethod)
+	public ClientHttpRequest createRequest(URI originalUri, HttpMethod httpMethod)
 			throws IOException {
-		//@formatter:off
-		ServiceInstance instance = LoadBalancerInterceptor.getThreadLocalServiceInstance();
+        String serviceId = originalUri.getHost();
+        ServiceInstance instance = loadBalancer.choose(serviceId);
+        URI uri = loadBalancer.reconstructURI(instance, originalUri);
+        //@formatter:off
 		IClientConfig clientConfig = clientFactory.getClientConfig(instance.getServiceId());
 		RestClient client = clientFactory.getClient(instance.getServiceId(), RestClient.class);
 		HttpRequest request = HttpRequest.newBuilder()
@@ -99,14 +105,14 @@ public class RibbonClientHttpRequestFactory implements ClientHttpRequestFactory
 		@Override
 		protected ClientHttpResponse executeInternal(HttpHeaders headers)
 				throws IOException {
-			HttpResponse response;
-			try {
-				response = client.execute(request, config);
-				return new RibbonHttpResponse(response);
-			}
-			catch (Exception e) {
-				throw new IOException(e);
-			}
+            // use execute here so stats are collected
+            return loadBalancer.execute(this.config.getClientName(), new LoadBalancerRequest<ClientHttpResponse>() {
+                @Override
+                public ClientHttpResponse apply(ServiceInstance instance) throws Exception {
+                    HttpResponse response = client.execute(request, config);
+                    return new RibbonHttpResponse(response);
+                }
+            });
 		}
 	}
 

spring-cloud-netflix-core/src/test/java/org/springframework/cloud/netflix/ribbon/RibbonClientPreprocessorOverridesIntegrationTests.java

@@ -46,35 +46,40 @@ public class RibbonClientPreprocessorOverridesIntegrationTests {
 
 	@Test
 	public void ruleOverridesToRandom() throws Exception {
-		RandomRule.class.cast(getLoadBalancer().getRule());
+		RandomRule.class.cast(getLoadBalancer("foo").getRule());
+        RoundRobinRule.class.cast(getLoadBalancer("bar").getRule());
 	}
 
 	@Test
 	public void pingOverridesToDummy() throws Exception {
-		DummyPing.class.cast(getLoadBalancer().getPing());
+		DummyPing.class.cast(getLoadBalancer("foo").getPing());
+        PingConstant.class.cast(getLoadBalancer("bar").getPing());
 	}
 
 	@Test
 	public void serverListOverridesToMy() throws Exception {
-		MyServiceList.class.cast(getLoadBalancer().getServerListImpl());
+		FooServiceList.class.cast(getLoadBalancer("foo").getServerListImpl());
+        BarServiceList.class.cast(getLoadBalancer("bar").getServerListImpl());
 	}
 
 	@SuppressWarnings("unchecked")
-	private ZoneAwareLoadBalancer<Server> getLoadBalancer() {
-		return (ZoneAwareLoadBalancer<Server>) this.factory
-				.getLoadBalancer("foo");
+	private ZoneAwareLoadBalancer<Server> getLoadBalancer(String name) {
+		return (ZoneAwareLoadBalancer<Server>) this.factory.getLoadBalancer(name);
 	}
 
 	@Test
 	public void serverListFilterOverride() throws Exception {
-		ServerListFilter<Server> filter = getLoadBalancer().getFilter();
-		assertEquals("MyTestZone",
+		ServerListFilter<Server> filter = getLoadBalancer("foo").getFilter();
+		assertEquals("FooTestZone",
 				ZonePreferenceServerListFilter.class.cast(filter)
 						.getZone());
 	}
 
 	@Configuration
-	@RibbonClient(name = "foo", configuration = FooConfiguration.class)
+    @RibbonClients({
+        @RibbonClient(name = "foo", configuration = FooConfiguration.class),
+        @RibbonClient(name = "bar", configuration = BarConfiguration.class)
+    })
 	@Import({ PropertyPlaceholderAutoConfiguration.class,
 			ArchaiusAutoConfiguration.class, EurekaClientAutoConfiguration.class,
 			RibbonAutoConfiguration.class})
@@ -83,11 +88,6 @@ public class RibbonClientPreprocessorOverridesIntegrationTests {
 
 	@Configuration
 	public static class FooConfiguration {
-
-		public FooConfiguration() {
-			System.out.println("here");
-		}
-
 		@Bean
 		public IRule ribbonRule() {
 			return new RandomRule();
@@ -100,21 +100,52 @@ public class RibbonClientPreprocessorOverridesIntegrationTests {
 
 		@Bean
 		public ServerList<Server> ribbonServerList(IClientConfig config) {
-			return new MyServiceList(config);
+			return new FooServiceList(config);
 		}
 
 		@Bean
 		public ZonePreferenceServerListFilter serverListFilter() {
 			ZonePreferenceServerListFilter filter = new ZonePreferenceServerListFilter();
-			filter.setZone("MyTestZone");
+			filter.setZone("FooTestZone");
 			return filter;
 		}
 	}
 
-	public static class MyServiceList extends ConfigurationBasedServerList {
-		public MyServiceList(IClientConfig config) {
+	public static class FooServiceList extends ConfigurationBasedServerList {
+		public FooServiceList(IClientConfig config) {
 			super.initWithNiwsConfig(config);
 		}
 	}
 
+    @Configuration
+    public static class BarConfiguration {
+
+        @Bean
+        public IRule ribbonRule() {
+            return new RoundRobinRule();
+        }
+
+        @Bean
+        public IPing ribbonPing() {
+            return new PingConstant();
+        }
+
+        @Bean
+        public ServerList<Server> ribbonServerList(IClientConfig config) {
+            return new BarServiceList(config);
+        }
+
+        @Bean
+        public ZonePreferenceServerListFilter serverListFilter() {
+            ZonePreferenceServerListFilter filter = new ZonePreferenceServerListFilter();
+            filter.setZone("BarTestZone");
+            return filter;
+        }
+    }
+
+    public static class BarServiceList extends ConfigurationBasedServerList {
+        public BarServiceList(IClientConfig config) {
+            super.initWithNiwsConfig(config);
+        }
+    }
 }