Add support loadBalancerKey in RibbonRoutingFilter of Zuul for Canary test (#2439)

* Add support loadBalancerKey in RibbonRoutingFilter

* Add documentation about how to provide a loadBalancerKey to Ribbon

* Add integrationtest for RibbonRoutingFilter's loadBalancerKey support

docs/src/main/asciidoc/spring-cloud-netflix.adoc

@@ -956,7 +956,33 @@ zuul:
     threadPoolKeyPrefix: zuulgw
 ----
 
+[[how-to-provdie-a-key-to-ribbon]]
+=== How to Provide a Key to Ribbon's `IRule`
 
+If you need to provide your own `IRule` implementation to handle a special routing requirement like a canary test,
+you probably want to pass some information to the `choose` method of `IRule`.
+
+.com.netflix.loadbalancer.IRule.java
+----
+public interface IRule{
+    public Server choose(Object key);
+         :
+----
+
+You can provide some information that will be used to choose a target server by your `IRule` implementation like
+the following:
+
+----
+RequestContext.getCurrentContext()
+              .set(FilterConstants.LOAD_BALANCER_KEY, "canary-test");
+----
+
+If you put any object into the `RequestContext` with a key `FilterConstants.LOAD_BALANCER_KEY`, it will
+be passed to the `choose` method of `IRule` implementation. Above code must be executed before `RibbonRoutingFilter`
+is executed and Zuul's pre filter is the best place to do that. You can easily access HTTP headers and query parameters
+via `RequestContext` in pre filter, so it can be used to determine `LOAD_BALANCER_KEY` that will be passed to Ribbon.
+If you don't put any value with `LOAD_BALANCER_KEY` in `RequestContext`, null will be passed as a parameter of `choose`
+method.
 
 [[spring-cloud-feign]]
 == Declarative REST Client: Feign

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

@@ -17,6 +17,7 @@
 
 package org.springframework.cloud.netflix.ribbon.support;
 
+import com.netflix.loadbalancer.reactive.LoadBalancerCommand;
 import org.springframework.cloud.netflix.ribbon.DefaultServerIntrospector;
 import org.springframework.cloud.netflix.ribbon.RibbonClientConfiguration;
 import org.springframework.cloud.netflix.ribbon.ServerIntrospector;
@@ -140,4 +141,11 @@ public abstract class AbstractLoadBalancingClient<S extends ContextAwareRequest,
 		}
 		return this.secure;
 	}
+
+	@Override
+	protected void customizeLoadBalancerCommandBuilder(S request, IClientConfig config, LoadBalancerCommand.Builder<T> builder) {
+		if (request.getLoadBalancerKey() != null) {
+			builder.withServerLocator(request.getLoadBalancerKey());
+		}
+	}
 }

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

@@ -43,6 +43,7 @@ public abstract class ContextAwareRequest extends ClientRequest implements HttpR
 		}
 		this.uri = context.uri();
 		this.isRetriable = context.getRetryable();
+		this.loadBalancerKey = context.getLoadBalancerKey();
 	}
 
 	public RibbonCommandContext getContext() {
@@ -68,7 +69,8 @@ public abstract class ContextAwareRequest extends ClientRequest implements HttpR
 		RibbonCommandContext commandContext = new RibbonCommandContext(this.context.getServiceId(),
 				this.context.getMethod(), uri.toString(), this.context.getRetryable(),
 				this.context.getHeaders(), this.context.getParams(), this.context.getRequestEntity(),
-				this.context.getRequestCustomizers(), this.context.getContentLength());
+				this.context.getRequestCustomizers(), this.context.getContentLength(),
+				this.context.getLoadBalancerKey());
 		return commandContext;
 	}
 }

spring-cloud-netflix-core/src/main/java/org/springframework/cloud/netflix/zuul/filters/route/RibbonCommandContext.java

