Request based sticky session (#860)

* Implement first draft for sticky-session load-balancing.

* Make setting instance cookie opt-in.

* Add tests.

* Make adding request cookie opt-in.

* Add docs and javadocs.

* Add default configuration with blocking discovery client.

* Fix docs after review.

spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/annotation/LoadBalancerClientConfiguration.java

@@ -98,6 +98,17 @@ public class LoadBalancerClientConfiguration {
 			return ServiceInstanceListSupplier.builder().withDiscoveryClient().withHealthChecks().build(context);
 		}
 
+		@Bean
+		@ConditionalOnBean(ReactiveDiscoveryClient.class)
+		@ConditionalOnMissingBean
+		@ConditionalOnProperty(value = "spring.cloud.loadbalancer.configurations",
+				havingValue = "request-based-sticky-session")
+		public ServiceInstanceListSupplier requestBasedStickySessionDiscoveryClientServiceInstanceListSupplier(
+				ConfigurableApplicationContext context) {
+			return ServiceInstanceListSupplier.builder().withDiscoveryClient().withRequestBasedStickySession()
+					.build(context);
+		}
+
 		@Bean
 		@ConditionalOnBean(ReactiveDiscoveryClient.class)
 		@ConditionalOnMissingBean
@@ -146,6 +157,17 @@ public class LoadBalancerClientConfiguration {
 					.build(context);
 		}
 
+		@Bean
+		@ConditionalOnBean(DiscoveryClient.class)
+		@ConditionalOnMissingBean
+		@ConditionalOnProperty(value = "spring.cloud.loadbalancer.configurations",
+				havingValue = "request-based-sticky-session")
+		public ServiceInstanceListSupplier requestBasedStickySessionDiscoveryClientServiceInstanceListSupplier(
+				ConfigurableApplicationContext context) {
+			return ServiceInstanceListSupplier.builder().withBlockingDiscoveryClient().withRequestBasedStickySession()
+					.build(context);
+		}
+
 		@Bean
 		@ConditionalOnBean(DiscoveryClient.class)
 		@ConditionalOnMissingBean

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

@@ -0,0 +1,121 @@
+/*
+ * Copyright 2012-2020 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.Collections;
+import java.util.List;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import reactor.core.publisher.Flux;
+
+import org.springframework.cloud.client.ServiceInstance;
+import org.springframework.cloud.client.loadbalancer.ClientRequestContext;
+import org.springframework.cloud.client.loadbalancer.Request;
+import org.springframework.cloud.client.loadbalancer.ServerHttpRequestContext;
+import org.springframework.cloud.client.loadbalancer.reactive.LoadBalancerProperties;
+import org.springframework.http.HttpCookie;
+import org.springframework.http.server.reactive.ServerHttpRequest;
+import org.springframework.web.reactive.function.client.ClientRequest;
+
+/**
+ * A session cookie based implementation of {@link ServiceInstanceListSupplier} that gives
+ * preference to the instance with an id specified in a request cookie.
+ *
+ * @author Olga Maciaszek-Sharma
+ * @since 3.0.0
+ */
+public class RequestBasedStickySessionServiceInstanceListSupplier extends DelegatingServiceInstanceListSupplier {
+
+	private static final Log LOG = LogFactory.getLog(RequestBasedStickySessionServiceInstanceListSupplier.class);
+
+	private final LoadBalancerProperties properties;
+
+	public RequestBasedStickySessionServiceInstanceListSupplier(ServiceInstanceListSupplier delegate,
+			LoadBalancerProperties properties) {
+		super(delegate);
+		this.properties = properties;
+	}
+
+	@Override
+	public String getServiceId() {
+		return delegate.getServiceId();
+	}
+
+	@Override
+	public Flux<List<ServiceInstance>> get() {
+		return delegate.get();
+	}
+
+	@SuppressWarnings("rawtypes")
+	@Override
+	public Flux<List<ServiceInstance>> get(Request request) {
+		String instanceIdCookieName = properties.getStickySession().getInstanceIdCookieName();
+		Object context = request.getContext();
+		if ((context instanceof ClientRequestContext)) {
+			ClientRequest originalRequest = ((ClientRequestContext) context).getClientRequest();
+			// We expect there to be one value in this cookie
+			String cookie = originalRequest.cookies().getFirst(instanceIdCookieName);
+			if (cookie != null) {
+				return get().map(serviceInstances -> selectInstance(serviceInstances, cookie));
+			}
+			if (LOG.isDebugEnabled()) {
+				LOG.debug("Cookie not found. Returning all instances returned by delegate.");
+			}
+			return get();
+		}
+		if ((context instanceof ServerHttpRequestContext)) {
+			ServerHttpRequest originalRequest = ((ServerHttpRequestContext) context).getClientRequest();
+			HttpCookie cookie = originalRequest.getCookies().getFirst(instanceIdCookieName);
+			if (cookie != null) {
+				return get().map(serviceInstances -> selectInstance(serviceInstances, cookie.getValue()));
+			}
+			if (LOG.isDebugEnabled()) {
+				LOG.debug("Cookie not found. Returning all instances returned by delegate.");
+			}
+			return get();
+		}
+		if (LOG.isDebugEnabled()) {
+			LOG.debug("Searching for instances based on cookie not supported for ClientRequestContext type."
+					+ " Returning all instances returned by delegate.");
+		}
+		// If no cookie is available, we return all the instances provided by the
+		// delegate.
+		return get();
+	}
+
+	private List<ServiceInstance> selectInstance(List<ServiceInstance> serviceInstances, String cookie) {
+		for (ServiceInstance serviceInstance : serviceInstances) {
+			if (cookie.equals(serviceInstance.getInstanceId())) {
+				if (LOG.isDebugEnabled()) {
+					LOG.debug(String.format("Returning the service instance: %s. Found for cookie: %s",
+							serviceInstance.toString(), cookie));
+				}
+				return Collections.singletonList(serviceInstance);
+			}
+		}
+		// If the instances cannot be found based on the cookie,
+		// we return all the instances provided by the delegate.
+		if (LOG.isDebugEnabled()) {
+			LOG.debug(String.format(
+					"Service instance for cookie: %s not found. Returning all instances returned by delegate.",
+					cookie));
+		}
+		return serviceInstances;
+	}
+
+}

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

@@ -157,6 +157,20 @@ public final class ServiceInstanceListSupplierBuilder {
 		return this;
 	}
 
+	/**
+	 * Adds a {@link RequestBasedStickySessionServiceInstanceListSupplier} to the
+	 * {@link ServiceInstanceListSupplier} hierarchy.
+	 * @return the {@link ServiceInstanceListSupplierBuilder} object
+	 */
+	public ServiceInstanceListSupplierBuilder withRequestBasedStickySession() {
+		DelegateCreator creator = (context, delegate) -> {
+			LoadBalancerProperties properties = context.getBean(LoadBalancerProperties.class);
+			return new RequestBasedStickySessionServiceInstanceListSupplier(delegate, properties);
+		};
+		this.creators.add(creator);
+		return this;
+	}
+
 	/**
 	 * If {@link LoadBalancerCacheManager} is available in the context, wraps created
 	 * {@link ServiceInstanceListSupplier} hierarchy with a

spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/annotation/LoadBalancerClientConfigurationTests.java

@@ -30,6 +30,7 @@ import org.springframework.cloud.loadbalancer.core.CachingServiceInstanceListSup
 import org.springframework.cloud.loadbalancer.core.DelegatingServiceInstanceListSupplier;
 import org.springframework.cloud.loadbalancer.core.DiscoveryClientServiceInstanceListSupplier;
 import org.springframework.cloud.loadbalancer.core.HealthCheckServiceInstanceListSupplier;
+import org.springframework.cloud.loadbalancer.core.RequestBasedStickySessionServiceInstanceListSupplier;
 import org.springframework.cloud.loadbalancer.core.RetryAwareServiceInstanceListSupplier;
 import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
 import org.springframework.cloud.loadbalancer.core.ZonePreferenceServiceInstanceListSupplier;
@@ -111,6 +112,19 @@ class LoadBalancerClientConfigurationTests {
 				});
 	}
 
+	@Test
+	void shouldInstantiateRequestBasedStickySessionServiceInstanceListSupplierTests() {
+		reactiveDiscoveryClientRunner.withUserConfiguration(TestConfig.class)
+				.withPropertyValues("spring.cloud.loadbalancer.configurations=request-based-sticky-session")
+				.run(context -> {
+					ServiceInstanceListSupplier supplier = context.getBean(ServiceInstanceListSupplier.class);
+					then(supplier).isInstanceOf(RequestBasedStickySessionServiceInstanceListSupplier.class);
+					ServiceInstanceListSupplier delegate = ((DelegatingServiceInstanceListSupplier) supplier)
+							.getDelegate();
+					then(delegate).isInstanceOf(DiscoveryClientServiceInstanceListSupplier.class);
+				});
+	}
+
 	@Test
 	void shouldInstantiateDefaultBlockingServiceInstanceListSupplierWhenConfigurationsPropertyNotSet() {
 		blockingDiscoveryClientRunner.run(context -> {

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

@@ -0,0 +1,164 @@
+/*
+ * Copyright 2012-2020 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.List;
+
+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.ClientRequestContext;
+import org.springframework.cloud.client.loadbalancer.DefaultRequest;
+import org.springframework.cloud.client.loadbalancer.DefaultRequestContext;
+import org.springframework.cloud.client.loadbalancer.Request;
+import org.springframework.cloud.client.loadbalancer.ServerHttpRequestContext;
+import org.springframework.cloud.client.loadbalancer.reactive.LoadBalancerProperties;
+import org.springframework.http.HttpCookie;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.server.reactive.ServerHttpRequest;
+import org.springframework.util.LinkedMultiValueMap;
+import org.springframework.util.MultiValueMap;
+import org.springframework.web.reactive.function.client.ClientRequest;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/**
+ * Tests for {@link RequestBasedStickySessionServiceInstanceListSupplier}.
+ *
+ * @author Olga Maciaszek-Sharma
+ */
+class RequestBasedStickySessionServiceInstanceListSupplierTests {
+
+	private final DiscoveryClientServiceInstanceListSupplier delegate = mock(
+			DiscoveryClientServiceInstanceListSupplier.class);
+
+	private final LoadBalancerProperties properties = new LoadBalancerProperties();
+
+	private final RequestBasedStickySessionServiceInstanceListSupplier supplier = new RequestBasedStickySessionServiceInstanceListSupplier(
+			delegate, properties);
+
+	private final ClientRequest clientRequest = mock(ClientRequest.class);
+
+	private final ServerHttpRequest serverHttpRequest = mock(ServerHttpRequest.class);
+
+	private final ServiceInstance first = serviceInstance("test-1");
+
+	private final ServiceInstance second = serviceInstance("test-2");
+
+	private final ServiceInstance third = serviceInstance("test-3");
+
+	@BeforeEach
+	void setUp() {
+		when(delegate.get()).thenReturn(Flux.just(Arrays.asList(first, second, third)));
+	}
+
+	@Test
+	void shouldReturnInstanceBasedOnCookieFromClientRequest() {
+		HttpHeaders headers = new HttpHeaders();
+		headers.add(properties.getStickySession().getInstanceIdCookieName(), "test-1");
+		when(clientRequest.cookies()).thenReturn(headers);
+		Request<ClientRequestContext> request = new DefaultRequest<>(new ClientRequestContext(clientRequest));
+
+		List<ServiceInstance> serviceInstances = supplier.get(request).blockFirst();
+
+		assertThat(serviceInstances).hasSize(1);
+		assertThat(serviceInstances.get(0).getInstanceId()).isEqualTo("test-1");
+	}
+
+	@Test
+	void shouldReturnAllDelegateInstancesIfInstanceBasedOnCookieFromClientRequestNotFound() {
+		HttpHeaders headers = new HttpHeaders();
+		headers.add(properties.getStickySession().getInstanceIdCookieName(), "test-4");
+		when(clientRequest.cookies()).thenReturn(headers);
+		Request<ClientRequestContext> request = new DefaultRequest<>(new ClientRequestContext(clientRequest));
+
+		List<ServiceInstance> serviceInstances = supplier.get(request).blockFirst();
+
+		assertThat(serviceInstances).hasSize(3);
+	}
+
+	@Test
+	void shouldReturnAllInstancesFromDelegateIfClientRequestHasNoCookie() {
+		when(clientRequest.cookies()).thenReturn(new HttpHeaders());
+		Request<ClientRequestContext> request = new DefaultRequest<>(new ClientRequestContext(clientRequest));
+
+		List<ServiceInstance> serviceInstances = supplier.get(request).blockFirst();
+
+		assertThat(serviceInstances).hasSize(3);
+	}
+
+	@Test
+	void shouldReturnInstanceBasedOnCookieFromServerHttpRequest() {
+		MultiValueMap<String, HttpCookie> cookies = new LinkedMultiValueMap<>();
+		cookies.add(properties.getStickySession().getInstanceIdCookieName(),
+				new HttpCookie(properties.getStickySession().getInstanceIdCookieName(), "test-1"));
+		when(serverHttpRequest.getCookies()).thenReturn(cookies);
+		Request<ServerHttpRequestContext> request = new DefaultRequest<>(
+				new ServerHttpRequestContext(serverHttpRequest));
+
+		List<ServiceInstance> serviceInstances = supplier.get(request).blockFirst();
+
+		assertThat(serviceInstances).hasSize(1);
+		assertThat(serviceInstances.get(0).getInstanceId()).isEqualTo("test-1");
+	}
+
+	@Test
+	void shouldReturnAllDelegateInstancesIfInstanceBasedOnCookieFromServerHttpRequestNotFound() {
+		MultiValueMap<String, HttpCookie> cookies = new LinkedMultiValueMap<>();
+		cookies.add(properties.getStickySession().getInstanceIdCookieName(),
+				new HttpCookie(properties.getStickySession().getInstanceIdCookieName(), "test-4"));
+		when(serverHttpRequest.getCookies()).thenReturn(cookies);
+		Request<ServerHttpRequestContext> request = new DefaultRequest<>(
+				new ServerHttpRequestContext(serverHttpRequest));
+
+		List<ServiceInstance> serviceInstances = supplier.get(request).blockFirst();
+
+		assertThat(serviceInstances).hasSize(3);
+	}
+
+	@Test
+	void shouldReturnAllInstancesFromDelegateIfServerHttpRequestHasNoCookie() {
+		MultiValueMap<String, HttpCookie> cookies = new LinkedMultiValueMap<>();
+		when(serverHttpRequest.getCookies()).thenReturn(cookies);
+		Request<ServerHttpRequestContext> request = new DefaultRequest<>(
+				new ServerHttpRequestContext(serverHttpRequest));
+
+		List<ServiceInstance> serviceInstances = supplier.get(request).blockFirst();
+
+		assertThat(serviceInstances).hasSize(3);
+	}
+
+	@Test
+	void shouldReturnAllInstancesFromDelegateIfNotSupportedRequestContext() {
+		Request<DefaultRequestContext> request = new DefaultRequest<>(new DefaultRequestContext(clientRequest));
+
+		List<ServiceInstance> serviceInstances = supplier.get(request).blockFirst();
+
+		assertThat(serviceInstances).hasSize(3);
+	}
+
+	private DefaultServiceInstance serviceInstance(String instanceId) {
+		return new DefaultServiceInstance(instanceId, "test", "http://test.test", 9080, false);
+	}
+
+}

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

@@ -41,23 +41,23 @@ import static org.mockito.Mockito.when;
  */
 class ZonePreferenceServiceInstanceListSupplierTests {
 
-	private DiscoveryClientServiceInstanceListSupplier delegate = mock(
+	private final DiscoveryClientServiceInstanceListSupplier delegate = mock(
 			DiscoveryClientServiceInstanceListSupplier.class);
 
-	private LoadBalancerZoneConfig zoneConfig = new LoadBalancerZoneConfig(null);
+	private final LoadBalancerZoneConfig zoneConfig = new LoadBalancerZoneConfig(null);
 
-	private ZonePreferenceServiceInstanceListSupplier supplier = new ZonePreferenceServiceInstanceListSupplier(delegate,
-			zoneConfig);
+	private final ZonePreferenceServiceInstanceListSupplier supplier = new ZonePreferenceServiceInstanceListSupplier(
+			delegate, zoneConfig);
 
-	private ServiceInstance first = serviceInstance("test-1", buildZoneMetadata("zone1"));
+	private final ServiceInstance first = serviceInstance("test-1", buildZoneMetadata("zone1"));
 
-	private ServiceInstance second = serviceInstance("test-2", buildZoneMetadata("zone1"));
+	private final ServiceInstance second = serviceInstance("test-2", buildZoneMetadata("zone1"));
 
-	private ServiceInstance third = serviceInstance("test-3", buildZoneMetadata("zone2"));
+	private final ServiceInstance third = serviceInstance("test-3", buildZoneMetadata("zone2"));
 
-	private ServiceInstance fourth = serviceInstance("test-4", buildZoneMetadata("zone3"));
+	private final ServiceInstance fourth = serviceInstance("test-4", buildZoneMetadata("zone3"));
 
-	private ServiceInstance fifth = serviceInstance("test-5", buildZoneMetadata(null));
+	private final ServiceInstance fifth = serviceInstance("test-5", buildZoneMetadata(null));
 
 	@Test
 	void shouldFilterInstancesByZone() {
@@ -103,7 +103,7 @@ class ZonePreferenceServiceInstanceListSupplierTests {
 	}
 
 	private DefaultServiceInstance serviceInstance(String instanceId, Map<String, String> metadata) {
-		return new DefaultServiceInstance("test", instanceId, "http://test.test", 9080, false, metadata);
+		return new DefaultServiceInstance(instanceId, "test", "http://test.test", 9080, false, metadata);
 	}
 
 	private Map<String, String> buildZoneMetadata(String zone) {