Call `get(Request request)` on delegates in WeightedServiceInstanceListSupplier.

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);
 		}
 

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<ServiceInstance> 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,
+		DelegateCreator creator = (context, delegate) -> {
-				weightFunction);
+			ReactiveLoadBalancer.Factory<ServiceInstance> loadBalancerClientFactory = context
+					.getBean(LoadBalancerClientFactory.class);
+			return new WeightedServiceInstanceListSupplier(delegate, weightFunction, loadBalancerClientFactory);
+		};
 		this.creators.add(creator);
 		return this;
 	}

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(delegate, WeightedServiceInstanceListSupplier::metadataWeightFunction);
-		this.weightFunction = 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<ServiceInstance> loadBalancerClientFactory) {
+		this(delegate, WeightedServiceInstanceListSupplier::metadataWeightFunction, loadBalancerClientFactory);
+	}
+
+	public WeightedServiceInstanceListSupplier(ServiceInstanceListSupplier delegate, WeightFunction weightFunction,
+			ReactiveLoadBalancer.Factory<ServiceInstance> loadBalancerClientFactory) {
+		super(delegate);
+		this.weightFunction = weightFunction;
+		callGetWithRequestOnDelegates = loadBalancerClientFactory.getProperties(getServiceId())
+				.isCallGetWithRequestOnDelegates();
+	}
+
 	@Override
 	public Flux<List<ServiceInstance>> get() {
 		return delegate.get().map(this::expandByWeight);
 	}
 
+	@Override
+	public Flux<List<ServiceInstance>> get(Request request) {
+		if (callGetWithRequestOnDelegates) {
+			return delegate.get(request).map(this::expandByWeight);
+		}
+		return get();
+	}
+
 	private List<ServiceInstance> expandByWeight(List<ServiceInstance> instances) {
 		if (instances.size() == 0) {
 			return instances;

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);
 				});

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<DefaultRequestContext> 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<ServiceInstance> serviceInstances = Objects.requireNonNull(supplier.get(request).blockFirst());
+		Map<String, Integer> 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<String, String> metadata) {
 		return new DefaultServiceInstance(instanceId, "test", "localhost", 8080, false, metadata);
 	}