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.internal.http.timers;
17
18 import java.time.Duration;
19 import java.util.Optional;
20 import java.util.concurrent.CompletableFuture;
21 import java.util.concurrent.ScheduledExecutorService;
22 import java.util.concurrent.ScheduledFuture;
23 import java.util.concurrent.TimeUnit;
24 import java.util.function.Supplier;
25 import software.amazon.awssdk.annotations.SdkInternalApi;
26 import software.amazon.awssdk.core.exception.SdkClientException;
27 import software.amazon.awssdk.utils.OptionalUtils;
28
29 @SdkInternalApi
30 public final class TimerUtils {
31
32     private TimerUtils() {
33     }
34
35     /**
36      * Schedule a {@link TimeoutTask} and exceptional completes a {@link CompletableFuture} with the provide exception
37      * if not otherwise completed before the given timeout.
38      *
39      * @param completableFuture the completableFuture to be timed
40      * @param timeoutExecutor the executor to execute the {@link TimeoutTask}
41      * @param exceptionSupplier the exception to thrown after timeout
42      * @param timeoutInMills the timeout in milliseconds.
43      * @param <T> the type of the {@link CompletableFuture}
44      * @return a {@link TimeoutTracker}
45      */
46     public static <T> TimeoutTracker timeAsyncTaskIfNeeded(CompletableFuture<T> completableFuture,
47                                                            ScheduledExecutorService timeoutExecutor,
48                                                            Supplier<SdkClientException> exceptionSupplier,
49                                                            long timeoutInMills) {
50         if (timeoutInMills <= 0) {
51             return NoOpTimeoutTracker.INSTANCE;
52         }
53
54         TimeoutTask timeoutTask = new AsyncTimeoutTask(completableFuture, exceptionSupplier);
55
56         ScheduledFuture<?> scheduledFuture =
57             timeoutExecutor.schedule(timeoutTask,
58                                      timeoutInMills,
59                                      TimeUnit.MILLISECONDS);
60         TimeoutTracker timeoutTracker = new ApiCallTimeoutTracker(timeoutTask, scheduledFuture);
61
62         completableFuture.whenComplete((o, t) -> timeoutTracker.cancel());
63
64         return timeoutTracker;
65     }
66
67     /**
68      * Schedule a {@link TimeoutTask} that aborts the task if not otherwise completed before the given timeout.
69      *
70      * @param timeoutExecutor the executor to execute the {@link TimeoutTask}
71      * @param timeoutInMills the timeout in milliseconds.
72      * @param threadToInterrupt the thread to interrupt
73      * @return a {@link TimeoutTracker}
74      */
75     public static TimeoutTracker timeSyncTaskIfNeeded(ScheduledExecutorService timeoutExecutor,
76                                                       long timeoutInMills,
77                                                       Thread threadToInterrupt) {
78         if (timeoutInMills <= 0) {
79             return NoOpTimeoutTracker.INSTANCE;
80         }
81
82         SyncTimeoutTask timeoutTask = new SyncTimeoutTask(threadToInterrupt);
83
84         ScheduledFuture<?> scheduledFuture =
85             timeoutExecutor.schedule(timeoutTask,
86                                      timeoutInMills,
87                                      TimeUnit.MILLISECONDS);
88         return new ApiCallTimeoutTracker(timeoutTask, scheduledFuture);
89     }
90
91     public static long resolveTimeoutInMillis(Supplier<Optional<Duration>> supplier, Duration fallback) {
92         return OptionalUtils.firstPresent(supplier.get(), () -> fallback)
93                             .map(Duration::toMillis)
94                             .orElse(0L);
95     }
96 }
97