Add hint-based instance filtering (#923)

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

@@ -0,0 +1,102 @@
+/*
+ * Copyright 2012-2021 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 reactor.core.publisher.Flux;
+
+import org.springframework.cloud.client.ServiceInstance;
+import org.springframework.cloud.client.loadbalancer.HintRequestContext;
+import org.springframework.cloud.client.loadbalancer.LoadBalancerProperties;
+import org.springframework.cloud.client.loadbalancer.Request;
+import org.springframework.cloud.client.loadbalancer.RequestDataContext;
+import org.springframework.http.HttpHeaders;
+import org.springframework.util.StringUtils;
+
+/**
+ * A {@link ServiceInstanceListSupplier} implementation that uses hints to filter service
+ * instances provided by the delegate.
+ *
+ * @author Olga Maciaszek-Sharma
+ * @since 3.0.2
+ */
+public class HintBasedServiceInstanceListSupplier extends DelegatingServiceInstanceListSupplier {
+
+	private final LoadBalancerProperties properties;
+
+	public HintBasedServiceInstanceListSupplier(ServiceInstanceListSupplier delegate,
+			LoadBalancerProperties properties) {
+		super(delegate);
+		this.properties = properties;
+	}
+
+	@Override
+	public Flux<List<ServiceInstance>> get() {
+		return delegate.get();
+	}
+
+	@Override
+	public Flux<List<ServiceInstance>> get(Request request) {
+		return get().map(instances -> filteredByHint(instances, getHint(request.getContext())));
+	}
+
+	private String getHint(Object requestContext) {
+		if (requestContext == null) {
+			return null;
+		}
+		String hint = null;
+		if (requestContext instanceof RequestDataContext) {
+			hint = getHintFromHeader((RequestDataContext) requestContext);
+		}
+		if (!StringUtils.hasText(hint) && requestContext instanceof HintRequestContext) {
+			hint = ((HintRequestContext) requestContext).getHint();
+		}
+		return hint;
+	}
+
+	private String getHintFromHeader(RequestDataContext context) {
+		if (context.getClientRequest() != null) {
+			HttpHeaders headers = context.getClientRequest().getHeaders();
+			if (headers != null) {
+				return headers.getFirst(properties.getHintHeaderName());
+			}
+		}
+		return null;
+	}
+
+	private List<ServiceInstance> filteredByHint(List<ServiceInstance> instances, String hint) {
+		if (!StringUtils.hasText(hint)) {
+			return instances;
+		}
+		List<ServiceInstance> filteredInstances = new ArrayList<>();
+		for (ServiceInstance serviceInstance : instances) {
+			if (serviceInstance.getMetadata().getOrDefault("hint", "").equals(hint)) {
+				filteredInstances.add(serviceInstance);
+			}
+		}
+		if (filteredInstances.size() > 0) {
+			return filteredInstances;
+		}
+
+		// If instances cannot be found based on hint,
+		// we return all instances retrieved for given service id.
+		return instances;
+	}
+
+}

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

@@ -48,11 +48,6 @@ public class RequestBasedStickySessionServiceInstanceListSupplier extends Delega
 		this.properties = properties;
 	}
 
