support get lbKey from request

ribbon-httpclient/src/test/java/com/netflix/niws/client/http/RestClientTest.java

@@ -20,6 +20,7 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
@@ -137,6 +138,33 @@ public void testDelete() throws Exception {
         assertStatusIsOk(response.getStatus());
     }
 
+    @Test
+    public void testExecuteWithSetLBKey() throws Exception {
+        final String srcKey = "66666";
+        RestClient client = (RestClient) ClientFactory.getNamedClient("allservices");
+        BaseLoadBalancer lb = new BaseLoadBalancer() {
+            public Server chooseServer(Object key) {
+                if (key != null) {
+                    assertEquals(srcKey, key.toString());
+                }
+                return super.chooseServer(key);
+            }
+        };
+        final Server[] servers = new Server[]{new Server("localhost", server.getServerPort())};
+        lb.addServers(Arrays.asList(servers));
+        client.setLoadBalancer(lb);
+        assertEquals(200, client.executeWithLoadBalancer(HttpRequest.newBuilder().uri(new URI("/")).build()).getStatus());
+        assertEquals(200, client.executeWithLoadBalancer(HttpRequest.newBuilder().uri(new URI("/")).loadBalancerKey(srcKey).build()).getStatus());
+        client.setLoadBalancer(new BaseLoadBalancer() {
+            public Server chooseServer(Object key) {
+                assertNotEquals(srcKey, key.toString());
+                addServers(Arrays.asList(servers));
+                return super.chooseServer(key);
+            }
+        });
+        assertEquals(200, client.executeWithLoadBalancer(HttpRequest.newBuilder().uri(new URI("/")).loadBalancerKey("77777").build()).getStatus());
+    }
+
     private void assertStatusIsOk(int status) {
         assertTrue(status == 200 || status == 302);
     }

ribbon-loadbalancer/src/main/java/com/netflix/client/AbstractLoadBalancerAwareClient.java

@@ -128,7 +128,8 @@ protected LoadBalancerCommand<T> buildLoadBalancerCommand(final S request, final
 		LoadBalancerCommand.Builder<T> builder = LoadBalancerCommand.<T>builder()
 				.withLoadBalancerContext(this)
 				.withRetryHandler(handler)
-				.withLoadBalancerURI(request.getUri());
+				.withLoadBalancerURI(request.getUri())
+				.withServerLocator(request.getLoadBalancerKey());
 		customizeLoadBalancerCommandBuilder(request, config, builder);
 		return builder.build();
 	}