Request based sticky session

docs/src/main/asciidoc/_configprops.adoc

@@ -46,6 +46,8 @@
 |spring.cloud.loadbalancer.retry.retryable-status-codes |  | A {@link Set} of status codes that should trigger a retry.
 |spring.cloud.loadbalancer.ribbon.enabled | `true` | Causes `RibbonLoadBalancerClient` to be used by default.
 |spring.cloud.loadbalancer.service-discovery.timeout |  | String representation of Duration of the timeout for calls to service discovery.
+|spring.cloud.loadbalancer.sticky-session.add-service-instance-cookie | `false` | Indicates whether a cookie with the newly selected instance should be added by SC LoadBalancer.
+|spring.cloud.loadbalancer.sticky-session.instance-id-cookie-name | `sc-lb-instance-id` | The name of the cookie holding the preferred instance id.
 |spring.cloud.loadbalancer.zone |  | Spring Cloud LoadBalancer zone.
 |spring.cloud.refresh.enabled | `true` | Enables autoconfiguration for the refresh scope and associated features.
 |spring.cloud.refresh.extra-refreshable | `true` | Additional class names for beans to post process into refresh scope.

docs/src/main/asciidoc/spring-cloud-commons.adoc

@@ -993,6 +993,32 @@ public class CustomLoadBalancerConfiguration {
 
 TIP: This is also a replacement for Zookeeper `StickyRule`.
 
+=== Request-based Sticky Session for LoadBalancer
+
+You can set up the LoadBalancer in such a way that it prefers the instance with `instanceId` provided in a request cookie. We currently support this if the request is being passed to the LoadBalancer through either `ClientRequestContext` or `ServerHttpRequestContext`, which are used by the SC LoadBalancer exchange filter functions and filters.
+
+For that, you need to use the `RequestBasedStickySessionServiceInstanceListSupplier`. You can configure it either by setting the value of `spring.cloud.loadbalancer.configurations` to `request-based-sticky-session` or by providing your own `ServiceInstanceListSupplier` bean -- for example:
+
+[[health-check-based-custom-loadbalancer-configuration]]
+[source,java,indent=0]
+----
+public class CustomLoadBalancerConfiguration {
+
+	@Bean
+	public ServiceInstanceListSupplier discoveryClientServiceInstanceListSupplier(
+			ConfigurableApplicationContext context) {
+		return ServiceInstanceListSupplier.builder()
+					.withDiscoveryClient()
+					.withRequestBasedStickySession()
+					.build(context);
+	    }
+	}
+----
+
+For that functionality, it is useful to have the selected service instance (which can be different from the one in the original request cookie if that one is not available) to be updated before sending the request forward. To do that, set the value of `spring.cloud.loadbalancer.sticky-session.add-service-instance-cookie` to `true`.
+
+By default, the name of the cookie is `sc-lb-instance-id`. You can modify it by changing the value of the `spring.cloud.loadbalancer.instance-id-cookie-name` property.
+
 [[spring-cloud-loadbalancer-hints]]
 === Spring Cloud LoadBalancer Hints
 

spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ExchangeFilterFunctionUtils.java

@@ -19,6 +19,8 @@
 import java.net.URI;
 import java.util.Map;
 
+import org.springframework.cloud.client.ServiceInstance;
+import org.springframework.cloud.client.loadbalancer.LoadBalancerUriTools;
 import org.springframework.web.reactive.function.client.ClientRequest;
 import org.springframework.web.reactive.function.client.ExchangeFilterFunction;
 
@@ -40,10 +42,17 @@ static String getHint(String serviceId, Map<String, String> hints) {
 		return hintPropertyValue != null ? hintPropertyValue : defaultHint;
 	}
 
-	static ClientRequest buildClientRequest(ClientRequest request, URI uri) {
-		return ClientRequest.create(request.method(), uri).headers(headers -> headers.addAll(request.headers()))
-				.cookies(cookies -> cookies.addAll(request.cookies()))
-				.attributes(attributes -> attributes.putAll(request.attributes())).body(request.body()).build();
+	static ClientRequest buildClientRequest(ClientRequest request, ServiceInstance serviceInstance,
+			String instanceIdCookieName, boolean addServiceInstanceCookie) {
+		URI originalUrl = request.url();
+		return ClientRequest.create(request.method(), LoadBalancerUriTools.reconstructURI(serviceInstance, originalUrl))
+				.headers(headers -> headers.addAll(request.headers())).cookies(cookies -> {
+					cookies.addAll(request.cookies());
+					if (!(instanceIdCookieName == null || instanceIdCookieName.length() == 0)
+							&& addServiceInstanceCookie) {
+						cookies.add(instanceIdCookieName, serviceInstance.getInstanceId());
+					}
+				}).attributes(attributes -> attributes.putAll(request.attributes())).body(request.body()).build();
 	}
 
 	static String serviceInstanceUnavailableMessage(String serviceId) {

spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/LoadBalancerProperties.java

@@ -53,6 +53,11 @@ public class LoadBalancerProperties {
 	 */
 	private Retry retry = new Retry();
 
+	/**
+	 * Properties for LoadBalancer sticky-session.
+	 */
+	private StickySession stickySession = new StickySession();
+
 	public HealthCheck getHealthCheck() {
 		return healthCheck;
 	}
@@ -77,6 +82,45 @@ public void setRetry(Retry retry) {
 		this.retry = retry;
 	}
 
+	public StickySession getStickySession() {
+		return stickySession;
+	}
+
+	public void setStickySession(StickySession stickySession) {
+		this.stickySession = stickySession;
+	}
+
+	public static class StickySession {
+
+		/**
+		 * The name of the cookie holding the preferred instance id.
+		 */
+		private String instanceIdCookieName = "sc-lb-instance-id";
+
+		/**
+		 * Indicates whether a cookie with the newly selected instance should be added by
+		 * SC LoadBalancer.
+		 */
+		private boolean addServiceInstanceCookie = false;
+
+		public String getInstanceIdCookieName() {
+			return instanceIdCookieName;
+		}
+
+		public void setInstanceIdCookieName(String instanceIdCookieName) {
+			this.instanceIdCookieName = instanceIdCookieName;
+		}
+
+		public boolean isAddServiceInstanceCookie() {
+			return addServiceInstanceCookie;
+		}
+
+		public void setAddServiceInstanceCookie(boolean addServiceInstanceCookie) {
+			this.addServiceInstanceCookie = addServiceInstanceCookie;
+		}
+
+	}
+
 	public static class HealthCheck {
 
 		/**

spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ReactorLoadBalancerExchangeFilterFunction.java

@@ -38,7 +38,6 @@
 import org.springframework.web.reactive.function.client.ExchangeFilterFunction;
 import org.springframework.web.reactive.function.client.ExchangeFunction;
 
-import static org.springframework.cloud.client.loadbalancer.LoadBalancerUriTools.reconstructURI;
 import static org.springframework.cloud.client.loadbalancer.reactive.ExchangeFilterFunctionUtils.buildClientRequest;
 import static org.springframework.cloud.client.loadbalancer.reactive.ExchangeFilterFunctionUtils.getHint;
 import static org.springframework.cloud.client.loadbalancer.reactive.ExchangeFilterFunctionUtils.serviceInstanceUnavailableMessage;
@@ -101,7 +100,10 @@ public Mono<ClientResponse> filter(ClientRequest clientRequest, ExchangeFunction
 				LOG.debug(String.format("LoadBalancer has retrieved the instance for service %s: %s", serviceId,
 						instance.getUri()));
 			}
-			ClientRequest newRequest = buildClientRequest(clientRequest, reconstructURI(instance, originalUrl));
+			LoadBalancerProperties.StickySession stickySessionProperties = properties.getStickySession();
+			ClientRequest newRequest = buildClientRequest(clientRequest, instance,
+					stickySessionProperties.getInstanceIdCookieName(),
+					stickySessionProperties.isAddServiceInstanceCookie());
 			return next.exchange(newRequest)
 					.doOnError(throwable -> supportedLifecycleProcessors.forEach(
 							lifecycle -> lifecycle.onComplete(new CompletionContext<ClientResponse, ServiceInstance>(

spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/RetryableLoadBalancerExchangeFilterFunction.java

@@ -46,7 +46,6 @@
 import org.springframework.web.reactive.function.client.ExchangeFilterFunction;
 import org.springframework.web.reactive.function.client.ExchangeFunction;
 
-import static org.springframework.cloud.client.loadbalancer.LoadBalancerUriTools.reconstructURI;
 import static org.springframework.cloud.client.loadbalancer.reactive.ExchangeFilterFunctionUtils.buildClientRequest;
 import static org.springframework.cloud.client.loadbalancer.reactive.ExchangeFilterFunctionUtils.getHint;
 import static org.springframework.cloud.client.loadbalancer.reactive.ExchangeFilterFunctionUtils.serviceInstanceUnavailableMessage;
@@ -123,14 +122,17 @@ public Mono<ClientResponse> filter(ClientRequest clientRequest, ExchangeFunction
 				LOG.debug(String.format("LoadBalancer has retrieved the instance for service %s: %s", serviceId,
 						instance.getUri()));
 			}
-			ClientRequest newRequest = buildClientRequest(clientRequest, reconstructURI(instance, originalUrl));
+			LoadBalancerProperties.StickySession stickySessionProperties = properties.getStickySession();
+			ClientRequest newRequest = buildClientRequest(clientRequest, instance,
+					stickySessionProperties.getInstanceIdCookieName(),
+					stickySessionProperties.isAddServiceInstanceCookie());
 			return next.exchange(newRequest)
 					.doOnError(throwable -> supportedLifecycleProcessors.forEach(
 							lifecycle -> lifecycle.onComplete(new CompletionContext<ClientResponse, ServiceInstance>(
 									CompletionContext.Status.FAILED, throwable, lbResponse))))
 					.doOnSuccess(clientResponse -> supportedLifecycleProcessors.forEach(
-							lifecycle -> lifecycle.onComplete(new CompletionContext<ClientResponse, ServiceInstance>(
-									CompletionContext.Status.SUCCESS, lbResponse, clientResponse))))
+							lifecycle -> lifecycle.onComplete(new CompletionContext<>(CompletionContext.Status.SUCCESS,
+									lbResponse, clientResponse))))
 					.map(clientResponse -> {
 						loadBalancerRetryContext.setClientResponse(clientResponse);
 						if (shouldRetrySameServiceInstance(loadBalancerRetryContext)) {

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

@@ -98,6 +98,17 @@ public ServiceInstanceListSupplier healthCheckDiscoveryClientServiceInstanceList
 			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 ServiceInstanceListSupplier healthCheckDiscoveryClientServiceInstanceList
 					.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 ServiceInstanceListSupplierBuilder withZonePreference() {
 		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.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 @@ void shouldInstantiateHealthCheckServiceInstanceListSupplier() {
 				});
 	}
 
+	@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 -> {