pl.allegro.tech.hermes.consumers.consumer.sender.timeout.FutureAsyncTimeout Maven / Gradle / Ivy

Go to download
package pl.allegro.tech.hermes.consumers.consumer.sender.timeout;

import java.time.Duration;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.function.Function;

/**
 * see http://www.nurkiewicz.com/2014/12/asynchronous-timeouts-with.html
 */
public class FutureAsyncTimeout {

    private final ScheduledExecutorService executor;
    private final Function failure;

    public FutureAsyncTimeout(Function failure, ScheduledExecutorService scheduledExecutorService) {
        this.executor = scheduledExecutorService;
        this.failure = failure;
    }

    public CompletableFuture within(CompletableFuture future, Duration duration) {
        return future.applyToEither(failAfter(duration), Function.identity());
    }

    private CompletableFuture failAfter(Duration duration) {
        final CompletableFuture promise = new CompletableFuture<>();
        executor.schedule(() -> {
            TimeoutException ex = new TimeoutException("Timeout after " + duration);
            return promise.complete(failure.apply(ex));
        }, duration.toMillis(), TimeUnit.MILLISECONDS);
        return promise;
    }

    public void shutdown() {
        executor.shutdown();
    }
}