Added classes for async load balancer interceptor

spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/AsyncLoadBalancerAutoConfiguration.java

@@ -0,0 +1,86 @@
+/*
+ * Copyright 2013-2016 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.client.loadbalancer;
+
+import org.springframework.beans.factory.SmartInitializingSingleton;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
+import org.springframework.boot.context.properties.EnableConfigurationProperties;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.http.client.AsyncClientHttpRequestInterceptor;
+import org.springframework.web.client.AsyncRestTemplate;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * Auto configuration for Ribbon (client side load balancing).
+ *
+ * @author Rob Worsnop
+ */
+@Configuration
+@ConditionalOnClass(AsyncRestTemplate.class)
+@ConditionalOnBean(LoadBalancerClient.class)
+@EnableConfigurationProperties(LoadBalancerRetryProperties.class)
+public class AsyncLoadBalancerAutoConfiguration {
+
+	@LoadBalanced
+	@Autowired(required = false)
+	private List<AsyncRestTemplate> restTemplates = Collections.emptyList();
+
+	@Bean
+	public SmartInitializingSingleton loadBalancedRestTemplateInitializer(
+			final List<AsyncRestTemplateCustomizer> customizers) {
+		return new SmartInitializingSingleton() {
+			@Override
+			public void afterSingletonsInstantiated() {
+				for (AsyncRestTemplate restTemplate : AsyncLoadBalancerAutoConfiguration.this.restTemplates) {
+					for (AsyncRestTemplateCustomizer customizer : customizers) {
+						customizer.customize(restTemplate);
+					}
+				}
+			}
+		};
+	}
+
+	@Configuration
+	static class LoadBalancerInterceptorConfig {
+		@Bean
+		public AsyncLoadBalancerInterceptor ribbonInterceptor(LoadBalancerClient loadBalancerClient) {
+			return new AsyncLoadBalancerInterceptor(loadBalancerClient);
+		}
+
+		@Bean
+		@ConditionalOnMissingBean
+		public AsyncRestTemplateCustomizer restTemplateCustomizer(
+				final AsyncLoadBalancerInterceptor loadBalancerInterceptor) {
+			return new AsyncRestTemplateCustomizer() {
+				@Override
+				public void customize(AsyncRestTemplate restTemplate) {
+					List<AsyncClientHttpRequestInterceptor> list = new ArrayList<>(
+							restTemplate.getInterceptors());
+					list.add(loadBalancerInterceptor);
+					restTemplate.setInterceptors(list);
+				}
+			};
+		}
+	}
+}

spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/AsyncLoadBalancerInterceptor.java

@@ -0,0 +1,58 @@
+/*
+ * Copyright 2013-2016 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.client.loadbalancer;
+
+import org.springframework.cloud.client.ServiceInstance;
+import org.springframework.http.HttpRequest;
+import org.springframework.http.client.AsyncClientHttpRequestExecution;
+import org.springframework.http.client.AsyncClientHttpRequestInterceptor;
+import org.springframework.http.client.ClientHttpResponse;
+import org.springframework.util.concurrent.ListenableFuture;
+
+import java.io.IOException;
+import java.net.URI;
+
+/**
+ * @author Rob Worsnop
+ */
+public class AsyncLoadBalancerInterceptor implements AsyncClientHttpRequestInterceptor {
+
+	private LoadBalancerClient loadBalancer;
+
+	public AsyncLoadBalancerInterceptor(LoadBalancerClient loadBalancer) {
+		this.loadBalancer = loadBalancer;
+	}
+
+	@Override
+	public ListenableFuture<ClientHttpResponse> intercept(final HttpRequest request, final byte[] body,
+			final AsyncClientHttpRequestExecution execution) throws IOException {
+		final URI originalUri = request.getURI();
+		String serviceName = originalUri.getHost();
+		return this.loadBalancer.execute(serviceName,
+				new LoadBalancerRequest<ListenableFuture<ClientHttpResponse>>() {
+					@Override
+					public ListenableFuture<ClientHttpResponse> apply(final ServiceInstance instance)
+							throws Exception {
+						HttpRequest serviceRequest = new ServiceRequestWrapper(request,
+								instance, loadBalancer);
+						return execution.executeAsync(serviceRequest, body);
+					}
+
+				});
+	}
+}
+

spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/AsyncRestTemplateCustomizer.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2013-2016 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.client.loadbalancer;
+
+import org.springframework.web.client.AsyncRestTemplate;
+
+/**
+ * @author Rob Worsnop
+ */
+public interface AsyncRestTemplateCustomizer {
+	void customize(AsyncRestTemplate restTemplate);
+}
+

spring-cloud-commons/src/main/resources/META-INF/spring.factories

@@ -3,6 +3,7 @@ org.springframework.boot.autoconfigure.EnableAutoConfiguration=\
 org.springframework.cloud.client.CommonsClientAutoConfiguration,\
 org.springframework.cloud.client.discovery.noop.NoopDiscoveryClientAutoConfiguration,\
 org.springframework.cloud.client.hypermedia.CloudHypermediaAutoConfiguration,\
+org.springframework.cloud.client.loadbalancer.AsyncLoadBalancerAutoConfiguration,\
 org.springframework.cloud.client.loadbalancer.LoadBalancerAutoConfiguration,\
 org.springframework.cloud.client.serviceregistry.ServiceRegistryAutoConfiguration,\
 org.springframework.cloud.commons.util.UtilAutoConfiguration

spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/AsyncLoadBalancerAutoConfigurationTests.java

