Zoned ServiceInstanceListSupplier (#658)

* Get ServiceInstance zone from metadata.

* Add ZonePreferenceServiceInstanceListSupplier.

* Add javadocs and license entries.

* Add tests.

* Add documentation.

* Documentation fix.

* Fix after code review.

spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/core/ZonePreferenceServiceInstanceListSupplier.java

@@ -0,0 +1,93 @@
+/*
+ * Copyright 2012-2019 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
+ *
+ *      https://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.loadbalancer.core;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+import reactor.core.publisher.Flux;
+
+import org.springframework.cloud.client.ServiceInstance;
+import org.springframework.core.env.Environment;
+
+/**
+ * An implementation of {@link ServiceInstanceListSupplier} that filters instances
+ * retrieved by the delegate by zone. The zone is retrieved from the
+ * <code>spring.cloud.loadbalancer.zone</code> property. If the zone is not set or no instances are found for the
+ * requested zone, all instances retrieved by the delegate are returned.
+ *
+ * @author Olga Maciaszek-Sharma
+ * @since 2.2.1
+ */
+public class ZonePreferenceServiceInstanceListSupplier
+		implements ServiceInstanceListSupplier {
+
+	private final String ZONE = "zone";
+
+	private final ServiceInstanceListSupplier delegate;
+
+	private final Environment environment;
+
+	private String zone;
+
+	public ZonePreferenceServiceInstanceListSupplier(ServiceInstanceListSupplier delegate,
+			Environment environment) {
+		this.delegate = delegate;
+		this.environment = environment;
+	}
+
+	@Override
+	public String getServiceId() {
+		return delegate.getServiceId();
+	}
+
+	@Override
+	public Flux<List<ServiceInstance>> get() {
+		return delegate.get().map(this::filteredByZone);
+	}
+
+	private List<ServiceInstance> filteredByZone(List<ServiceInstance> serviceInstances) {
+		if (zone == null) {
+			zone = environment.getProperty("spring.cloud.loadbalancer.zone");
+		}
+		if (zone != null) {
+			List<ServiceInstance> filteredInstances = new ArrayList<>();
+			for (ServiceInstance serviceInstance : serviceInstances) {
+				String instanceZone = getZone(serviceInstance);
+				if (zone.equalsIgnoreCase(instanceZone)) {
+					filteredInstances.add(serviceInstance);
+				}
+			}
+			if (filteredInstances.size() > 0) {
+				return filteredInstances;
+			}
+		}
+		// If the zone is not set or there are no zone-specific instances available,
+		// we return all instances retrieved for given service id.
+		return serviceInstances;
+	}
+
+	private String getZone(ServiceInstance serviceInstance) {
+		Map<String, String> metadata = serviceInstance.getMetadata();
+		if (metadata != null) {
+			return metadata.get(ZONE);
+		}
+		return null;
+	}
+
+}

spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/core/ZonePreferenceServiceInstanceListSupplierTests.java

@@ -0,0 +1,124 @@
+/*
+ * Copyright 2012-2019 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
+ *
+ *      https://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.loadbalancer.core;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Flux;
+
+import org.springframework.cloud.client.DefaultServiceInstance;
+import org.springframework.cloud.client.ServiceInstance;
+import org.springframework.core.env.Environment;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/**
+ * Tests for {@link ZonePreferenceServiceInstanceListSupplier}.
+ *
+ * @author Olga Maciaszek-Sharma
+ */
+class ZonePreferenceServiceInstanceListSupplierTests {
+
+	private DiscoveryClientServiceInstanceListSupplier delegate = mock(
+			DiscoveryClientServiceInstanceListSupplier.class);
+
+	private Environment environment = mock(Environment.class);
+
+	private ZonePreferenceServiceInstanceListSupplier supplier = new ZonePreferenceServiceInstanceListSupplier(
+			delegate, environment);
+
+	private ServiceInstance first = serviceInstance("test-1", buildZoneMetadata("zone1"));
+
+	private ServiceInstance second = serviceInstance("test-2",
+			buildZoneMetadata("zone1"));
+
+	private ServiceInstance third = serviceInstance("test-3", buildZoneMetadata("zone2"));
+
+	private ServiceInstance fourth = serviceInstance("test-4",
+			buildZoneMetadata("zone3"));
+
+	private ServiceInstance fifth = serviceInstance("test-5", buildZoneMetadata(null));
+
+	@Test
+	void shouldFilterInstancesByZone() {
+		when(environment.getProperty("spring.cloud.loadbalancer.zone"))
+				.thenReturn("zone1");
+		when(delegate.get()).thenReturn(Flux.just(Arrays.asList(first, second, third, fourth, fifth)));
+
+		List<ServiceInstance> filtered = supplier.get().blockFirst();
+
+		assertThat(filtered).hasSize(2);
+		assertThat(filtered).contains(first, second);
+		assertThat(filtered).doesNotContain(third);
+		assertThat(filtered).doesNotContain(fourth);
+		assertThat(filtered).doesNotContain(fifth);
+	}
+
+	@Test
+	void shouldReturnAllInstancesIfNoZoneInstances() {
+		when(environment.getProperty("spring.cloud.loadbalancer.zone"))
+				.thenReturn("zone1");
+		when(delegate.get()).thenReturn(Flux.just(Arrays.asList(third, fourth)));
+
+		List<ServiceInstance> filtered = supplier.get().blockFirst();
+
+		assertThat(filtered).hasSize(2);
+		assertThat(filtered).contains(third, fourth);
+	}
+
+	@Test
+	void shouldNotThrowNPEIfNullInstanceMetadata() {
+		when(environment.getProperty("spring.cloud.loadbalancer.zone"))
+				.thenReturn("zone1");
+		when(delegate.get()).thenReturn(
+				Flux.just(Collections.singletonList(serviceInstance("test-6", null))));
+		assertThatCode(() -> supplier.get().blockFirst()).doesNotThrowAnyException();
+	}
+
+	@Test
+	void shouldReturnAllInstancesIfNoZone() {
+		when(environment.getProperty("spring.cloud.loadbalancer.zone")).thenReturn(null);
+		when(delegate.get())
+				.thenReturn(Flux.just(Arrays.asList(first, second, third, fourth)));
+
+		List<ServiceInstance> filtered = supplier.get().blockFirst();
+
+		assertThat(filtered).hasSize(4);
+		assertThat(filtered).contains(first, second, third, fourth);
+	}
+
+	private DefaultServiceInstance serviceInstance(String instanceId,
+			Map<String, String> metadata) {
+		return new DefaultServiceInstance("test", instanceId, "http://test.test", 9080,
+				false, metadata);
+	}
+
+	private Map<String, String> buildZoneMetadata(String zone) {
+		Map<String, String> metadata = new HashMap<>();
+		metadata.put("zone", zone);
+		return metadata;
+	}
+
+}