DEV: make the retry context aware of the policies that have a maximum number of attempts set

Author: e-ivaldi

src/main/java/org/springframework/retry/RetryContext.java

@@ -61,6 +61,13 @@ public interface RetryContext extends AttributeAccessor {
 	 */
 	String NO_RECOVERY = "context.no-recovery";
 
+	/**
+	 * Retry context attribute that represent the maximum number of attempts for policies
+	 * that provide a maximum number of attempts before failure. For other policies the
+	 * value returned is {@link RetryPolicy#NO_MAXIMUM_ATTEMPTS_SET}
+	 */
+	String MAX_ATTEMPTS = "context.max-attempts";
+
 	/**
 	 * Signal to the framework that no more attempts should be made to try or retry the
 	 * current {@link RetryCallback}.

src/main/java/org/springframework/retry/RetryPolicy.java

@@ -30,6 +30,12 @@
  */
 public interface RetryPolicy extends Serializable {
 
+	/**
+	 * The value returned by {@link RetryPolicy#getMaxAttempts()} when the policy doesn't
+	 * provide a maximum number of attempts before failure
+	 */
+	int NO_MAXIMUM_ATTEMPTS_SET = -1;
+
 	/**
 	 * @param context the current retry status
 	 * @return true if the operation can proceed
@@ -59,4 +65,14 @@ public interface RetryPolicy extends Serializable {
 	 */
 	void registerThrowable(RetryContext context, Throwable throwable);
 
+	/**
+	 * Called to understand if the policy has a fixed number of maximum attempts before
+	 * failure
+	 * @return -1 if the policy doesn't provide a fixed number of maximum attempts before
+	 * failure, the number of maximum attempts before failure otherwise
+	 */
+	default int getMaxAttempts() {
+		return NO_MAXIMUM_ATTEMPTS_SET;
+	}
+
 }

src/main/java/org/springframework/retry/policy/CompositeRetryPolicy.java

@@ -145,6 +145,22 @@ public void registerThrowable(RetryContext context, Throwable throwable) {
 		((RetryContextSupport) context).registerThrowable(throwable);
 	}
 
+	/**
+	 * @return the lower 'maximum number of attempts before failure' between all policies
+	 * that have a 'maximum number of attempts before failure' set, if at least one is
+	 * present among the policies, return {@link RetryPolicy#NO_MAXIMUM_ATTEMPTS_SET}
+	 * otherwise
+	 */
+	@Override
+	public int getMaxAttempts() {
+		return Arrays.stream(policies)
+			.map(RetryPolicy::getMaxAttempts)
+			.filter(maxAttempts -> maxAttempts != NO_MAXIMUM_ATTEMPTS_SET)
+			.sorted()
+			.findFirst()
+			.orElse(NO_MAXIMUM_ATTEMPTS_SET);
+	}
+
 	private static class CompositeRetryContext extends RetryContextSupport {
 
 		RetryContext[] contexts;

src/main/java/org/springframework/retry/policy/MaxAttemptsRetryPolicy.java

@@ -74,6 +74,7 @@ public void setMaxAttempts(int maxAttempts) {
 	 * The maximum number of attempts before failure.
 	 * @return the maximum number of attempts
 	 */
+	@Override
 	public int getMaxAttempts() {
 		return this.maxAttempts;
 	}

src/main/java/org/springframework/retry/policy/SimpleRetryPolicy.java

@@ -194,6 +194,7 @@ public void maxAttemptsSupplier(Supplier<Integer> maxAttemptsSupplier) {
 	 * The maximum number of attempts before failure.
 	 * @return the maximum number of attempts
 	 */
+	@Override
 	public int getMaxAttempts() {
 		if (this.maxAttemptsSupplier != null) {
 			return this.maxAttemptsSupplier.get();

src/main/java/org/springframework/retry/support/RetryTemplate.java

@@ -296,6 +296,10 @@ protected <T, E extends Throwable> T doExecute(RetryCallback<T, E> retryCallback
 				throw new TerminatedRetryException("Retry terminated abnormally by interceptor before first attempt");
 			}
 
+			if (!context.hasAttribute(RetryContext.MAX_ATTEMPTS)) {
+				context.setAttribute(RetryContext.MAX_ATTEMPTS, retryPolicy.getMaxAttempts());
+			}
+
 			// Get or Start the backoff context...
 			BackOffContext backOffContext = null;
 			Object resource = context.getAttribute("backOffContext");

src/test/java/org/springframework/retry/annotation/EnableRetryTests.java

@@ -294,7 +294,7 @@ void runtimeExpressions() throws Exception {
 		ExpressionService service = context.getBean(ExpressionService.class);
 		service.service6();
 		RuntimeConfigs runtime = context.getBean(RuntimeConfigs.class);
-		verify(runtime, times(5)).getMaxAttempts();
+		verify(runtime, times(6)).getMaxAttempts();
 		verify(runtime, times(2)).getInitial();
 		verify(runtime, times(2)).getMax();
 		verify(runtime, times(2)).getMult();

src/test/java/org/springframework/retry/policy/CompositeRetryPolicyTests.java

@@ -160,4 +160,23 @@ public boolean canRetry(RetryContext context) {
 		assertThat(policy.canRetry(context)).isTrue();
 	}
 
+	@Test
+	public void testMaximumAttemptsForNonSuitablePolicies() {
+		CompositeRetryPolicy policy = new CompositeRetryPolicy();
+		policy.setOptimistic(true);
+		policy.setPolicies(new RetryPolicy[] { new NeverRetryPolicy(), new NeverRetryPolicy() });
+
+		assertThat(policy.getMaxAttempts()).isEqualTo(RetryPolicy.NO_MAXIMUM_ATTEMPTS_SET);
+	}
+
+	@Test
+	public void testMaximumAttemptsForSuitablePolicies() {
+		CompositeRetryPolicy policy = new CompositeRetryPolicy();
+		policy.setOptimistic(true);
+		policy.setPolicies(
+				new RetryPolicy[] { new SimpleRetryPolicy(6), new SimpleRetryPolicy(3), new SimpleRetryPolicy(4) });
+
+		assertThat(policy.getMaxAttempts()).isEqualTo(3);
+	}
+
 }

src/test/java/org/springframework/retry/support/RetryTemplateTests.java

@@ -26,11 +26,13 @@
 import org.springframework.retry.RetryCallback;
 import org.springframework.retry.RetryContext;
 import org.springframework.retry.RetryListener;
+import org.springframework.retry.RetryPolicy;
 import org.springframework.retry.TerminatedRetryException;
 import org.springframework.retry.backoff.BackOffContext;
 import org.springframework.retry.backoff.BackOffInterruptedException;
 import org.springframework.retry.backoff.BackOffPolicy;
 import org.springframework.retry.backoff.StatelessBackOffPolicy;
+import org.springframework.retry.policy.AlwaysRetryPolicy;
 import org.springframework.retry.policy.NeverRetryPolicy;
 import org.springframework.retry.policy.SimpleRetryPolicy;
 
@@ -333,6 +335,30 @@ public <T, E extends Throwable> void onSuccess(RetryContext context, RetryCallba
 		assertThat(callCount.get()).isEqualTo(2);
 	}
 
+	@Test
+	public void testContextForPolicyWithMaximumNumberOfAttempts() throws Throwable {
+		RetryTemplate retryTemplate = new RetryTemplate();
+		RetryPolicy retryPolicy = new SimpleRetryPolicy(2);
+		retryTemplate.setRetryPolicy(retryPolicy);
+
+		Integer result = retryTemplate.execute((RetryCallback<Integer, Throwable>) context -> (Integer) context
+			.getAttribute(RetryContext.MAX_ATTEMPTS), context -> RetryPolicy.NO_MAXIMUM_ATTEMPTS_SET);
+
+		assertThat(result).isEqualTo(2);
+	}
+
+	@Test
+	public void testContextForPolicyWithNoMaximumNumberOfAttempts() throws Throwable {
+		RetryTemplate retryTemplate = new RetryTemplate();
+		RetryPolicy retryPolicy = new AlwaysRetryPolicy();
+		retryTemplate.setRetryPolicy(retryPolicy);
+
+		Integer result = retryTemplate.execute((RetryCallback<Integer, Throwable>) context -> (Integer) context
+			.getAttribute(RetryContext.MAX_ATTEMPTS), context -> RetryPolicy.NO_MAXIMUM_ATTEMPTS_SET);
+
+		assertThat(result).isEqualTo(RetryPolicy.NO_MAXIMUM_ATTEMPTS_SET);
+	}
+
 	private static class MockRetryCallback implements RetryCallback<Object, Exception> {
 
 		private int attempts;