From 64d7a57190f01f2115fd07abb09a63bbf8b0d80a Mon Sep 17 00:00:00 2001
From: Spencer Gibb <spencer@gibb.us>
Date: Thu, 26 Mar 2015 10:46:15 -0600
Subject: [PATCH] RestTemplateCustomizer to customize the @LoadBalanced
 RestTemplate.

By default, creates a RestTemplateCustomizer adds a LoadBalancerInterceptor

fixes gh-16
---
 .../LoadBalancerAutoConfiguration.java        | 20 +++++++++++++++----
 .../loadbalancer/LoadBalancerInterceptor.java |  7 -------
 .../loadbalancer/RestTemplateCustomizer.java  | 10 ++++++++++
 3 files changed, 26 insertions(+), 11 deletions(-)
 create mode 100644 spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/RestTemplateCustomizer.java

diff --git a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerAutoConfiguration.java b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerAutoConfiguration.java
index 48aa4ecb..09dab1e8 100644
--- a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerAutoConfiguration.java
+++ b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerAutoConfiguration.java
@@ -21,6 +21,7 @@ import java.util.List;
 
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.http.client.ClientHttpRequestInterceptor;
@@ -39,14 +40,25 @@ public class LoadBalancerAutoConfiguration {
 
 	@Bean
     @LoadBalanced
-	public RestTemplate loadBalancedRestTemplate(LoadBalancerInterceptor loadBalancerInterceptor) {
+	public RestTemplate loadBalancedRestTemplate(RestTemplateCustomizer customizer) {
 		RestTemplate restTemplate = new RestTemplate();
-		List<ClientHttpRequestInterceptor> list = new ArrayList<>();
-		list.add(loadBalancerInterceptor);
-		restTemplate.setInterceptors(list);
+        customizer.customize(restTemplate);
 		return restTemplate;
 	}
 
+    @Bean
+    @ConditionalOnMissingBean
+    public RestTemplateCustomizer restTemplateCustomizer(final LoadBalancerInterceptor loadBalancerInterceptor) {
+        return new RestTemplateCustomizer() {
+            @Override
+            public void customize(RestTemplate restTemplate) {
+                List<ClientHttpRequestInterceptor> list = new ArrayList<>();
+                list.add(loadBalancerInterceptor);
+                restTemplate.setInterceptors(list);
+            }
+        };
+    }
+
 	@Bean
 	public LoadBalancerInterceptor ribbonInterceptor(LoadBalancerClient loadBalancerClient) {
 		return new LoadBalancerInterceptor(loadBalancerClient);
diff --git a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerInterceptor.java b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerInterceptor.java
index bb5c0caa..5f1877b1 100644
--- a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerInterceptor.java
+++ b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/LoadBalancerInterceptor.java
@@ -34,8 +34,6 @@ public class LoadBalancerInterceptor implements ClientHttpRequestInterceptor {
 
 	private LoadBalancerClient loadBalancer;
 
-    private static final ThreadLocal<ServiceInstance> instanceHolder = new ThreadLocal<>();
-
 	public LoadBalancerInterceptor(LoadBalancerClient loadBalancer) {
 		this.loadBalancer = loadBalancer;
 	}
@@ -51,7 +49,6 @@ public class LoadBalancerInterceptor implements ClientHttpRequestInterceptor {
 					@Override
 					public ClientHttpResponse apply(final ServiceInstance instance)
 							throws Exception {
-                        instanceHolder.set(instance);
 						HttpRequest serviceRequest = new ServiceRequestWrapper(request,
 								instance);
 						return execution.execute(serviceRequest, body);
@@ -78,8 +75,4 @@ public class LoadBalancerInterceptor implements ClientHttpRequestInterceptor {
 
 	}
 
-    public static ServiceInstance getThreadLocalServiceInstance() {
-        return instanceHolder.get();
-    }
-
 }
diff --git a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/RestTemplateCustomizer.java b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/RestTemplateCustomizer.java
new file mode 100644
index 00000000..0ed9bfe1
--- /dev/null
+++ b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/RestTemplateCustomizer.java
@@ -0,0 +1,10 @@
+package org.springframework.cloud.client.loadbalancer;
+
+import org.springframework.web.client.RestTemplate;
+
+/**
+ * @author Spencer Gibb
+ */
+public interface RestTemplateCustomizer {
+    void customize(RestTemplate restTemplate);
+}