diff --git a/docs/src/main/asciidoc/_configprops.adoc b/docs/src/main/asciidoc/_configprops.adoc index 71794eac7..0c6f1d88a 100644 --- a/docs/src/main/asciidoc/_configprops.adoc +++ b/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. diff --git a/docs/src/main/asciidoc/spring-cloud-commons.adoc b/docs/src/main/asciidoc/spring-cloud-commons.adoc index 1c5ef7c02..b864b1616 100644 --- a/docs/src/main/asciidoc/spring-cloud-commons.adoc +++ b/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 diff --git a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ExchangeFilterFunctionUtils.java b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ExchangeFilterFunctionUtils.java index 18e7124ec..b5f6134ea 100644 --- a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ExchangeFilterFunctionUtils.java +++ b/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 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) { diff --git a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/LoadBalancerProperties.java b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/LoadBalancerProperties.java index 218bf534a..1a36cc2e9 100644 --- a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/LoadBalancerProperties.java +++ b/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 { /** diff --git a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ReactorLoadBalancerExchangeFilterFunction.java b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ReactorLoadBalancerExchangeFilterFunction.java index ec2e150bb..7ce1df5f1 100644 --- a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/ReactorLoadBalancerExchangeFilterFunction.java +++ b/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 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( diff --git a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/RetryableLoadBalancerExchangeFilterFunction.java b/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/RetryableLoadBalancerExchangeFilterFunction.java index 8864e9733..abb7c23ee 100644 --- a/spring-cloud-commons/src/main/java/org/springframework/cloud/client/loadbalancer/reactive/RetryableLoadBalancerExchangeFilterFunction.java +++ b/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 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( CompletionContext.Status.FAILED, throwable, lbResponse)))) .doOnSuccess(clientResponse -> supportedLifecycleProcessors.forEach( - lifecycle -> lifecycle.onComplete(new CompletionContext( - CompletionContext.Status.SUCCESS, lbResponse, clientResponse)))) + lifecycle -> lifecycle.onComplete(new CompletionContext<>(CompletionContext.Status.SUCCESS, + lbResponse, clientResponse)))) .map(clientResponse -> { loadBalancerRetryContext.setClientResponse(clientResponse); if (shouldRetrySameServiceInstance(loadBalancerRetryContext)) { diff --git a/spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/annotation/LoadBalancerClientConfiguration.java b/spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/annotation/LoadBalancerClientConfiguration.java index fb3fd9599..61481e00a 100644 --- a/spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/annotation/LoadBalancerClientConfiguration.java +++ b/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 diff --git a/spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/core/RequestBasedStickySessionServiceInstanceListSupplier.java b/spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/core/RequestBasedStickySessionServiceInstanceListSupplier.java new file mode 100644 index 000000000..4022b00a6 --- /dev/null +++ b/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> get() { + return delegate.get(); + } + + @SuppressWarnings("rawtypes") + @Override + public Flux> 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 selectInstance(List 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; + } + +} diff --git a/spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/core/ServiceInstanceListSupplierBuilder.java b/spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/core/ServiceInstanceListSupplierBuilder.java index 7fd327f05..a1fe379cb 100644 --- a/spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/core/ServiceInstanceListSupplierBuilder.java +++ b/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 diff --git a/spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/annotation/LoadBalancerClientConfigurationTests.java b/spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/annotation/LoadBalancerClientConfigurationTests.java index 1379b62e9..6d8becdc0 100644 --- a/spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/annotation/LoadBalancerClientConfigurationTests.java +++ b/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 -> { diff --git a/spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/core/RequestBasedStickySessionServiceInstanceListSupplierTests.java b/spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/core/RequestBasedStickySessionServiceInstanceListSupplierTests.java new file mode 100644 index 000000000..3a4df6634 --- /dev/null +++ b/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 request = new DefaultRequest<>(new ClientRequestContext(clientRequest)); + + List 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 request = new DefaultRequest<>(new ClientRequestContext(clientRequest)); + + List serviceInstances = supplier.get(request).blockFirst(); + + assertThat(serviceInstances).hasSize(3); + } + + @Test + void shouldReturnAllInstancesFromDelegateIfClientRequestHasNoCookie() { + when(clientRequest.cookies()).thenReturn(new HttpHeaders()); + Request request = new DefaultRequest<>(new ClientRequestContext(clientRequest)); + + List serviceInstances = supplier.get(request).blockFirst(); + + assertThat(serviceInstances).hasSize(3); + } + + @Test + void shouldReturnInstanceBasedOnCookieFromServerHttpRequest() { + MultiValueMap cookies = new LinkedMultiValueMap<>(); + cookies.add(properties.getStickySession().getInstanceIdCookieName(), + new HttpCookie(properties.getStickySession().getInstanceIdCookieName(), "test-1")); + when(serverHttpRequest.getCookies()).thenReturn(cookies); + Request request = new DefaultRequest<>( + new ServerHttpRequestContext(serverHttpRequest)); + + List serviceInstances = supplier.get(request).blockFirst(); + + assertThat(serviceInstances).hasSize(1); + assertThat(serviceInstances.get(0).getInstanceId()).isEqualTo("test-1"); + } + + @Test + void shouldReturnAllDelegateInstancesIfInstanceBasedOnCookieFromServerHttpRequestNotFound() { + MultiValueMap cookies = new LinkedMultiValueMap<>(); + cookies.add(properties.getStickySession().getInstanceIdCookieName(), + new HttpCookie(properties.getStickySession().getInstanceIdCookieName(), "test-4")); + when(serverHttpRequest.getCookies()).thenReturn(cookies); + Request request = new DefaultRequest<>( + new ServerHttpRequestContext(serverHttpRequest)); + + List serviceInstances = supplier.get(request).blockFirst(); + + assertThat(serviceInstances).hasSize(3); + } + + @Test + void shouldReturnAllInstancesFromDelegateIfServerHttpRequestHasNoCookie() { + MultiValueMap cookies = new LinkedMultiValueMap<>(); + when(serverHttpRequest.getCookies()).thenReturn(cookies); + Request request = new DefaultRequest<>( + new ServerHttpRequestContext(serverHttpRequest)); + + List serviceInstances = supplier.get(request).blockFirst(); + + assertThat(serviceInstances).hasSize(3); + } + + @Test + void shouldReturnAllInstancesFromDelegateIfNotSupportedRequestContext() { + Request request = new DefaultRequest<>(new DefaultRequestContext(clientRequest)); + + List serviceInstances = supplier.get(request).blockFirst(); + + assertThat(serviceInstances).hasSize(3); + } + + private DefaultServiceInstance serviceInstance(String instanceId) { + return new DefaultServiceInstance(instanceId, "test", "http://test.test", 9080, false); + } + +} diff --git a/spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/core/ZonePreferenceServiceInstanceListSupplierTests.java b/spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/core/ZonePreferenceServiceInstanceListSupplierTests.java index c6ff15639..4fe7d39f6 100644 --- a/spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/core/ZonePreferenceServiceInstanceListSupplierTests.java +++ b/spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/core/ZonePreferenceServiceInstanceListSupplierTests.java @@ -41,23 +41,23 @@ */ 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 @@ void shouldReturnAllInstancesIfNoZone() { } private DefaultServiceInstance serviceInstance(String instanceId, Map 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 buildZoneMetadata(String zone) {