Add service instance cookie mechanism for blocking load-balanced requests. (#918)

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

@@ -1045,8 +1045,6 @@ For that functionality, it is useful to have the selected service instance (whic
 
 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.
 
-WARNING: Passing the selected service instance id as a cookie is currently only supported for the reactive `@LoadBalanced Webclient`, `ReactorLoadBalancerExchangeFilterFunction` and `RetryableLoadBalancerExchangeFilterFunction`.
-
 [[spring-cloud-loadbalancer-hints]]
 === Spring Cloud LoadBalancer Hints
 

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

@@ -0,0 +1,60 @@
+/*
+ * 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.client.loadbalancer;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.springframework.cloud.client.ServiceInstance;
+import org.springframework.http.HttpCookie;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.HttpRequest;
+import org.springframework.util.StringUtils;
+
+/**
+ * A {@link LoadBalancerRequestTransformer} that allows passing the {@code} instanceId) of
+ * the {@link ServiceInstance} selected by the {@link LoadBalancerClient} in a cookie.
+ *
+ * @author Olga Maciaszek-Sharma
+ * @since 3.0.2
+ */
+public class LoadBalancerServiceInstanceCookieTransformer implements LoadBalancerRequestTransformer {
+
+	private final LoadBalancerProperties.StickySession stickySessionProperties;
+
+	public LoadBalancerServiceInstanceCookieTransformer(LoadBalancerProperties.StickySession stickySessionProperties) {
+		this.stickySessionProperties = stickySessionProperties;
+	}
+
+	@Override
+	public HttpRequest transformRequest(HttpRequest request, ServiceInstance instance) {
+		if (instance == null) {
+			return request;
+		}
+		String instanceIdCookieName = stickySessionProperties.getInstanceIdCookieName();
+		if (!StringUtils.hasText(instanceIdCookieName)) {
+			return request;
+		}
+		HttpHeaders headers = request.getHeaders();
+		List<String> cookieHeaders = new ArrayList<>(request.getHeaders().getOrEmpty(HttpHeaders.COOKIE));
+		String serviceInstanceCookie = new HttpCookie(instanceIdCookieName, instance.getInstanceId()).toString();
+		cookieHeaders.add(serviceInstanceCookie);
+		headers.put(HttpHeaders.COOKIE, cookieHeaders);
+		return request;
+	}
+
+}

spring-cloud-commons/src/test/java/org/springframework/cloud/client/loadbalancer/LoadBalancerServiceInstanceCookieTransformerTests.java

@@ -0,0 +1,87 @@
+/*
+ * 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.client.loadbalancer;
+
+import org.junit.jupiter.api.Test;
+
+import org.springframework.cloud.client.DefaultServiceInstance;
+import org.springframework.cloud.client.ServiceInstance;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.HttpRequest;
+import org.springframework.mock.http.client.MockClientHttpRequest;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Tests for {@link LoadBalancerServiceInstanceCookieTransformer}.
+ *
+ * @author Olga Maciaszek-Sharma
+ */
+class LoadBalancerServiceInstanceCookieTransformerTests {
+
+	LoadBalancerProperties.StickySession stickySessionProperties = new LoadBalancerProperties().getStickySession();
+
+	LoadBalancerServiceInstanceCookieTransformer transformer = new LoadBalancerServiceInstanceCookieTransformer(
+			stickySessionProperties);
+
+	ServiceInstance serviceInstance = new DefaultServiceInstance("test-01", "test", "host", 8080, false);
+
+	HttpRequest request = new MockClientHttpRequest();
+
+	@Test
+	void shouldAddServiceInstanceCookieHeader() {
+		HttpRequest newRequest = transformer.transformRequest(request, serviceInstance);
+
+		assertThat(newRequest.getHeaders().get(HttpHeaders.COOKIE)).hasSize(1);
+		assertThat(newRequest.getHeaders().get(HttpHeaders.COOKIE)).containsExactly("sc-lb-instance-id=test-01");
+	}
+
+	@Test
+	void shouldAppendServiceInstanceCookieHeaderIfCookiesPresent() {
+		request.getHeaders().add(HttpHeaders.COOKIE, "testCookieName=testCookieValue");
+
+		HttpRequest newRequest = transformer.transformRequest(request, serviceInstance);
+
+		assertThat(newRequest.getHeaders().get(HttpHeaders.COOKIE)).containsExactly("testCookieName=testCookieValue",
+				"sc-lb-instance-id=test-01");
+	}
+
+	@Test
+	void shouldReturnPassedRequestWhenNoServiceInstance() {
+		HttpRequest newRequest = transformer.transformRequest(request, null);
+
+		assertThat(newRequest.getHeaders()).doesNotContainKey(HttpHeaders.COOKIE);
+	}
+
+	@Test
+	void shouldReturnPassedRequestWhenNullServiceInstanceCookieName() {
+		stickySessionProperties.setInstanceIdCookieName(null);
+		HttpRequest newRequest = transformer.transformRequest(request, serviceInstance);
+
+		assertThat(newRequest.getHeaders()).doesNotContainKey(HttpHeaders.COOKIE);
+	}
+
+	@Test
+	void shouldReturnPassedRequestWhenEmptyServiceInstanceCookieName() {
+		stickySessionProperties.setInstanceIdCookieName("");
+		HttpRequest newRequest = transformer.transformRequest(request, serviceInstance);
+
+		assertThat(newRequest.getHeaders()).doesNotContainKey(HttpHeaders.COOKIE);
+
+	}
+
+}

spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/config/BlockingLoadBalancerClientAutoConfiguration.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2020 the original author or authors.
+ * 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.
@@ -21,11 +21,13 @@ import org.springframework.boot.autoconfigure.AutoConfigureBefore;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.cloud.client.loadbalancer.AsyncLoadBalancerAutoConfiguration;
 import org.springframework.cloud.client.loadbalancer.LoadBalancedRetryFactory;
 import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
 import org.springframework.cloud.client.loadbalancer.LoadBalancerProperties;
+import org.springframework.cloud.client.loadbalancer.LoadBalancerServiceInstanceCookieTransformer;
 import org.springframework.cloud.loadbalancer.annotation.LoadBalancerClients;
 import org.springframework.cloud.loadbalancer.blocking.client.BlockingLoadBalancerClient;
 import org.springframework.cloud.loadbalancer.blocking.retry.BlockingLoadBalancedRetryFactory;
@@ -57,6 +59,15 @@ public class BlockingLoadBalancerClientAutoConfiguration {
 		return new BlockingLoadBalancerClient(loadBalancerClientFactory, properties);
 	}
 
+	@Bean
+	@ConditionalOnProperty(value = "spring.cloud.loadbalancer.sticky-session.add-service-instance-cookie",
+			havingValue = "true")
+	@ConditionalOnMissingBean(LoadBalancerServiceInstanceCookieTransformer.class)
+	public LoadBalancerServiceInstanceCookieTransformer loadBalancerServiceInstanceCookieTransformer(
+			LoadBalancerProperties properties) {
+		return new LoadBalancerServiceInstanceCookieTransformer(properties.getStickySession());
+	}
+
 	@Configuration
 	@ConditionalOnClass(RetryTemplate.class)
 	@EnableConfigurationProperties(LoadBalancerProperties.class)