How to wrap a supplier with a java future that returns a value on timeout but keeps running in the background (with option to cancel)

Question

As an example, I have a supplier that might take time to run:

Supplier<Integer> numLoader = sneaky(() -> {
    Thread.sleep(10000);
    System.out.println("5 Calculated!");
    return 5;
});

* sneaky is just a utility to convert to runtime exception.

I'd like to be able to do something like this:

Future<Integer> future = createFutureValueOnTimeout(-1, numLoader);
// numLoader takes 10 seconds to complete so -1 is returned.
int num = future.get(1000, TimeUnit.MILLISECONDS);
if (resourcesAreLow()) {
    future.cancel(true);
}
doSomethingWithTheValue(num);

I also have a partial implementation for createFutureValueOnTimeout:

private static <V> Future<V> createFutureValueOnTimeout(V v, Supplier<V> supplier) {
    CompletableFuture<V> completableFuture = CompletableFuture.supplyAsync(supplier);
    return new FutureDecorator<V>(completableFuture) {
        @Override
        public V get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException {
            return completableFuture.completeOnTimeout(v, timeout, unit).get();
        }
    };
}

The problem is that when calling cancel, the sleep isn't interrupted.

1. How can I get the cancel to work?
2. Is there an easier way to return a value on timeout?

Complete test:

public class TimeoutTest {
    @SneakyThrows
    @Test
    public void testTimeout() {
        int loadTimeMillis = 10000;
        Supplier<Integer> numLoader = () -> {
            try {
                // Simulate long operation
                Thread.sleep(loadTimeMillis);
            } catch (InterruptedException e) {
                System.out.println("Interrupted! message: " + e.getMessage());
                throw Lombok.sneakyThrow(e);
            }
            System.out.println("5 Calculated!");
            return 5;
        };

        Future<Integer> future = createFutureValueOnTimeout(-1, numLoader);
        
        long start = System.currentTimeMillis();

        // numLoader takes 10 seconds to complete so -1 is returned.
        int num = future.get(1000, TimeUnit.MILLISECONDS);

        System.out.println("Got: num: " + num + ". time: " + (System.currentTimeMillis() - start));

        if (resourcesAreLow()) {
            future.cancel(true);
        }
        // Don't stop the test. Give time for the cancel to kick in.
        Thread.sleep(loadTimeMillis);
        System.out.println("Finished. Time: " + (System.currentTimeMillis() - start));
    }

    private boolean resourcesAreLow() {
        return true;
    }

    private static <V> Future<V> createFutureValueOnTimeout(V v, Supplier<V> supplier) {
        CompletableFuture<V> completableFuture = CompletableFuture.supplyAsync(supplier);
        return new FutureDecorator<V>(completableFuture) {
            @Override
            public V get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException {
                return completableFuture.completeOnTimeout(v, timeout, unit).get();
            }
        };
    }

    private static class FutureDecorator<V> implements Future<V> {
        private final Future<V> inner;

        private FutureDecorator(Future<V> inner) {this.inner = inner;}

        @Override
        public boolean cancel(boolean mayInterruptIfRunning) {
            return inner.cancel(mayInterruptIfRunning);
        }

        @Override
        public boolean isCancelled() {
            return inner.isCancelled();
        }

        @Override
        public boolean isDone() {
            return inner.isDone();
        }

        @Override
        public V get() throws InterruptedException, ExecutionException {
            return inner.get();
        }

        @Override
        public V get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
            return inner.get(timeout, unit);
        }
    }
}

Output: (Notice the lack of the Interrupted! message)

Got: num: -1. time: 1007
5 Calculated!
Finished. Time: 11021

Answer 1

You can combine the Executor/Future API which supports cancelation with CompletableFuture:

public static <R> CompletableFuture<R> withInterruptionSupport(Callable<R> c) {
    CompletableFuture<R> cf = new CompletableFuture<>();
    FutureTask<R> ft = new FutureTask<R>(c) {
        @Override
        protected void set(R v) {
            super.set(v);
            cf.complete(v);
        }
        @Override
        protected void setException(Throwable t) {
            super.setException(t);
            cf.completeExceptionally(t);
        }
    };
    cf.defaultExecutor().execute(ft);
    cf.whenComplete((x,y) -> ft.cancel(true));
    return cf;
}

Since support for interruption in the actual function typically implies dealing with InterruptedException, it’s convenient to use Callable instead of Supplier, so it is allowed to throw this exception.

The Future which supports cancelation with interruption is unconditionally canceled whenever the CompletableFuture completes, which works without problems as whenever the completion stems from the task itself, the future is already completed and the subsequent cancelation will be ignored.

This means, we do not need to distinguish between the different possibilities of completion here. Not only does completeOnTimeout work, you can also invoke cancel(…) on the CompletableFuture and it will interrupt the evaluation of the Callable (the boolean argument still is irrelevant though). Even calling complete with an alternative result without waiting for a timeout would interrupt the now-obsolete evaluation.

So, the following works:

for(int timeout: new int[] { 5, 15 }) {
    System.out.println("with timeout of " + timeout);
    Integer i = withInterruptionSupport(() -> {
            Thread.sleep(10000);
            System.out.println("5 Calculated!");
            return 5;
        })
        .completeOnTimeout(42, timeout, TimeUnit.SECONDS)
        .join();
    System.out.println("got " + i);
}
ForkJoinPool.commonPool().awaitQuiescence(1, TimeUnit.DAYS);

with timeout of 5
got 42
with timeout of 15
5 Calculated!
got 5