Move LoadBalancerServiceInstanceCookieTransformer to loadbalancer module. (#920)

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

@@ -27,10 +27,10 @@ import org.springframework.cloud.client.loadbalancer.AsyncLoadBalancerAutoConfig
 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;
+import org.springframework.cloud.loadbalancer.core.LoadBalancerServiceInstanceCookieTransformer;
 import org.springframework.cloud.loadbalancer.support.LoadBalancerClientFactory;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;

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

@@ -0,0 +1,63 @@
+/*
+ * 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 org.springframework.cloud.client.ServiceInstance;
+import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
+import org.springframework.cloud.client.loadbalancer.LoadBalancerProperties;
+import org.springframework.cloud.client.loadbalancer.LoadBalancerRequestTransformer;
+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-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/core/LoadBalancerServiceInstanceCookieTransformerTests.java

@@ -0,0 +1,88 @@
+/*
+ * 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 org.junit.jupiter.api.Test;
+
+import org.springframework.cloud.client.DefaultServiceInstance;
+import org.springframework.cloud.client.ServiceInstance;
+import org.springframework.cloud.client.loadbalancer.LoadBalancerProperties;
+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);
+
+	}
+
+}