1 /*
2  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License").
5  * You may not use this file except in compliance with the License.
6  * A copy of the License is located at
7  *
8  *  http://aws.amazon.com/apache2.0
9  *
10  * or in the "license" file accompanying this file. This file is distributed
11  * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12  * express or implied. See the License for the specific language governing
13  * permissions and limitations under the License.
14  */
15
16 package software.amazon.awssdk.core.retry.conditions;
17
18 import java.util.Arrays;
19 import java.util.HashSet;
20 import java.util.Set;
21 import java.util.function.Predicate;
22 import java.util.stream.Collectors;
23 import software.amazon.awssdk.annotations.SdkPublicApi;
24 import software.amazon.awssdk.core.exception.SdkException;
25 import software.amazon.awssdk.core.retry.RetryPolicyContext;
26 import software.amazon.awssdk.utils.ToString;
27
28 /**
29  * Retry condition implementation that retries if the exception or the cause of the exception matches the classes defined.
30  */
31 @SdkPublicApi
32 public final class RetryOnExceptionsCondition implements RetryCondition {
33
34     private final Set<Class<? extends Exception>> exceptionsToRetryOn;
35
36     /**
37      * @param exceptionsToRetryOn Exception classes to retry on.
38      */
39     private RetryOnExceptionsCondition(Set<Class<? extends Exception>> exceptionsToRetryOn) {
40         this.exceptionsToRetryOn = new HashSet<>(exceptionsToRetryOn);
41     }
42
43     /**
44      * @param context Context about the state of the last request and information about the number of requests made.
45      * @return True if the exception class matches one of the whitelisted exceptions or if the cause of the exception matches the
46      *     whitelisted exception.
47      */
48     @Override
49     public boolean shouldRetry(RetryPolicyContext context) {
50
51         SdkException exception = context.exception();
52         if (exception == null) {
53             return false;
54         }
55
56         Predicate<Class<? extends Exception>> isRetryableException =
57             ex -> ex.isAssignableFrom(exception.getClass());
58
59         Predicate<Class<? extends Exception>> hasRetrableCause =
60             ex -> exception.getCause() != null && ex.isAssignableFrom(exception.getCause().getClass());
61
62         return exceptionsToRetryOn.stream().anyMatch(isRetryableException.or(hasRetrableCause));
63     }
64
65     /**
66      * @param exceptionsToRetryOn Exception classes to retry on.
67      */
68     public static RetryOnExceptionsCondition create(Set<Class<? extends Exception>> exceptionsToRetryOn) {
69         return new RetryOnExceptionsCondition(exceptionsToRetryOn);
70     }
71
72     /**
73      * @param exceptionsToRetryOn Exception classes to retry on.
74      */
75     public static RetryOnExceptionsCondition create(Class<? extends Exception>... exceptionsToRetryOn) {
76         return new RetryOnExceptionsCondition(Arrays.stream(exceptionsToRetryOn).collect(Collectors.toSet()));
77     }
78
79     @Override
80     public boolean equals(Object o) {
81         if (this == o) {
82             return true;
83         }
84         if (o == null || getClass() != o.getClass()) {
85             return false;
86         }
87
88         RetryOnExceptionsCondition that = (RetryOnExceptionsCondition) o;
89
90         return exceptionsToRetryOn.equals(that.exceptionsToRetryOn);
91     }
92
93     @Override
94     public int hashCode() {
95         return exceptionsToRetryOn.hashCode();
96     }
97
98     @Override
99     public String toString() {
100         return ToString.builder("RetryOnExceptionsCondition")
101                        .add("exceptionsToRetryOn", exceptionsToRetryOn)
102                        .build();
103     }
104 }
105