Handle avoiding previous instance with ServiceInstanceListSupplier in place of LoadBalancer.

Author: OlgaMaciaszek

docs/src/main/asciidoc/_configprops.adoc

@@ -35,7 +35,7 @@
 |spring.cloud.loadbalancer.health-check.interval | 25s | Interval for rerunning the HealthCheck scheduler.
 |spring.cloud.loadbalancer.health-check.path |  | 
 |spring.cloud.loadbalancer.hint |  | Allows setting the value of <code>hint</code> that is passed on to the LoadBalancer request and can subsequently be used in {@link ReactiveLoadBalancer} implementations.
-|spring.cloud.loadbalancer.retry.avoidPreviousInstance.enabled | true | Instantiates AvoidPreviousInstanceRoundRobinLoadBalancer if Spring-Retry is in the classpath.
+|spring.cloud.loadbalancer.retry.avoidPreviousInstance | true | Enables wrapping ServiceInstanceListSupplier beans with `RetryAwareServiceInstanceListSupplier` if Spring-Retry is in the classpath.
 |spring.cloud.loadbalancer.retry.enabled | true | 
 |spring.cloud.loadbalancer.retry.max-retries-on-next-service-instance | 1 | Number of retries to be executed on the next <code>ServiceInstance</code>. A <code>ServiceInstance</code> is chosen before each retry call.
 |spring.cloud.loadbalancer.retry.max-retries-on-same-service-instance | 0 | Number of retries to be executed on the same <code>ServiceInstance</code>.

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

@@ -449,7 +449,7 @@ You can set:
 - `spring.cloud.loadbalancer.retry.maxRetriesOnNextServiceInstance` - indicates how many times a request should be retried a newly selected `ServiceInstance`
 - `spring.cloud.loadbalancer.retry.retryableStatusCodes` - the status codes on which to always retry a failed request.
 
-NOTE: For load-balanced retries, by default, we instantiate `AvoidPreviousInstanceRoundRobinLoadBalancer` to select a different instance from the one previously chosen, if available. You can switch back to using the basic `RoundRobinLoadBalancer` by setting the value of `spring.cloud.loadbalancer.retry.avoidPreviousInstance.enabled` to `false`.
+NOTE: For load-balanced retries, by default, we wrap the `ServiceInstanceListSupplier` bean with `RetryAwareServiceInstanceListSupplier` to select a different instance from the one previously chosen, if available. You can disable this behaviour by setting the value of `spring.cloud.loadbalancer.retry.avoidPreviousInstance` to `false`.
 
 ====
 [source,java,indent=0]

spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/annotation/LoadBalancerClientConfiguration.java

@@ -16,6 +16,7 @@
 
 package org.springframework.cloud.loadbalancer.annotation;
 
+import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.condition.AllNestedConditions;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -27,8 +28,8 @@
 import org.springframework.cloud.client.ServiceInstance;
 import org.springframework.cloud.client.discovery.DiscoveryClient;
 import org.springframework.cloud.client.discovery.ReactiveDiscoveryClient;
-import org.springframework.cloud.loadbalancer.core.AvoidPreviousInstanceRoundRobinLoadBalancer;
 import org.springframework.cloud.loadbalancer.core.ReactorLoadBalancer;
+import org.springframework.cloud.loadbalancer.core.RetryAwareServiceInstanceListSupplier;
 import org.springframework.cloud.loadbalancer.core.RoundRobinLoadBalancer;
 import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
 import org.springframework.cloud.loadbalancer.support.LoadBalancerClientFactory;
@@ -39,6 +40,7 @@
 import org.springframework.context.annotation.Primary;
 import org.springframework.core.annotation.Order;
 import org.springframework.core.env.Environment;