-	@Override
-	public String getServiceId() {
-		return delegate.getServiceId();
-	}
-
 	@Override
 	public Flux<List<ServiceInstance>> get() {
 		return delegate.get();

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

@@ -239,6 +239,15 @@ public final class ServiceInstanceListSupplierBuilder {
 		return this;
 	}
 
+	public ServiceInstanceListSupplierBuilder withHints() {
+		DelegateCreator creator = (context, delegate) -> {
+			LoadBalancerProperties properties = context.getBean(LoadBalancerProperties.class);
+			return new HintBasedServiceInstanceListSupplier(delegate, properties);
+		};
+		creators.add(creator);
+		return this;
+	}
+
 	/**
 	 * Builds the {@link ServiceInstanceListSupplier} hierarchy.
 	 * @param context application context

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

@@ -0,0 +1,141 @@
+/*
+ * Copyright 2012-2021 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.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.junit.jupiter.api.BeforeEach;
+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.cloud.client.loadbalancer.DefaultRequest;
+import org.springframework.cloud.client.loadbalancer.LoadBalancerProperties;
+import org.springframework.cloud.client.loadbalancer.Request;
+import org.springframework.cloud.client.loadbalancer.RequestData;
+import org.springframework.cloud.client.loadbalancer.RequestDataContext;
+import org.springframework.mock.http.client.MockClientHttpRequest;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/**
+ * Tests for {@link HintBasedServiceInstanceListSupplier}.
+ *
+ * @author Olga Maciaszek-Sharma
+ */
+class HintBasedServiceInstanceListSupplierTests {
+
+	private final DiscoveryClientServiceInstanceListSupplier delegate = mock(
+			DiscoveryClientServiceInstanceListSupplier.class);
+
+	private final LoadBalancerProperties properties = new LoadBalancerProperties();
+
+	private final RequestDataContext requestContext = new RequestDataContext(
+			new RequestData(new MockClientHttpRequest()));
+
+	private final HintBasedServiceInstanceListSupplier supplier = new HintBasedServiceInstanceListSupplier(delegate,
+			properties);
+
+	private final ServiceInstance first = serviceInstance("test-1", buildHintMetadata("test1"));
+
+	private final ServiceInstance second = serviceInstance("test-2", buildHintMetadata("test2"));
+
+	private final ServiceInstance third = serviceInstance("test-3", new HashMap<>());
+
+	@BeforeEach
+	void setUp() {
+		properties.setHintHeaderName("X-Test");
+		when(delegate.get()).thenReturn(Flux.just(Arrays.asList(first, second, third)));
+	}
+
+	@Test
+	void shouldReturnInstancesForHintFromHeaderWhenAvailable() {
+		requestContext.setHint("test1");
+		requestContext.getClientRequest().getHeaders().add("X-Test", "test2");
+		Request<RequestDataContext> request = new DefaultRequest<>(requestContext);
+
+		List<ServiceInstance> filtered = supplier.get(request).blockFirst();
+
+		assertThat(filtered).hasSize(1);
+		assertThat(filtered.get(0).getInstanceId()).isEqualTo("test-2");
+	}
+
+	@Test
+	void shouldReturnInstancesForHintFromPropertiesWhenNoHintHeader() {
+		requestContext.setHint("test1");
+		Request<RequestDataContext> request = new DefaultRequest<>(requestContext);
+
+		List<ServiceInstance> filtered = supplier.get(request).blockFirst();
+
+		assertThat(filtered).hasSize(1);
+		assertThat(filtered.get(0).getInstanceId()).isEqualTo("test-1");
+	}
+
+	@Test
+	void shouldReturnAllInstancesWhenNoHint() {
+		Request<RequestDataContext> request = new DefaultRequest<>(requestContext);
+
+		List<ServiceInstance> filtered = supplier.get(request).blockFirst();
+
+		assertThat(filtered).hasSize(3);
+	}
+
+	@Test
+	void shouldReturnAllInstancesWhenHintNotMatched() {
+		requestContext.getClientRequest().getHeaders().add("X-Test", "testX");
+		Request<RequestDataContext> request = new DefaultRequest<>(requestContext);
+
+		List<ServiceInstance> filtered = supplier.get(request).blockFirst();
+
+		assertThat(filtered).hasSize(3);
+	}
+
+	@Test
+	void shouldReturnAllInstancesWhenRequestContextNull() {
+		Request<RequestDataContext> request = new DefaultRequest<>(null);
+
+		List<ServiceInstance> filtered = supplier.get(request).blockFirst();
+
+		assertThat(filtered).hasSize(3);
+	}
+
+	@Test
+	void shouldReturnAllInstancesWhenClientRequestNull() {
+		Request<RequestDataContext> request = new DefaultRequest<>(new RequestDataContext(null));
+
+		List<ServiceInstance> filtered = supplier.get(request).blockFirst();
+
+		assertThat(filtered).hasSize(3);
+	}
+
+	private DefaultServiceInstance serviceInstance(String instanceId, Map<String, String> metadata) {
+		return new DefaultServiceInstance(instanceId, "test", "http://test.test", 9080, false, metadata);
+	}
+
+	private Map<String, String> buildHintMetadata(String zone) {
+		Map<String, String> metadata = new HashMap<>();
+		metadata.put("hint", zone);
+		return metadata;
+	}
+
+}