Make load balancer prefer the current zone if there is one

spring-cloud-netflix-core/pom.xml

@@ -113,6 +113,10 @@
 			<artifactId>neo4j-cypher-compiler-2.1</artifactId>
 			<version>2.1.2</version>
 		</dependency>
+		<dependency>
+			<groupId>org.apache.activemq</groupId>
+			<artifactId>activemq-client</artifactId>
+		</dependency>
 	</dependencies>
 
 </project>

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

 package org.springframework.cloud.netflix.ribbon;
 
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
 import org.springframework.cloud.netflix.eureka.EurekaClientAutoConfiguration;
+import org.springframework.cloud.netflix.ribbon.eureka.EurekaRibbonInitializer;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.http.client.ClientHttpRequestInterceptor;
 import org.springframework.web.client.RestTemplate;
 
+import com.netflix.loadbalancer.BaseLoadBalancer;
+
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 /**
@@ -19,6 +24,12 @@ import java.util.List;
 @AutoConfigureAfter(EurekaClientAutoConfiguration.class)
 public class RibbonAutoConfiguration {
 	
+	@Autowired(required=false)
+	private List<BaseLoadBalancer> balancers = Collections.emptyList();
+
+	@Autowired(required=false)
+	private EurekaRibbonInitializer initializer;
+
     @Bean
     @ConditionalOnMissingBean(RestTemplate.class)
     public RestTemplate restTemplate(RibbonInterceptor ribbonInterceptor) {
@@ -32,7 +43,7 @@ public class RibbonAutoConfiguration {
     @Bean
     @ConditionalOnMissingBean(LoadBalancerClient.class)
     public LoadBalancerClient loadBalancerClient() {
-        return new RibbonLoadBalancerClient();
+        return new RibbonLoadBalancerClient(balancers);
     }
 
     @Bean

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

 package org.springframework.cloud.netflix.ribbon;
 
-import com.netflix.client.ClientFactory;
-import com.netflix.loadbalancer.ILoadBalancer;
-import com.netflix.loadbalancer.Server;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.cloud.client.ServiceInstance;
 import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
 
+import com.netflix.client.ClientFactory;
+import com.netflix.loadbalancer.BaseLoadBalancer;
+import com.netflix.loadbalancer.ILoadBalancer;
+import com.netflix.loadbalancer.Server;
+
 /**
  * @author Spencer Gibb
+ * @author Dave Syer
  */
 public class RibbonLoadBalancerClient implements LoadBalancerClient {
+
 	@Autowired
 	private ServerListInitializer serverListInitializer;
 
+	private Map<String, ILoadBalancer> balancers = new HashMap<String, ILoadBalancer>();
+
+	public RibbonLoadBalancerClient(List<BaseLoadBalancer> balancers) {
+		for (BaseLoadBalancer balancer : balancers) {
+			this.balancers.put(balancer.getName(), balancer);
+		}
+	}
+
 	@Override
 	public ServiceInstance choose(String serviceId) {
 		serverListInitializer.initialize(serviceId);
-        ILoadBalancer loadBalancer = ClientFactory.getNamedLoadBalancer(serviceId);
-        Server server = loadBalancer.chooseServer(null);
+		ILoadBalancer loadBalancer = this.balancers.get(serviceId);
+		if (loadBalancer == null) {
+			loadBalancer = ClientFactory.getNamedLoadBalancer(serviceId);
+		}
+		Server server = loadBalancer.chooseServer("default");
 		if (server == null) {
-            throw new IllegalStateException("Unable to locate ILoadBalancer for service: "+ serviceId);
+			throw new IllegalStateException(
+					"Unable to locate ILoadBalancer for service: " + serviceId);
 		}
 		return new RibbonServer(serviceId, server);
 	}
@@ -46,7 +66,7 @@ public class RibbonLoadBalancerClient implements LoadBalancerClient {
 
 		@Override
 		public String getIpAddress() {
-            return null; //TODO: ribbon doesn't supply ip
+			return null; // TODO: ribbon doesn't supply ip
 		}
 
 		@Override

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

+/*
+ * Copyright 2013-2014 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.ribbon.eureka;
+
+import java.net.MalformedURLException;
+import java.net.URL;
+import java.util.List;
+
+import org.springframework.util.StringUtils;
+
+import com.netflix.loadbalancer.Server;
+import com.netflix.loadbalancer.ServerList;
+
+/**
+ * @author Dave Syer
+ *
+ */
+public class DomainExtractingServerList implements ServerList<Server> {
+
+	private ServerList<Server> list;
+
+	public DomainExtractingServerList(ServerList<Server> list) {
+		this.list = list;
+	}
+
+	@Override
+	public List<Server> getInitialListOfServers() {
+		List<Server> servers = setZones(list.getInitialListOfServers());
+		return servers;
+	}
+
+	@Override
+	public List<Server> getUpdatedListOfServers() {
+		List<Server> servers = setZones(list.getUpdatedListOfServers());
+		return servers;
+	}
+
+	private List<Server> setZones(List<Server> servers) {
+		for (Server server : servers) {
+			if (!server.getZone().equals("default")) {
+				String zone = extractApproximateZone(server.getId());
+				server.setZone(zone);
+			}
+		}
+		return servers;
+	}
+
+	private String extractApproximateZone(String id) {
+		try {
+			URL url = new URL("http://" + id);
+			String host = url.getHost();
+			if (!host.contains(".")) {
+				return host;
+			}
+			String[] split = StringUtils.split(host, ".");
+			StringBuilder builder = new StringBuilder(split[1]);
+			for (int i=2; i<split.length; i++) {
+				builder.append(".").append(split[i]);
+			}
+			return builder.toString();
+		}
+		catch (MalformedURLException e) {
+		}
+		return "defaultZone";
+	}
+
+}
+

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

 package org.springframework.cloud.netflix.ribbon.eureka;
 
-import com.netflix.config.ConfigurationManager;
-import com.netflix.niws.loadbalancer.DiscoveryEnabledNIWSServerList;
+import static com.netflix.client.config.CommonClientConfigKey.DeploymentContextBasedVipAddresses;
+import static com.netflix.client.config.CommonClientConfigKey.EnableZoneAffinity;
+import static com.netflix.client.config.CommonClientConfigKey.NFLoadBalancerRuleClassName;
+import static com.netflix.client.config.CommonClientConfigKey.NIWSServerListClassName;
+import static com.netflix.client.config.CommonClientConfigKey.NIWSServerListFilterClassName;
+
 import org.springframework.cloud.netflix.ribbon.ServerListInitializer;
 
-import static com.netflix.client.config.CommonClientConfigKey.*;
+import com.netflix.appinfo.AmazonInfo;
+import com.netflix.appinfo.EurekaInstanceConfig;
+import com.netflix.client.ClientFactory;
+import com.netflix.config.ConfigurationManager;
+import com.netflix.config.DeploymentContext.ContextKey;
+import com.netflix.loadbalancer.DynamicServerListLoadBalancer;
+import com.netflix.loadbalancer.ILoadBalancer;
+import com.netflix.loadbalancer.Server;
+import com.netflix.loadbalancer.ServerList;
+import com.netflix.loadbalancer.ZoneAvoidanceRule;
+import com.netflix.niws.loadbalancer.DiscoveryEnabledNIWSServerList;
 
 /**
+ * Convenience class that sets up some configuration defaults for eureka-discovered ribbon
+ * clients.
+ * 
  * @author Spencer Gibb
+ * @author Dave Syer
  */
 public class EurekaRibbonInitializer implements ServerListInitializer {
 
+	private EurekaInstanceConfig instance;
+
+	public EurekaRibbonInitializer(EurekaInstanceConfig instance) {
+		this.instance = instance;
+	}
+
 	@Override
 	public void initialize(String serviceId) {
-        //TODO: should this look more like hibernate spring boot props?
-        //TODO: only set the property if it hasn't already been set?
-        setProp(serviceId, NIWSServerListClassName.key(), DiscoveryEnabledNIWSServerList.class.getName());
-        setProp(serviceId, DeploymentContextBasedVipAddresses.key(), serviceId); //FIXME: what should this be?
+		if (instance != null
+				&& ConfigurationManager.getDeploymentContext().getValue(ContextKey.zone) == null) {
+			// You can set this with archaius.deployment.* (maybe requires
+			// custom deployment context)?
+			String zone = instance.getMetadataMap().get("zone");
+			if (zone == null && instance.getDataCenterInfo() instanceof AmazonInfo) {
+				AmazonInfo info = (AmazonInfo) instance.getDataCenterInfo();
+				zone = info.getMetadata().get(AmazonInfo.MetaDataKey.availabilityZone);
+			}
+			if (zone != null) {
+				ConfigurationManager.getDeploymentContext().setValue(ContextKey.zone,
+						zone);
+			}
+		}
+		// TODO: should this look more like hibernate spring boot props?
+		// TODO: only set the property if it hasn't already been set?
+		setProp(serviceId, NIWSServerListClassName.key(),
+				DiscoveryEnabledNIWSServerList.class.getName());
+		// FIXME: what should this be?
+		setProp(serviceId, DeploymentContextBasedVipAddresses.key(), serviceId);
+		setProp(serviceId, NFLoadBalancerRuleClassName.key(),
+				ZoneAvoidanceRule.class.getName());
+		setProp(serviceId, NIWSServerListFilterClassName.key(),
+				ZonePreferenceServerListFilter.class.getName());
+		setProp(serviceId, EnableZoneAffinity.key(), "true");
+		ILoadBalancer loadBalancer = ClientFactory.getNamedLoadBalancer(serviceId);
+		wrapServerList(loadBalancer);
+	}
+
+	private void wrapServerList(ILoadBalancer balancer) {
+		if (balancer instanceof DynamicServerListLoadBalancer) {
+			@SuppressWarnings("unchecked")
+			DynamicServerListLoadBalancer<Server> dynamic = (DynamicServerListLoadBalancer<Server>) balancer;
+			ServerList<Server> list = dynamic.getServerListImpl();
+			if (!(list instanceof DomainExtractingServerList)
+					&& !(instance.getDataCenterInfo() instanceof AmazonInfo)) {
+				// This is optional: you can use the native Eureka AWS features by making
+				// the EurekaInstanceConfig.dataCenterInfo an AmazonInfo
+				dynamic.setServerListImpl(new DomainExtractingServerList(list));
+			}
+		}
 	}
 
 	protected void setProp(String serviceId, String suffix, String value) {
-        //how to set the namespace properly?
+		// how to set the namespace properly?
 		String namespace = "ribbon";
-        ConfigurationManager.getConfigInstance().setProperty(serviceId + "."+ namespace +"." + suffix, value);
+		ConfigurationManager.getConfigInstance().setProperty(
+				serviceId + "." + namespace + "." + suffix, value);
 	}
+
 }

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

@@ -15,7 +15,10 @@
  */
 package org.springframework.cloud.netflix.ribbon.eureka;
 
+import com.netflix.appinfo.EurekaInstanceConfig;
 import com.netflix.niws.loadbalancer.DiscoveryEnabledNIWSServerList;
+
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnExpression;
@@ -36,8 +39,11 @@ import org.springframework.context.annotation.Configuration;
 @AutoConfigureAfter(RibbonAutoConfiguration.class)
 public class RibbonEurekaAutoConfiguration {
 	
+	@Autowired(required=false)
+	private EurekaInstanceConfig instance;
+
     @Bean
     public ServerListInitializer serverListInitializer() {
-        return new EurekaRibbonInitializer();
+        return new EurekaRibbonInitializer(instance);
     }
 }

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

+/*
+ * Copyright 2013-2014 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.ribbon.eureka;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import com.netflix.client.config.IClientConfig;
+import com.netflix.config.ConfigurationManager;
+import com.netflix.config.DeploymentContext.ContextKey;
+import com.netflix.loadbalancer.Server;
+import com.netflix.loadbalancer.ZoneAffinityServerListFilter;
+
+/**
+ * A filter that actively prefers the local zone (as defined by the deployment context, or
+ * the Eureka instance metadata).
+ * 
+ * @author Dave Syer
+ *
+ */
+public class ZonePreferenceServerListFilter extends ZoneAffinityServerListFilter<Server> {
+
+	private String zone;
+
+	@Override
+	public void initWithNiwsConfig(IClientConfig niwsClientConfig) {
+		super.initWithNiwsConfig(niwsClientConfig);
+		if (ConfigurationManager.getDeploymentContext() != null) {
+			zone = ConfigurationManager.getDeploymentContext().getValue(ContextKey.zone);
+		}
+	}
+
+	@Override
+	public List<Server> getFilteredListOfServers(List<Server> servers) {
+		List<Server> output = super.getFilteredListOfServers(servers);
+		if (zone != null && output.size() == servers.size()) {
+			List<Server> local = new ArrayList<Server>();
+			for (Server server : output) {
+				if (zone.equalsIgnoreCase(server.getZone())) {
+					local.add(server);
+				}
+			}
+			if (!local.isEmpty()) {
+				return local;
+			}
+		}
+		return output;
+	}
+
+}

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

+/*
+ * Copyright 2013-2014 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.ribbon.eureka;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import org.junit.After;
+import org.junit.Test;
+import org.springframework.cloud.netflix.eureka.EurekaInstanceConfigBean;
+
+import com.netflix.client.ClientFactory;
+import com.netflix.config.ConfigurationManager;
+import com.netflix.config.DeploymentContext.ContextKey;
+import com.netflix.loadbalancer.ILoadBalancer;
+import com.netflix.loadbalancer.Server;
+import com.netflix.loadbalancer.ZoneAwareLoadBalancer;
+
+/**
+ * @author Dave Syer
+ *
+ */
+public class EurekaRibbonInitializerTests {
+	
+	@After
+	public void close() {
+		ConfigurationManager.getDeploymentContext().setValue(ContextKey.zone, "");
+	}
+
+	@Test
+	public void basicConfigurationCreatedForLoadBalancer() {
+		EurekaInstanceConfigBean instance = new EurekaInstanceConfigBean();
+		instance.getMetadataMap().put("zone", "foo");
+		EurekaRibbonInitializer initializer = new EurekaRibbonInitializer(
+				instance);
+		initializer.initialize("service");
+		ILoadBalancer balancer = ClientFactory.getNamedLoadBalancer("service");
+		assertNotNull(balancer);
+		@SuppressWarnings("unchecked")
+		ZoneAwareLoadBalancer<Server> aware = (ZoneAwareLoadBalancer<Server>) balancer;
+		assertTrue(aware.getServerListImpl() instanceof DomainExtractingServerList);
+		assertEquals("foo", ConfigurationManager.getDeploymentContext().getValue(ContextKey.zone));
+	}
+
+}

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

+/*
+ * Copyright 2013-2014 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.ribbon.eureka;
+
+import static org.junit.Assert.assertEquals;
+
+import java.util.Arrays;
+import java.util.List;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.springframework.test.util.ReflectionTestUtils;
+
+import com.netflix.loadbalancer.Server;
+
+/**
+ * @author Dave Syer
+ *
+ */
+public class ZonePreferenceServerListFilterTests {
+
+	private Server dsyer = new Server("dsyer", 8080);
+	private Server localhost = new Server("localhost", 8080);
+	
+	@Before
+	public void init() {
+		dsyer.setZone("dsyer");
+		localhost.setZone("localhost");
+	}
+
+	@Test
+	public void noZoneSet() {
+		ZonePreferenceServerListFilter filter = new ZonePreferenceServerListFilter();
+		List<Server> result = filter.getFilteredListOfServers(Arrays.asList(localhost));
+		assertEquals(1, result.size());
+	}
+
+	@Test
+	public void withZoneSetAndNoMatches() {
+		ZonePreferenceServerListFilter filter = new ZonePreferenceServerListFilter();
+		ReflectionTestUtils.setField(filter, "zone", "dsyer");
+		List<Server> result = filter.getFilteredListOfServers(Arrays.asList(localhost));
+		assertEquals(1, result.size());
+	}
+
+	@Test
+	public void withZoneSetAndMatches() {
+		ZonePreferenceServerListFilter filter = new ZonePreferenceServerListFilter();
+		ReflectionTestUtils.setField(filter, "zone", "dsyer");
+		List<Server> result = filter.getFilteredListOfServers(Arrays.asList(dsyer, localhost));
+		assertEquals(1, result.size());
+	}
+
+}