@@ -32,6 +32,7 @@ import java.util.Objects;
 
 /**
  * @author Spencer Gibb
+ * @author Yongsung Yoon
  */
 public class RibbonCommandContext {
 	private final String serviceId;
@@ -43,6 +44,7 @@ public class RibbonCommandContext {
 	private final List<RibbonRequestCustomizer> requestCustomizers;
 	private InputStream requestEntity;
 	private Long contentLength;
+	private Object loadBalancerKey;
 
 	/**
 	 * Kept for backwards compatibility with Spring Cloud Sleuth 1.x versions
@@ -52,7 +54,7 @@ public class RibbonCommandContext {
 								String uri, Boolean retryable, MultiValueMap<String, String> headers,
 								MultiValueMap<String, String> params, InputStream requestEntity) {
 		this(serviceId, method, uri, retryable, headers, params, requestEntity,
-			new ArrayList<RibbonRequestCustomizer>(), null);
+			new ArrayList<RibbonRequestCustomizer>(), null, null);
 	}
 
 	public RibbonCommandContext(String serviceId, String method, String uri,
@@ -60,13 +62,22 @@ public class RibbonCommandContext {
 								MultiValueMap<String, String> params, InputStream requestEntity,
 								List<RibbonRequestCustomizer> requestCustomizers) {
 		this(serviceId, method, uri, retryable, headers, params, requestEntity,
-			requestCustomizers, null);
+			requestCustomizers, null, null);
 	}
 
 	public RibbonCommandContext(String serviceId, String method, String uri,
 								Boolean retryable, MultiValueMap<String, String> headers,
 								MultiValueMap<String, String> params, InputStream requestEntity,
 								List<RibbonRequestCustomizer> requestCustomizers, Long contentLength) {
+		this(serviceId, method, uri, retryable, headers, params, requestEntity,
+			requestCustomizers, contentLength, null);
+	}
+
+	public RibbonCommandContext(String serviceId, String method, String uri,
+								Boolean retryable, MultiValueMap<String, String> headers,
+								MultiValueMap<String, String> params, InputStream requestEntity,
+								List<RibbonRequestCustomizer> requestCustomizers, Long contentLength,
+								Object loadBalancerKey) {
 		Assert.notNull(serviceId, "serviceId may not be null");
 		Assert.notNull(method, "method may not be null");
 		Assert.notNull(uri, "uri may not be null");
@@ -82,6 +93,7 @@ public class RibbonCommandContext {
 		this.requestEntity = requestEntity;
 		this.requestCustomizers = requestCustomizers;
 		this.contentLength = contentLength;
+		this.loadBalancerKey = loadBalancerKey;
 	}
 
 	public URI uri() {
@@ -155,6 +167,14 @@ public class RibbonCommandContext {
 		this.contentLength = contentLength;
 	}
 
+	public Object getLoadBalancerKey() {
+		return loadBalancerKey;
+	}
+
+	public void setLoadBalancerKey(Object loadBalancerKey) {
+		this.loadBalancerKey = loadBalancerKey;
+	}
+
 	@Override
 	public boolean equals(Object o) {
 		if (this == o)
@@ -169,13 +189,14 @@ public class RibbonCommandContext {
 			.equals(params, that.params) && Objects
 			.equals(requestEntity, that.requestEntity) && Objects
 			.equals(requestCustomizers, that.requestCustomizers) && Objects
-			.equals(contentLength, that.contentLength);
+			.equals(contentLength, that.contentLength) && Objects
+			.equals(loadBalancerKey, that.loadBalancerKey);
 	}
 
 	@Override
 	public int hashCode() {
 		return Objects.hash(serviceId, method, uri, retryable, headers, params,
-			requestEntity, requestCustomizers, contentLength);
+			requestEntity, requestCustomizers, contentLength, loadBalancerKey);
 	}
 
 	@Override
@@ -190,6 +211,7 @@ public class RibbonCommandContext {
 		sb.append(", requestEntity=").append(requestEntity);
 		sb.append(", requestCustomizers=").append(requestCustomizers);
 		sb.append(", contentLength=").append(contentLength);
+		sb.append(", loadBalancerKey=").append(loadBalancerKey);
 		sb.append('}');
 		return sb.toString();
 	}

spring-cloud-netflix-core/src/main/java/org/springframework/cloud/netflix/zuul/filters/route/RibbonRoutingFilter.java

@@ -22,7 +22,6 @@ import java.util.List;
 import java.util.Map;
 
 import javax.servlet.http.HttpServletRequest;
-import javax.servlet.http.HttpServletResponse;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -44,6 +43,7 @@ import static org.springframework.cloud.netflix.zuul.filters.support.FilterConst
 import static org.springframework.cloud.netflix.zuul.filters.support.FilterConstants.RIBBON_ROUTING_FILTER_ORDER;
 import static org.springframework.cloud.netflix.zuul.filters.support.FilterConstants.ROUTE_TYPE;
 import static org.springframework.cloud.netflix.zuul.filters.support.FilterConstants.SERVICE_ID_KEY;
+import static org.springframework.cloud.netflix.zuul.filters.support.FilterConstants.LOAD_BALANCER_KEY;
 
 /**
  * Route {@link ZuulFilter} that uses Ribbon, Hystrix and pluggable http clients to send requests.
@@ -135,6 +135,7 @@ public class RibbonRoutingFilter extends ZuulFilter {
 
 		String serviceId = (String) context.get(SERVICE_ID_KEY);
 		Boolean retryable = (Boolean) context.get(RETRYABLE_KEY);
+		Object loadBalancerKey = context.get(LOAD_BALANCER_KEY);
 
 		String uri = this.helper.buildZuulRequestURI(request);
 
@@ -144,7 +145,7 @@ public class RibbonRoutingFilter extends ZuulFilter {
 		long contentLength = useServlet31 ? request.getContentLengthLong(): request.getContentLength();
 
 		return new RibbonCommandContext(serviceId, verb, uri, retryable, headers, params,
-				requestEntity, this.requestCustomizers, contentLength);
+				requestEntity, this.requestCustomizers, contentLength, loadBalancerKey);
 	}
 
 	protected ClientHttpResponse forward(RibbonCommandContext context) throws Exception {

spring-cloud-netflix-core/src/main/java/org/springframework/cloud/netflix/zuul/filters/support/FilterConstants.java

@@ -70,6 +70,11 @@ public class FilterConstants {
 	 */
 	public static final String SERVICE_ID_KEY = "serviceId";
 
+	/**
+	 * Zuul {@link com.netflix.zuul.context.RequestContext} key for use in {@link org.springframework.cloud.netflix.zuul.filters.route.RibbonRoutingFilter}
+	 */
+	public static final String LOAD_BALANCER_KEY = "loadBalancerKey";
+
 	// ORDER constants -----------------------------------
 
 	/**

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

@@ -56,6 +56,7 @@ public class ContextAwareRequestTests {
 		doReturn(new URI("http://foo")).when(context).uri();
 		doReturn("foo").when(context).getServiceId();
 		doReturn(new LinkedMultiValueMap<>()).when(context).getParams();
+		doReturn("testLoadBalancerKey").when(context).getLoadBalancerKey();
 		request = new TestContextAwareRequest(context);
 	}
 
@@ -97,6 +98,18 @@ public class ContextAwareRequestTests {
 		assertEquals(headers, request.getHeaders());
 	}
 
+	@Test
+	public void getLoadBalancerKey() throws Exception {
+		assertEquals("testLoadBalancerKey", request.getLoadBalancerKey());
+
+		RibbonCommandContext defaultContext = mock(RibbonCommandContext.class);
+		doReturn(new LinkedMultiValueMap()).when(defaultContext).getHeaders();
+		doReturn(null).when(defaultContext).getLoadBalancerKey();
+		ContextAwareRequest defaultRequest = new TestContextAwareRequest(defaultContext);
+
+		assertNull(defaultRequest.getLoadBalancerKey());
+	}
+
 	static class TestContextAwareRequest extends ContextAwareRequest {
 
 		public TestContextAwareRequest(RibbonCommandContext context) {

spring-cloud-netflix-core/src/test/java/org/springframework/cloud/netflix/zuul/filters/route/RibbonCommandContextTest.java

@@ -27,9 +27,11 @@ import org.springframework.util.LinkedMultiValueMap;
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.util.Collections;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 
 /**
@@ -91,4 +93,18 @@ public class RibbonCommandContextTest {
 				new ByteArrayInputStream(TEST_CONTENT),
 				Lists.newArrayList(requestCustomizer));
 	}
+
+	@Test
+	public void testNullSafetyWithNullableParameters() throws Exception {
+		LinkedMultiValueMap headers = new LinkedMultiValueMap();
+		LinkedMultiValueMap params = new LinkedMultiValueMap();
+
+		RibbonCommandContext testContext = new RibbonCommandContext("serviceId",
+				HttpMethod.POST.toString(), "/my/route", true, headers, params,
+				new ByteArrayInputStream(TEST_CONTENT), Collections.<RibbonRequestCustomizer>emptyList(),
+				null, null);
+
+		assertNotEquals(0, testContext.hashCode());
+		assertNotNull(testContext.toString());
+	}
 }

spring-cloud-netflix-core/src/test/java/org/springframework/cloud/netflix/zuul/filters/route/RibbonRoutingFilterLoadBalancerKeyIntegrationTests.java

@@ -0,0 +1,178 @@
+/*
+ * Copyright 2013-2017 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.netflix.zuul.filters.route;
+
+import com.netflix.loadbalancer.AvailabilityFilteringRule;
+import com.netflix.loadbalancer.IRule;
+import com.netflix.loadbalancer.Server;
+import com.netflix.loadbalancer.ServerList;
+import com.netflix.zuul.ZuulFilter;
+import com.netflix.zuul.context.RequestContext;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.boot.context.embedded.LocalServerPort;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.boot.test.web.client.TestRestTemplate;
+import org.springframework.cloud.netflix.ribbon.RibbonClient;
+import org.springframework.cloud.netflix.ribbon.StaticServerList;
+import org.springframework.cloud.netflix.zuul.EnableZuulProxy;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.http.*;
+import org.springframework.test.annotation.DirtiesContext;
+import org.springframework.test.context.junit4.SpringRunner;
+import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RequestMethod;
+import org.springframework.web.bind.annotation.RestController;
+
+import javax.servlet.http.HttpServletRequest;
+
+import static org.junit.Assert.assertEquals;
+import static org.springframework.cloud.netflix.zuul.filters.support.FilterConstants.LOAD_BALANCER_KEY;
+import static org.springframework.cloud.netflix.zuul.filters.support.FilterConstants.PRE_TYPE;
+
+/**
+ * @author Yongsung Yoon
+ */
+@RunWith(SpringRunner.class)
+@SpringBootTest(classes = CanaryTestZuulProxyApplication.class,
+				webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT,
+				value = { "zuul.routes.simple.path: /simple/**" })
+@DirtiesContext
+public class RibbonRoutingFilterLoadBalancerKeyIntegrationTests {
+
+	@Autowired
+	private TestRestTemplate testRestTemplate;
+
+	@Before
+	public void setTestRequestContext() {
+		RequestContext context = new RequestContext();
+		RequestContext.testSetCurrentContext(context);
+	}
+
+	@After
+	public void clear() {
+		RequestContext.getCurrentContext().clear();
+	}
+
+	@Test
+	public void invokeWithUserDefinedCanaryHeader() {
+		HttpHeaders headers = new HttpHeaders();
+		headers.set("X-Canary-Test", "true");
+
+		ResponseEntity<String> result = testRestTemplate.exchange("/simple/hello", HttpMethod.GET,
+				new HttpEntity<>(headers), String.class);
+		assertEquals(HttpStatus.OK, result.getStatusCode());
+		assertEquals("canary", result.getBody());
+	}
+
+	@Test
+	public void invokeWithoutUserDefinedCanaryHeader() {
+		HttpHeaders headers = new HttpHeaders();
+		ResponseEntity<String> result = testRestTemplate.exchange("/simple/hello", HttpMethod.GET,
+				new HttpEntity<>(headers), String.class);
+		assertEquals(HttpStatus.INTERNAL_SERVER_ERROR, result.getStatusCode());
+	}
+}
+
+@Configuration
+@EnableAutoConfiguration
+@RestController
+@EnableZuulProxy
+@RibbonClient(name = "simple", configuration = CanaryTestRibbonClientConfiguration.class)
+class CanaryTestZuulProxyApplication {
+
+	@RequestMapping(value = "/hello", method = RequestMethod.GET)
+	public String hello() {
+		return "canary";
+	}
+
+	@Bean
+	public ZuulFilter testCanarySupportPreFilter() {
+		return new ZuulFilter() {
+			@Override
+			public Object run() {
+				RequestContext context = RequestContext.getCurrentContext();
+				if (checkIfCanaryRequest(context)) {
+					context.set(LOAD_BALANCER_KEY, "canary"); // set loadBalancerKey for IRule
+				}
+				return null;
+			}
+
+			private boolean checkIfCanaryRequest(RequestContext context) {
+				HttpServletRequest request = context.getRequest();
+				String canaryHeader = request.getHeader("X-Canary-Test"); // user defined header
+
+				if ((canaryHeader != null) && (canaryHeader.equalsIgnoreCase("true"))) {
+					return true;
+				}
+				return false;
+			}
+
+			@Override
+			public boolean shouldFilter() {
+				return true;
+			}
+
+			@Override
+			public String filterType() {
+				return PRE_TYPE;
+			}
+
+			@Override
+			public int filterOrder() {
+				return 0;
+			}
+		};
+	}
+}
+
+@Configuration
+class CanaryTestRibbonClientConfiguration {
+
+	@LocalServerPort
+	private  int port;
+
+	private static Server testCanaryInstance;
+
+	@Bean
+	public ServerList<Server> ribbonServerList() {
+		return new StaticServerList<>(new Server("normal-routing-notexist-localhost", this.port));
+	}
+
+	@Bean
+	public IRule canaryTestRule() {
+		if (testCanaryInstance == null) {
+			testCanaryInstance = new Server("localhost", port); // use test server as a canary instance
+		}
+		return new TestCanaryRule();
+	}
+
+	public static class TestCanaryRule extends AvailabilityFilteringRule {
+		@Override
+		public Server choose(Object key) {
+			if ((key != null) && (key.equals("canary"))) {
+				return testCanaryInstance;  // choose test canary server instead of normal servers.
+			}
+			return super.choose(key); // normal routing
+		}
+	}
+}

spring-cloud-netflix-core/src/test/java/org/springframework/cloud/netflix/zuul/filters/route/RibbonRoutingFilterTests.java

@@ -19,22 +19,73 @@ package org.springframework.cloud.netflix.zuul.filters.route;
 
 import java.util.Collections;
 
+import com.netflix.zuul.context.RequestContext;
+import org.junit.After;
+import org.junit.Before;
 import org.junit.Test;
 import org.springframework.cloud.netflix.ribbon.support.RibbonRequestCustomizer;
 import org.springframework.cloud.netflix.zuul.filters.ProxyRequestHelper;
+import org.springframework.mock.web.MockHttpServletRequest;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.mock;
+import static org.springframework.cloud.netflix.zuul.filters.support.FilterConstants.LOAD_BALANCER_KEY;
+import static org.springframework.cloud.netflix.zuul.filters.support.FilterConstants.SERVICE_ID_KEY;
 
 /**
  * @author Spencer Gibb
+ * @author Yongsung Yoon
  */
 public class RibbonRoutingFilterTests {
+
+	private RequestContext requestContext;
+	private RibbonRoutingFilter filter;
+
+	@Before
+	public void setUp() throws Exception {
+		setUpRequestContext();
+		setupRibbonRoutingFilter();
+	}
+
+	@After
+	public void tearDown() throws Exception {
+		requestContext.unset();
+	}
+
 	@Test
 	public void useServlet31Works() {
-		RibbonCommandFactory factory = mock(RibbonCommandFactory.class);
-		RibbonRoutingFilter filter = new RibbonRoutingFilter(new ProxyRequestHelper(), factory,
-				Collections.<RibbonRequestCustomizer>emptyList());
 		assertThat(filter.isUseServlet31()).isTrue();
 	}
+
+	@Test
+	public void testLoadBalancerKeyToRibbonCommandContext() throws Exception {
+		final String testKey = "testLoadBalancerKey";
+		requestContext.set(LOAD_BALANCER_KEY, testKey);
+		RibbonCommandContext commandContext = filter.buildCommandContext(requestContext);
+
+		assertThat(commandContext.getLoadBalancerKey()).isEqualTo(testKey);
+	}
+
+	@Test
+	public void testNullLoadBalancerKeyToRibbonCommandContext() throws Exception {
+		requestContext.set(LOAD_BALANCER_KEY, null);
+		RibbonCommandContext commandContext = filter.buildCommandContext(requestContext);
+
+		assertThat(commandContext.getLoadBalancerKey()).isNull();
+	}
+
+	private void setUpRequestContext() {
+		requestContext = RequestContext.getCurrentContext();
+		MockHttpServletRequest mockRequest = new MockHttpServletRequest();
+		mockRequest.setMethod("GET");
+		mockRequest.setRequestURI("/foo/bar");
+		requestContext.setRequest(mockRequest);
+		requestContext.setRequestQueryParams(Collections.EMPTY_MAP);
+		requestContext.set(SERVICE_ID_KEY, "testServiceId");
+	}
+
+	private void setupRibbonRoutingFilter() {
+		RibbonCommandFactory factory = mock(RibbonCommandFactory.class);
+		filter = new RibbonRoutingFilter(new ProxyRequestHelper(), factory, Collections.<RibbonRequestCustomizer>emptyList());
+	}
 }

spring-cloud-netflix-dependencies/pom.xml

@@ -18,7 +18,7 @@
 		<eureka.version>1.7.0</eureka.version>
 		<feign.version>9.5.0</feign.version>
 		<hystrix.version>1.5.12</hystrix.version>
-		<ribbon.version>2.2.2</ribbon.version>
+		<ribbon.version>2.2.4</ribbon.version>
 		<servo.version>0.10.1</servo.version>
 		<zuul.version>1.3.0</zuul.version>
 		<rxjava.version>1.2.0</rxjava.version>