From e4e56e93cfdb007df0f089647ae40ce6864d30b3 Mon Sep 17 00:00:00 2001 From: Olga MaciaszekSharma Date: Wed, 28 Jun 2023 14:30:15 +0200 Subject: [PATCH] Call `get(Request request)` on delegates in WeightedServiceInstanceListSupplier. --- .../LoadBalancerClientConfiguration.java | 4 +-- .../ServiceInstanceListSupplierBuilder.java | 13 ++++++-- .../WeightedServiceInstanceListSupplier.java | 31 +++++++++++++++-- .../LoadBalancerClientConfigurationTests.java | 8 ++--- ...ghtedServiceInstanceListSupplierTests.java | 33 ++++++++++++++++++- 5 files changed, 76 insertions(+), 13 deletions(-) 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 ea9cd772..f0aae50c 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 @@ -139,7 +139,7 @@ public class LoadBalancerClientConfiguration { @ConditionalOnMissingBean @Conditional(WeightedConfigurationCondition.class) public ServiceInstanceListSupplier weightedServiceInstanceListSupplier(ConfigurableApplicationContext context) { - return ServiceInstanceListSupplier.builder().withDiscoveryClient().withWeighted().withCaching() + return ServiceInstanceListSupplier.builder().withDiscoveryClient().withCaching().withWeighted() .build(context); } @@ -204,7 +204,7 @@ public class LoadBalancerClientConfiguration { @ConditionalOnMissingBean @Conditional(WeightedConfigurationCondition.class) public ServiceInstanceListSupplier weightedServiceInstanceListSupplier(ConfigurableApplicationContext context) { - return ServiceInstanceListSupplier.builder().withBlockingDiscoveryClient().withWeighted().withCaching() + return ServiceInstanceListSupplier.builder().withBlockingDiscoveryClient().withCaching().withWeighted() .build(context); } 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 7906638a..74ee2e38 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 @@ -116,7 +116,11 @@ public final class ServiceInstanceListSupplierBuilder { * @return the {@link ServiceInstanceListSupplierBuilder} object */ public ServiceInstanceListSupplierBuilder withWeighted() { - DelegateCreator creator = (context, delegate) -> new WeightedServiceInstanceListSupplier(delegate); + DelegateCreator creator = (context, delegate) -> { + ReactiveLoadBalancer.Factory loadBalancerClientFactory = context + .getBean(LoadBalancerClientFactory.class); + return new WeightedServiceInstanceListSupplier(delegate, loadBalancerClientFactory); + }; this.creators.add(creator); return this; } @@ -129,8 +133,11 @@ public final class ServiceInstanceListSupplierBuilder { * @return the {@link ServiceInstanceListSupplierBuilder} object */ public ServiceInstanceListSupplierBuilder withWeighted(WeightFunction weightFunction) { - DelegateCreator creator = (context, delegate) -> new WeightedServiceInstanceListSupplier(delegate, - weightFunction); + DelegateCreator creator = (context, delegate) -> { + ReactiveLoadBalancer.Factory loadBalancerClientFactory = context + .getBean(LoadBalancerClientFactory.class); + return new WeightedServiceInstanceListSupplier(delegate, weightFunction, loadBalancerClientFactory); + }; this.creators.add(creator); return this; } diff --git a/spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/core/WeightedServiceInstanceListSupplier.java b/spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/core/WeightedServiceInstanceListSupplier.java index 527e6173..172038b5 100644 --- a/spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/core/WeightedServiceInstanceListSupplier.java +++ b/spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/core/WeightedServiceInstanceListSupplier.java @@ -1,5 +1,5 @@ /* - * Copyright 2012-2022 the original author or authors. + * Copyright 2012-2023 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. @@ -24,12 +24,15 @@ import org.apache.commons.logging.LogFactory; import reactor.core.publisher.Flux; import org.springframework.cloud.client.ServiceInstance; +import org.springframework.cloud.client.loadbalancer.Request; +import org.springframework.cloud.client.loadbalancer.reactive.ReactiveLoadBalancer; /** * A {@link ServiceInstanceListSupplier} implementation that uses weights to expand the * instances provided by delegate. * * @author Zhuozhi Ji + * @author Olga Maciaszek-Sharma */ public class WeightedServiceInstanceListSupplier extends DelegatingServiceInstanceListSupplier { @@ -41,9 +44,10 @@ public class WeightedServiceInstanceListSupplier extends DelegatingServiceInstan private final WeightFunction weightFunction; + private boolean callGetWithRequestOnDelegates; + public WeightedServiceInstanceListSupplier(ServiceInstanceListSupplier delegate) { - super(delegate); - this.weightFunction = WeightedServiceInstanceListSupplier::metadataWeightFunction; + this(delegate, WeightedServiceInstanceListSupplier::metadataWeightFunction); } public WeightedServiceInstanceListSupplier(ServiceInstanceListSupplier delegate, WeightFunction weightFunction) { @@ -51,11 +55,32 @@ public class WeightedServiceInstanceListSupplier extends DelegatingServiceInstan this.weightFunction = weightFunction; } + public WeightedServiceInstanceListSupplier(ServiceInstanceListSupplier delegate, + ReactiveLoadBalancer.Factory loadBalancerClientFactory) { + this(delegate, WeightedServiceInstanceListSupplier::metadataWeightFunction, loadBalancerClientFactory); + } + + public WeightedServiceInstanceListSupplier(ServiceInstanceListSupplier delegate, WeightFunction weightFunction, + ReactiveLoadBalancer.Factory loadBalancerClientFactory) { + super(delegate); + this.weightFunction = weightFunction; + callGetWithRequestOnDelegates = loadBalancerClientFactory.getProperties(getServiceId()) + .isCallGetWithRequestOnDelegates(); + } + @Override public Flux> get() { return delegate.get().map(this::expandByWeight); } + @Override + public Flux> get(Request request) { + if (callGetWithRequestOnDelegates) { + return delegate.get(request).map(this::expandByWeight); + } + return get(); + } + private List expandByWeight(List instances) { if (instances.size() == 0) { return instances; 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 858d5ce4..eda9dc1c 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 @@ -120,10 +120,10 @@ class LoadBalancerClientConfigurationTests { reactiveDiscoveryClientRunner.withUserConfiguration(TestConfig.class) .withPropertyValues("spring.cloud.loadbalancer.configurations=weighted").run(context -> { ServiceInstanceListSupplier supplier = context.getBean(ServiceInstanceListSupplier.class); - then(supplier).isInstanceOf(CachingServiceInstanceListSupplier.class); + then(supplier).isInstanceOf(WeightedServiceInstanceListSupplier.class); ServiceInstanceListSupplier delegate = ((DelegatingServiceInstanceListSupplier) supplier) .getDelegate(); - then(delegate).isInstanceOf(WeightedServiceInstanceListSupplier.class); + then(delegate).isInstanceOf(CachingServiceInstanceListSupplier.class); ServiceInstanceListSupplier secondDelegate = ((DelegatingServiceInstanceListSupplier) delegate) .getDelegate(); then(secondDelegate).isInstanceOf(DiscoveryClientServiceInstanceListSupplier.class); @@ -207,10 +207,10 @@ class LoadBalancerClientConfigurationTests { blockingDiscoveryClientRunner.withUserConfiguration(RestTemplateTestConfig.class) .withPropertyValues("spring.cloud.loadbalancer.configurations=weighted").run(context -> { ServiceInstanceListSupplier supplier = context.getBean(ServiceInstanceListSupplier.class); - then(supplier).isInstanceOf(CachingServiceInstanceListSupplier.class); + then(supplier).isInstanceOf(WeightedServiceInstanceListSupplier.class); ServiceInstanceListSupplier delegate = ((DelegatingServiceInstanceListSupplier) supplier) .getDelegate(); - then(delegate).isInstanceOf(WeightedServiceInstanceListSupplier.class); + then(delegate).isInstanceOf(CachingServiceInstanceListSupplier.class); then(((DelegatingServiceInstanceListSupplier) delegate).getDelegate()) .isInstanceOf(DiscoveryClientServiceInstanceListSupplier.class); }); diff --git a/spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/core/WeightedServiceInstanceListSupplierTests.java b/spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/core/WeightedServiceInstanceListSupplierTests.java index 9bd925e2..58c00682 100644 --- a/spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/core/WeightedServiceInstanceListSupplierTests.java +++ b/spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/core/WeightedServiceInstanceListSupplierTests.java @@ -1,5 +1,5 @@ /* - * Copyright 2012-2022 the original author or authors. + * Copyright 2012-2023 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. @@ -29,9 +29,15 @@ import reactor.core.publisher.Flux; import org.springframework.cloud.client.DefaultServiceInstance; import org.springframework.cloud.client.ServiceInstance; +import org.springframework.cloud.client.loadbalancer.DefaultRequest; +import org.springframework.cloud.client.loadbalancer.DefaultRequestContext; +import org.springframework.cloud.client.loadbalancer.LoadBalancerProperties; +import org.springframework.cloud.client.loadbalancer.Request; +import org.springframework.cloud.loadbalancer.support.LoadBalancerClientFactory; import static java.util.stream.Collectors.summingInt; import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import static org.springframework.cloud.loadbalancer.core.WeightedServiceInstanceListSupplier.DEFAULT_WEIGHT; @@ -40,6 +46,7 @@ import static org.springframework.cloud.loadbalancer.core.WeightedServiceInstanc * Tests for {@link WeightedServiceInstanceListSupplier}. * * @author Zhuozhi Ji + * @author Olga Maciaszek-Sharma */ class WeightedServiceInstanceListSupplierTests { @@ -177,6 +184,30 @@ class WeightedServiceInstanceListSupplierTests { assertThat(counter).containsEntry("test-3", DEFAULT_WEIGHT); } + @Test + void shouldCallGetRequestOnDelegate() { + LoadBalancerClientFactory loadBalancerClientFactory = mock(LoadBalancerClientFactory.class); + LoadBalancerProperties properties = new LoadBalancerProperties(); + properties.setCallGetWithRequestOnDelegates(true); + when(loadBalancerClientFactory.getProperties(any())).thenReturn(properties); + ServiceInstance one = serviceInstance("test-1", Collections.emptyMap()); + ServiceInstance two = serviceInstance("test-2", Collections.emptyMap()); + ServiceInstance three = serviceInstance("test-3", buildWeightMetadata(3)); + Request request = new DefaultRequest<>(new DefaultRequestContext()); + + when(delegate.get()).thenReturn(Flux.just(Arrays.asList(one, two, three))); + when(delegate.get(request)).thenReturn(Flux.just(Arrays.asList(one, two))); + WeightedServiceInstanceListSupplier supplier = new WeightedServiceInstanceListSupplier(delegate, + loadBalancerClientFactory); + + List serviceInstances = Objects.requireNonNull(supplier.get(request).blockFirst()); + Map counter = serviceInstances.stream() + .collect(Collectors.groupingBy(ServiceInstance::getInstanceId, summingInt(e -> 1))); + assertThat(counter).containsEntry("test-1", DEFAULT_WEIGHT); + assertThat(counter).containsEntry("test-2", DEFAULT_WEIGHT); + assertThat(counter).doesNotContainEntry("test-3", 3); + } + private ServiceInstance serviceInstance(String instanceId, Map metadata) { return new DefaultServiceInstance(instanceId, "test", "localhost", 8080, false, metadata); }