com.pivovarit.collectors.AsyncParallelCollector Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of parallel-collectors Show documentation
Parallel collection processing with customizable thread pools
There is a newer version: 3.2.0
package com.pivovarit.collectors;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.function.BiConsumer;
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.Stream;

import static com.pivovarit.collectors.BatchingStream.defaultBatchAmount;
import static com.pivovarit.collectors.BatchingStream.partitioned;
import static com.pivovarit.collectors.Dispatcher.unbounded;
import static java.util.Objects.requireNonNull;
import static java.util.stream.Collectors.collectingAndThen;
import static java.util.stream.Collectors.toList;

/**
 * @author Grzegorz Piwowarek
 */
final class AsyncParallelCollector
  implements Collector>, CompletableFuture> {

    private final Dispatcher dispatcher;
    private final Function mapper;
    private final Function>, CompletableFuture> processor;

    private final CompletableFuture result = new CompletableFuture<>();

    private AsyncParallelCollector(
      Function mapper,
      Dispatcher dispatcher,
      Function>, CompletableFuture> processor) {
        this.dispatcher = dispatcher;
        this.processor = processor;
        this.mapper = mapper;
    }

    @Override
    public Supplier>> supplier() {
        return ArrayList::new;
    }

    @Override
    public BinaryOperator>> combiner() {
        return (left, right) -> {
            throw new UnsupportedOperationException();
        };
    }

    @Override
    public BiConsumer>, T> accumulator() {
        return (acc, e) -> {
            startConsuming();
            acc.add(dispatcher.enqueue(() -> mapper.apply(e)));
        };
    }

    @Override
    public Function>, CompletableFuture> finisher() {
        return futures -> {
            dispatcher.stop();
            return processor.apply(toCombined(futures))
              .handle(result())
              .thenCompose(__ -> result);
        };
    }

    @Override
    public Set characteristics() {
        return Collections.emptySet();
    }

    private static  CompletableFuture> toCombined(List> futures) {
        return allOf(futures)
          .thenApply(__ -> futures.stream()
            .map(CompletableFuture::join));
    }

    private void startConsuming() {
        if (!dispatcher.isRunning()) {
            dispatcher.start().handle((__, ex) -> result.completeExceptionally(ex));
        }
    }

    private static  CompletableFuture allOf(List> futures) {
        CompletableFuture future = new CompletableFuture<>();

        futures.forEach(f -> f.handle((__, ex) -> ex != null && future.completeExceptionally(ex)));

        CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).thenAccept(future::complete);

        return future;
    }

    private BiFunction result() {
        return (c, ex) -> ex == null ? result.complete(c) : result.completeExceptionally(ex);
    }

    static  Collector>> collectingToStream(Function mapper, Executor executor) {
        requireNonNull(executor, "executor can't be null");
        requireNonNull(mapper, "mapper can't be null");
        return new AsyncParallelCollector<>(mapper, Dispatcher.limiting(executor), t -> t);
    }

    static  Collector>> collectingToStream(Function mapper, Executor executor, int parallelism) {
        requireNonNull(executor, "executor can't be null");
        requireNonNull(mapper, "mapper can't be null");
        requireValidParallelism(parallelism);
        return new AsyncParallelCollector<>(mapper, Dispatcher.limiting(executor, parallelism), t -> t);
    }

    static  Collector> collectingWithCollector(Collector collector, Function mapper, Executor executor) {
        requireNonNull(collector, "collector can't be null");
        requireNonNull(executor, "executor can't be null");
        requireNonNull(mapper, "mapper can't be null");
        return new AsyncParallelCollector<>(mapper, Dispatcher.limiting(executor), r -> r
          .thenApply(s -> s.collect(collector)));
    }

    static  Collector> collectingWithCollector(Collector collector, Function mapper, Executor executor, int parallelism) {
        requireNonNull(collector, "collector can't be null");
        requireNonNull(executor, "executor can't be null");
        requireNonNull(mapper, "mapper can't be null");
        requireValidParallelism(parallelism);
        return new AsyncParallelCollector<>(mapper, Dispatcher.limiting(executor, parallelism), r -> r
          .thenApply(s -> s.collect(collector)));
    }

    static  Collector>> collectingToStreamInBatches(Function mapper, Executor executor) {
        requireNonNull(executor, "executor can't be null");
        requireNonNull(mapper, "mapper can't be null");
        return collectingToStreamInBatches(mapper, executor, defaultBatchAmount());
    }

    static  Collector>> collectingToStreamInBatches(Function mapper, Executor executor, int parallelism) {
        requireNonNull(executor, "executor can't be null");
        requireNonNull(mapper, "mapper can't be null");
        requireValidParallelism(parallelism);

        return collectingAndThen(toList(), list -> partitioned(list, parallelism)
          .collect(new AsyncParallelCollector<>(batch(mapper), unbounded(executor), cf -> cf
            .thenApply(s -> s.flatMap(Collection::stream)))));
    }

    static  Collector> collectingWithCollectorInBatches(Collector collector, Function mapper, Executor executor) {
        requireNonNull(collector, "collector can't be null");
        requireNonNull(executor, "executor can't be null");
        requireNonNull(mapper, "mapper can't be null");
        return collectingWithCollectorInBatches(collector, mapper, executor, defaultBatchAmount());
    }

    static  Collector> collectingWithCollectorInBatches(Collector collector, Function mapper, Executor executor, int parallelism) {
        requireNonNull(collector, "collector can't be null");
        requireNonNull(executor, "executor can't be null");
        requireNonNull(mapper, "mapper can't be null");
        requireValidParallelism(parallelism);

        return batching(collector, mapper, executor, parallelism);
    }

    static  Collector> batching(Collector collector, Function mapper, Executor executor, int parallelism) {
        return collectingAndThen(toList(), list -> partitioned(list, parallelism)
          .collect(new AsyncParallelCollector<>(batch(mapper),
            unbounded(executor), cf -> cf.thenApply(s -> s.flatMap(Collection::stream).collect(collector)))));
    }

    static  Function, List> batch(Function mapper) {
        return batch -> batch.stream().map(mapper).collect(toList());
    }

    static void requireValidParallelism(int parallelism) {
        if (parallelism < 1) {
            throw new IllegalArgumentException("Parallelism can't be lower than 1");
        }
    }
}