@@ -0,0 +1,159 @@
+package org.springframework.cloud.client.loadbalancer;
+
+import lombok.SneakyThrows;
+import org.hamcrest.MatcherAssert;
+import org.junit.Test;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.builder.SpringApplicationBuilder;
+import org.springframework.cloud.client.DefaultServiceInstance;
+import org.springframework.cloud.client.ServiceInstance;
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.context.annotation.Primary;
+import org.springframework.http.client.AsyncClientHttpRequestInterceptor;
+import org.springframework.web.client.AsyncRestTemplate;
+
+import java.io.IOException;
+import java.net.URI;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+
+import static org.hamcrest.Matchers.notNullValue;
+import static org.hamcrest.Matchers.hasSize;
+import static org.hamcrest.Matchers.instanceOf;
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.empty;
+
+
+/**
+ * @author Rob Worsnop
+ */
+public class AsyncLoadBalancerAutoConfigurationTests {
+
+	@Test
+	public void restTemplateGetsLoadBalancerInterceptor() {
+		ConfigurableApplicationContext context = init(OneRestTemplate.class);
+		final Map<String, AsyncRestTemplate> restTemplates = context
+				.getBeansOfType(AsyncRestTemplate.class);
+
+		MatcherAssert.assertThat(restTemplates, is(notNullValue()));
+		MatcherAssert.assertThat(restTemplates.values(), hasSize(1));
+		AsyncRestTemplate restTemplate = restTemplates.values().iterator().next();
+		MatcherAssert.assertThat(restTemplate, is(notNullValue()));
+
+		assertLoadBalanced(restTemplate);
+	}
+
+	private void assertLoadBalanced(AsyncRestTemplate restTemplate) {
+		List<AsyncClientHttpRequestInterceptor> interceptors = restTemplate.getInterceptors();
+		MatcherAssert.assertThat(interceptors, hasSize(1));
+		AsyncClientHttpRequestInterceptor interceptor = interceptors.get(0);
+		MatcherAssert.assertThat(interceptor, is(instanceOf(AsyncLoadBalancerInterceptor.class)));
+	}
+
+	@Test
+	public void multipleRestTemplates() {
+		ConfigurableApplicationContext context = init(TwoRestTemplates.class);
+		final Map<String, AsyncRestTemplate> restTemplates = context
+				.getBeansOfType(AsyncRestTemplate.class);
+
+		MatcherAssert.assertThat(restTemplates, is(notNullValue()));
+		Collection<AsyncRestTemplate> templates = restTemplates.values();
+		MatcherAssert.assertThat(templates, hasSize(2));
+
+		TwoRestTemplates.Two two = context.getBean(TwoRestTemplates.Two.class);
+
+		MatcherAssert.assertThat(two.loadBalanced, is(notNullValue()));
+		assertLoadBalanced(two.loadBalanced);
+
+		MatcherAssert.assertThat(two.nonLoadBalanced, is(notNullValue()));
+		MatcherAssert.assertThat(two.nonLoadBalanced.getInterceptors(), is(empty()));
+	}
+
+	protected ConfigurableApplicationContext init(Class<?> config) {
+		return new SpringApplicationBuilder().web(false)
+				.properties("spring.aop.proxyTargetClass=true")
+				.sources(config, AsyncLoadBalancerAutoConfiguration.class).run();
+	}
+
+	@Configuration
+	protected static class OneRestTemplate {
+
+		@LoadBalanced
+		@Bean
+		AsyncRestTemplate loadBalancedRestTemplate() {
+			return new AsyncRestTemplate();
+		}
+
+		@Bean
+		LoadBalancerClient loadBalancerClient() {
+			return new NoopLoadBalancerClient();
+		}
+
+		@Bean
+		LoadBalancedRetryPolicyFactory loadBalancedRetryPolicyFactory() { return new LoadBalancedRetryPolicyFactory.NeverRetryFactory();}
+
+	}
+
+	@Configuration
+	protected static class TwoRestTemplates {
+
+		@Primary
+		@Bean
+		AsyncRestTemplate restTemplate() {
+			return new AsyncRestTemplate();
+		}
+
+		@LoadBalanced
+		@Bean
+		AsyncRestTemplate loadBalancedRestTemplate() {
+			return new AsyncRestTemplate();
+		}
+
+		@Bean
+		LoadBalancerClient loadBalancerClient() {
+			return new NoopLoadBalancerClient();
+		}
+
+		@Configuration
+		protected static class Two {
+			@Autowired
+			AsyncRestTemplate nonLoadBalanced;
+
+			@Autowired
+			@LoadBalanced
+			AsyncRestTemplate loadBalanced;
+		}
+
+	}
+
+	private static class NoopLoadBalancerClient implements LoadBalancerClient {
+		private final Random random = new Random();
+
+		@Override
+		public ServiceInstance choose(String serviceId) {
+			return new DefaultServiceInstance(serviceId, serviceId,
+					this.random.nextInt(40000), false);
+		}
+
+		@Override
+		@SneakyThrows
+		public <T> T execute(String serviceId, LoadBalancerRequest<T> request) {
+			return request.apply(choose(serviceId));
+		}
+
+		@Override
+		@SneakyThrows
+		public <T> T execute(String serviceId, ServiceInstance serviceInstance, LoadBalancerRequest<T> request) throws IOException {
+			return request.apply(choose(serviceId));
+		}
+
+		@Override
+		public URI reconstructURI(ServiceInstance instance, URI original) {
+			return DefaultServiceInstance.getUri(instance);
+		}
+	}
+}