+import org.springframework.retry.support.RetryTemplate;
 
 /**
  * @author Spencer Gibb
@@ -60,17 +62,6 @@ public ReactorLoadBalancer<ServiceInstance> reactorServiceInstanceLoadBalancer(E
 				loadBalancerClientFactory.getLazyProvider(name, ServiceInstanceListSupplier.class), name);
 	}
 
-	@Bean
-	@ConditionalOnClass(name = "org.springframework.retry.support.RetryTemplate")
-	@Conditional(OnAvoidPreviousInstanceAndRetryEnabledCondition.class)
-	@Primary
-	public ReactorLoadBalancer<ServiceInstance> avoidRetryingOnSameServiceInstanceLoadBalancer(Environment environment,
-			LoadBalancerClientFactory loadBalancerClientFactory) {
-		String name = environment.getProperty(LoadBalancerClientFactory.PROPERTY_NAME);
-		return new AvoidPreviousInstanceRoundRobinLoadBalancer(
-				loadBalancerClientFactory.getLazyProvider(name, ServiceInstanceListSupplier.class), name);
-	}
-
 	@Configuration(proxyBeanMethods = false)
 	@ConditionalOnReactiveDiscoveryEnabled
 	@Order(REACTIVE_SERVICE_INSTANCE_SUPPLIER_ORDER)
@@ -145,6 +136,24 @@ public ServiceInstanceListSupplier healthCheckDiscoveryClientServiceInstanceList
 
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	@ConditionalOnBlockingDiscoveryEnabled
+	@ConditionalOnClass(RetryTemplate.class)
+	@Conditional(OnAvoidPreviousInstanceAndRetryEnabledCondition.class)
+	@AutoConfigureAfter(BlockingSupportConfiguration.class)
+	@ConditionalOnBean(ServiceInstanceListSupplier.class)
+	public static class BlockingRetryConfiguration {
+
+		@Bean
+		@ConditionalOnBean(DiscoveryClient.class)
+		@Primary
+		public ServiceInstanceListSupplier retryAwareDiscoveryClientServiceInstanceListSupplier(
+				ServiceInstanceListSupplier delegate) {
+			return new RetryAwareServiceInstanceListSupplier(delegate);
+		}
+
+	}
+
 	static final class OnAvoidPreviousInstanceAndRetryEnabledCondition extends AllNestedConditions {
 
 		private OnAvoidPreviousInstanceAndRetryEnabledCondition() {
@@ -157,8 +166,8 @@ static class LoadBalancerRetryEnabled {
 
 		}
 
-		@ConditionalOnProperty(value = "spring.cloud.loadbalancer.retry.avoidPreviousInstance.enabled",
-				havingValue = "true", matchIfMissing = true)
+		@ConditionalOnProperty(value = "spring.cloud.loadbalancer.retry.avoidPreviousInstance", havingValue = "true",
+				matchIfMissing = true)
 		static class AvoidPreviousInstanceEnabled {
 
 		}

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

@@ -0,0 +1,85 @@
+/*
+ * 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.ArrayList;
+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.Request;
+import org.springframework.cloud.client.loadbalancer.RetryableRequestContext;
+
+/**
+ * A {@link ServiceInstanceListSupplier} implementation that avoids picking the same
+ * service instance while retrying requests.
+ *
+ * @author Olga Maciaszek-Sharma
+ * @since 3.0.0
+ */
+public class RetryAwareServiceInstanceListSupplier extends DelegatingServiceInstanceListSupplier {
+
+	private final Log LOG = LogFactory.getLog(RetryAwareServiceInstanceListSupplier.class);
+
+	public RetryAwareServiceInstanceListSupplier(ServiceInstanceListSupplier delegate) {
+		super(delegate);
+	}
+
+	@Override
+	public String getServiceId() {
+		return delegate.getServiceId();
+	}
+
+	@Override
+	public Flux<List<ServiceInstance>> get(Request request) {
+		if (!(request.getContext() instanceof RetryableRequestContext)) {
+			return get();
+		}
+		RetryableRequestContext context = (RetryableRequestContext) request.getContext();
+		ServiceInstance previousServiceInstance = context.getPreviousServiceInstance();
+		if (previousServiceInstance == null) {
+			return get();
+		}
+		return get().map(instances -> filteredByPreviousInstance(instances, previousServiceInstance));
+	}
+
+	private List<ServiceInstance> filteredByPreviousInstance(List<ServiceInstance> instances,
+			ServiceInstance previousServiceInstance) {
+		List<ServiceInstance> filteredInstances = new ArrayList<>(instances);
+		if (previousServiceInstance != null) {
+			filteredInstances.remove(previousServiceInstance);
+		}
+		if (filteredInstances.size() > 0) {
+			return filteredInstances;
+		}
+		if (LOG.isWarnEnabled()) {
+			LOG.warn(String.format(
+					"No instances found after removing previously used service instance from the search (%s). Returning all found instances.",
+					previousServiceInstance));
+		}
+		return instances;
+	}
+
+	@Override
+	public Flux<List<ServiceInstance>> get() {
+		return delegate.get();
+	}
+
+}

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

@@ -78,7 +78,7 @@ public RoundRobinLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> servic
 	public Mono<Response<ServiceInstance>> choose(Request request) {
 		ServiceInstanceListSupplier supplier = serviceInstanceListSupplierProvider
 				.getIfAvailable(NoopServiceInstanceListSupplier::new);
-		return supplier.get().next().map(this::getInstanceResponse);
+		return supplier.get(request).next().map(this::getInstanceResponse);
 	}
 
 	Response<ServiceInstance> getInstanceResponse(List<ServiceInstance> instances) {

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

@@ -24,6 +24,7 @@
 
 import org.springframework.cloud.client.DefaultServiceInstance;
 import org.springframework.cloud.client.ServiceInstance;
+import org.springframework.cloud.client.loadbalancer.Request;
 import org.springframework.core.env.Environment;
 
 import static org.springframework.cloud.loadbalancer.support.LoadBalancerClientFactory.PROPERTY_NAME;
@@ -38,6 +39,10 @@ public interface ServiceInstanceListSupplier extends Supplier<Flux<List<ServiceI
 
 	String getServiceId();
 
+	default Flux<List<ServiceInstance>> get(Request request) {
+		return get();
+	}
+
 	static ServiceInstanceListSupplierBuilder builder() {
 		return new ServiceInstanceListSupplierBuilder();
 	}

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

@@ -172,6 +172,12 @@ public ServiceInstanceListSupplierBuilder withCaching() {
 		return this;
 	}
 
+	public ServiceInstanceListSupplierBuilder withRetryAwareness() {
+		DelegateCreator creator = (context, delegate) -> new RetryAwareServiceInstanceListSupplier(delegate);
+		creators.add(creator);
+		return this;
+	}
+
 	/**
 	 * Builds the {@link ServiceInstanceListSupplier} hierarchy.
 	 * @param context application context

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

@@ -24,8 +24,8 @@
 		},
 		{
 			"defaultValue": true,
-			"name": "spring.cloud.loadbalancer.retry.avoidPreviousInstance.enabled",
-			"description": "Instantiates AvoidPreviousInstanceRoundRobinLoadBalancer if Spring-Retry is in the classpath.",
+			"name": "spring.cloud.loadbalancer.retry.avoidPreviousInstance",
+			"description": "Enables wrapping ServiceInstanceListSupplier beans with `RetryAwareServiceInstanceListSupplier` if Spring-Retry is in the classpath.",
 			"type": "java.lang.Boolean"
 		}
 	]