1515 */
1616package net .jodah .failsafe .spi ;
1717
18- import java .util .HashMap ;
18+ import java .util .Collections ;
1919import java .util .Iterator ;
2020import java .util .Map ;
2121import java .util .Map .Entry ;
22+ import java .util .TreeMap ;
2223import java .util .concurrent .CancellationException ;
2324import java .util .concurrent .CompletableFuture ;
2425import java .util .concurrent .Future ;
@@ -35,9 +36,9 @@ public class FailsafeFuture<R> extends CompletableFuture<R> {
3536
3637 // Mutable state guarded by "this"
3738
38- // The most recent executoin attempt
39+ // The most recent execution attempt
3940 private ExecutionInternal <R > newestExecution ;
40- // Functions to apply when this future is cancelled
41+ // Functions to apply when this future is cancelled for each policy index, in descending order
4142 private Map <Integer , BiConsumer <Boolean , ExecutionResult <R >>> cancelFunctions ;
4243 // Whether a cancel with interrupt has already occurred
4344 private boolean cancelledWithInterrupt ;
@@ -111,6 +112,10 @@ public synchronized void cancelDependencies(PolicyExecutor<R, ?> cancellingPolic
111112 int cancellingPolicyIndex =
112113 cancellingPolicyExecutor == null ? Integer .MAX_VALUE : cancellingPolicyExecutor .getPolicyIndex ();
113114 Iterator <Entry <Integer , BiConsumer <Boolean , ExecutionResult <R >>>> it = cancelFunctions .entrySet ().iterator ();
115+
116+ /* This iteration occurs in descending order to ensure that the {@code cancelResult} can be supplied to outer
117+ cancel functions before the inner supplier is cancelled, which would cause PolicyExecutors to complete with
118+ CancellationException rather than the expected {@code cancelResult}. */
114119 while (it .hasNext ()) {
115120 Map .Entry <Integer , BiConsumer <Boolean , ExecutionResult <R >>> entry = it .next ();
116121 if (cancellingPolicyIndex > entry .getKey ()) {
@@ -134,12 +139,13 @@ public synchronized void setExecution(ExecutionInternal<R> execution) {
134139
135140 /**
136141 * Sets a {@code cancelFn} to be called when a PolicyExecutor {@link #cancelDependencies(PolicyExecutor, boolean,
137- * ExecutionResult) cancels dependencies} or when this future is {@link #cancel(boolean) cancelled}.
142+ * ExecutionResult) cancels dependencies} with a policyIndex > the given {@code policyIndex}, or when this future is
143+ * {@link #cancel(boolean) cancelled}.
138144 */
139- public synchronized void setCancelFn (BiConsumer <Boolean , ExecutionResult <R >> cancelFn ) {
145+ public synchronized void setCancelFn (int policyIndex , BiConsumer <Boolean , ExecutionResult <R >> cancelFn ) {
140146 if (cancelFunctions == null )
141- cancelFunctions = new HashMap <>();
142- cancelFunctions .put (- 1 , cancelFn );
147+ cancelFunctions = new TreeMap <>(Collections . reverseOrder () );
148+ cancelFunctions .put (policyIndex , cancelFn );
143149 }
144150
145151 /**
@@ -150,7 +156,7 @@ public synchronized void setCancelFn(BiConsumer<Boolean, ExecutionResult<R>> can
150156 public synchronized void setCancelFn (PolicyExecutor <R , ?> policyExecutor ,
151157 BiConsumer <Boolean , ExecutionResult <R >> cancelFn ) {
152158 if (cancelFunctions == null )
153- cancelFunctions = new HashMap <>();
159+ cancelFunctions = new TreeMap <>(Collections . reverseOrder () );
154160 cancelFunctions .put (policyExecutor .getPolicyIndex (), cancelFn );
155161 }
156162
@@ -162,6 +168,6 @@ public synchronized void propagateCancellation(Future<R> future) {
162168 if (isCancelled ())
163169 future .cancel (cancelledWithInterrupt );
164170 else
165- setCancelFn ((mayInterrupt , cancelResult ) -> future .cancel (mayInterrupt ));
171+ setCancelFn (- 2 , (mayInterrupt , cancelResult ) -> future .cancel (mayInterrupt ));
166172 }
167173}
0 commit comments