CompletableFutures.java

package com.codepoetics.protonpack.collectors;

import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.BiFunction;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collector;
import java.util.stream.Collectors;
import java.util.Optional;

import static java.util.stream.Collectors.toList;

public final class CompletableFutures {

    private CompletableFutures() {
    }

    /**
     * Collector which converts a stream of CompletableFuture<T> into a CompletableFuture<List<T>>
     * @param <T> The type of value returned by each {@link CompletableFuture} in the stream.
     * @return A {@link CompletableFuture} which completes with a list of all the values returned by futures in the
     *         stream, once they have all completed, or fails if any future in the stream fails.
     */
    public static <T> Collector<CompletableFuture<T>, ?, CompletableFuture<List<T>>> toFutureList() {
        return Collectors.collectingAndThen(
                Collectors.<CompletableFuture<T>>toList(),
                futures -> {
                    AtomicLong resultsRemaining = new AtomicLong(futures.size());
                    CompletableFuture<List<T>> result = new CompletableFuture<>();

                    BiFunction<T, Throwable, Void> handler = (success, failure) -> {
                        if (failure == null) {
                            if (resultsRemaining.decrementAndGet() == 0) {
                                result.complete(
                                        futures.stream()
                                                .map(CompletableFutures::safeGet)
                                                .collect(toList()));
                            }
                        } else {
                            result.completeExceptionally(failure);
                        }
                        return null;
                    };

                    futures.forEach(future -> future.handle(handler));
                    return result;
                });
    }

    private static <T> T safeGet(CompletableFuture<T> future) {
        try {
            return future.get();
        } catch (InterruptedException | ExecutionException e) {
            throw new IllegalStateException("safeGet called on failed future: " + e);
        }
    }

    private interface FutureReducer<T, A> {

        static <T, A> FutureReducer<T, A> of(A identity, BiFunction<T, A, A> reducer, BinaryOperator<A> combiner) {
            return new SingleFutureReducer<>(identity, reducer, combiner);
        }

        CompletableFuture<A> complete();

        FutureReducer<T, A> addFuture(CompletableFuture<T> future);

        FutureReducer<T, A> combine(FutureReducer<T, A> other);
    }

    private static final class SingleFutureReducer<T, A> implements FutureReducer<T, A> {


        private AtomicLong futureCount = new AtomicLong(0L);
        private long resultsReceived = 0;
        private CompletableFuture<A> output = null;
        private Throwable exception = null;

        private A accumulator;
        private final BiFunction<T, A, A> reducer;
        private final BinaryOperator<A> combiner;

        private SingleFutureReducer(A accumulator, BiFunction<T, A, A> reducer, BinaryOperator<A> combiner) {
            this.accumulator = accumulator;
            this.reducer = reducer;
            this.combiner = combiner;
        }

        private synchronized void resultReceived(T result) {
            try {
                accumulator = reducer.apply(result, accumulator);
                resultsReceived += 1;
                completeIfReady();
            } catch (Exception e) {
                exceptionReceived(e);
            }
        }

        private synchronized void exceptionReceived(Throwable exception) {
            this.exception = exception;
            completeIfReady();
        }

        @Override
        public synchronized CompletableFuture<A> complete() {
            output = new CompletableFuture<A>();
            completeIfReady();
            return output;
        }

        private void completeIfReady() {
            if (output == null) {
                return;
            }
            if (exception != null) {
                output.completeExceptionally(exception);
            } else if (futureCount.get() == resultsReceived) {
                output.complete(accumulator);
            }
        }

        @Override
        public FutureReducer<T, A> addFuture(CompletableFuture<T> future) {
            futureCount.incrementAndGet();
            future.handle((result, error) -> {
                if (error != null) {
                    exceptionReceived(error);
                } else {
                    resultReceived(result);
                }
                return null;
            });
            return this;
        }

        @Override
        public FutureReducer<T, A> combine(FutureReducer<T, A> other) {
            return new CombinedFutureReducer<>(this, other, combiner);
        }
    }

    private static final class CombinedFutureReducer<T, A> implements FutureReducer<T, A> {
        private final FutureReducer<T, A> left;
        private final FutureReducer<T, A> right;
        private final BinaryOperator<A> combiner;

        private CombinedFutureReducer(FutureReducer<T, A> left, FutureReducer<T, A> right, BinaryOperator<A> combiner) {
            this.left = left;
            this.right = right;
            this.combiner = combiner;
        }

        @Override
        public CompletableFuture<A> complete() {
            return left.complete().thenCombine(right.complete(), combiner);
        }

        @Override
        public FutureReducer<T, A> addFuture(CompletableFuture<T> future) {
            throw new IllegalStateException("Cannot add futures after combination");
        }

        @Override
        public FutureReducer<T, A> combine(FutureReducer<T, A> other) {
            return new CombinedFutureReducer<>(this, other, combiner);
        }
    }

    public static <A> Collector<CompletableFuture<A>, ?, CompletableFuture<Optional<A>>> reducing(BinaryOperator<A> reducer) {
        return toFuture(
                Optional::empty,

                (left, maybeRight) -> maybeRight.isPresent()
                    ? maybeRight.map(right -> reducer.apply(right, left))
                    : Optional.of(left),

                (maybeLeft, maybeRight) -> maybeLeft.isPresent()
                    ? maybeLeft.flatMap(left -> maybeRight.map(right -> reducer.apply(left, right)))
                    : maybeRight);
    }

    public static <T, A> Collector<CompletableFuture<T>, ?, CompletableFuture<A>> toFuture(
            Supplier<A> identitySupplier, BiFunction<T, A, A> reducer, BinaryOperator<A> combiner) {
        return toFuture(identitySupplier, reducer, combiner, Function.identity());
    }

    public static <T, A, R> Collector<CompletableFuture<T>, ?, CompletableFuture<R>> toFuture(
            Supplier<A> identitySupplier, BiFunction<T, A, A> reducer, BinaryOperator<A> combiner, Function<? super A, ? extends R> completer) {
        return Collector.of(
                    () -> FutureReducer.of(identitySupplier.get(), reducer, combiner),
                    FutureReducer::addFuture,
                    FutureReducer::combine,
                    fr -> fr.complete().thenApply(completer));
    }

}