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.AbstractMap;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.function.BiConsumer;
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 java.util.Objects.requireNonNull;
import static java.util.concurrent.CompletableFuture.allOf;
import static java.util.function.Function.identity;
import static java.util.stream.Collectors.toCollection;
import static java.util.stream.Collectors.toMap;

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

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

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

    private AsyncParallelCollector(
      Function function,
      Function>, CompletableFuture> processor,
      Executor executor,
      int parallelism) {
        this.dispatcher = new Dispatcher<>(executor, parallelism);
        this.processor = processor;
        this.function = function;
    }

    private AsyncParallelCollector(
      Function function,
      Function>, CompletableFuture> processor,
      Executor executor) {
        this.dispatcher = new Dispatcher<>(executor);
        this.processor = processor;
        this.function = function;
    }

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

    @Override
    public BinaryOperator>> combiner() {
        return (left, right) -> {
            left.addAll(right);
            return left;
        };
    }

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

    @Override
    public Function>, CompletableFuture> finisher() {
        return futures -> {
            dispatcher.stop();

            processor.apply(combined(futures))
              .whenComplete((c, throwable) -> {
                  if (throwable == null) {
                      result.complete(c);
                  } else {
                      result.completeExceptionally(throwable);
                  }
              });

            return result;
        };
    }

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

    private static  CompletableFuture> combined(List> futures) {
        return allOf(futures.toArray(new CompletableFuture[0]))
          .thenApply(__ -> futures.stream()
            .map(CompletableFuture::join));
    }

    private void startConsuming() {
        if (!dispatcher.isRunning()) {
            dispatcher.start()
              .exceptionally(throwable -> {
                  result.completeExceptionally(throwable);
                  return null;
              });
        }
    }

    private static > Function>, CompletableFuture> toCollectionStrategy(Supplier collectionFactory) {
        return result -> result.thenApply(futures -> futures.collect(toCollection(collectionFactory)));
    }

    static > Collector> collectingToCollection(Function mapper, Supplier collectionSupplier, Executor executor) {
        requireNonNull(collectionSupplier, "collectionSupplier can't be null");
        requireNonNull(executor, "executor can't be null");
        requireNonNull(mapper, "mapper can't be null");
        return new AsyncParallelCollector<>(mapper, toCollectionStrategy(collectionSupplier), executor);
    }

    static > Collector> collectingToCollection(Function mapper, Supplier collectionSupplier, Executor executor, int parallelism) {
        requireNonNull(collectionSupplier, "collectionSupplier can't be null");
        requireNonNull(executor, "executor can't be null");
        requireNonNull(mapper, "mapper can't be null");
        requireValidParallelism(parallelism);
        return new AsyncParallelCollector<>(mapper, toCollectionStrategy(collectionSupplier), executor, parallelism);
    }

    static  Collector>> collectingToMap(Function keyMapper, Function valueMapper, Executor executor) {
        return collectingToMap(keyMapper, valueMapper, defaultMapSupplier(), uniqueKeyMerger(), executor);
    }

    static  Collector>> collectingToMap(Function keyMapper, Function valueMapper, Executor executor, int parallelism) {
        return collectingToMap(keyMapper, valueMapper, defaultMapSupplier(), uniqueKeyMerger(), executor, parallelism);
    }

    static  Collector>> collectingToMap(Function keyMapper, Function valueMapper, BinaryOperator merger, Executor executor) {
        return collectingToMap(keyMapper, valueMapper, defaultMapSupplier(), merger, executor);
    }

    static  Collector>> collectingToMap(Function keyMapper, Function valueMapper, BinaryOperator merger, Executor executor, int parallelism) {
        return collectingToMap(keyMapper, valueMapper, defaultMapSupplier(), merger, executor, parallelism);
    }

    static  Collector>> collectingToMap(Function keyMapper, Function valueMapper, Supplier> mapSupplier, Executor executor) {
        return collectingToMap(keyMapper, valueMapper, mapSupplier, uniqueKeyMerger(), executor);
    }

    static  Collector>> collectingToMap(Function keyMapper, Function valueMapper, Supplier> mapSupplier, Executor executor, int parallelism) {
        return collectingToMap(keyMapper, valueMapper, mapSupplier, uniqueKeyMerger(), executor, parallelism);
    }

    static  Collector>> collectingToMap(Function keyMapper, Function valueMapper, Supplier> mapSupplier, BinaryOperator merger, Executor executor) {
        requireNonNull(executor, "executor can't be null");
        requireNonNull(keyMapper, "keyMapper can't be null");
        requireNonNull(valueMapper, "valueMapper can't be null");
        requireNonNull(merger, "merger can't be null");
        requireNonNull(mapSupplier, "mapSupplier can't be null");
        return new AsyncParallelCollector<>(toEntry(keyMapper, valueMapper), toMapStrategy(merger, mapSupplier), executor);
    }

    static  Collector>> collectingToMap(Function keyMapper, Function valueMapper, Supplier> mapSupplier, BinaryOperator merger, Executor executor, int parallelism) {
        requireNonNull(executor, "executor can't be null");
        requireNonNull(keyMapper, "keyMapper can't be null");
        requireNonNull(valueMapper, "valueMapper can't be null");
        requireNonNull(merger, "merger can't be null");
        requireNonNull(mapSupplier, "mapSupplier can't be null");
        requireValidParallelism(parallelism);
        return new AsyncParallelCollector<>(toEntry(keyMapper, valueMapper), toMapStrategy(merger, mapSupplier), executor, parallelism);
    }

    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, identity(), executor);
    }

    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, identity(), executor, parallelism);
    }

    static  Supplier> defaultListSupplier() {
        return ArrayList::new;
    }

    static  Supplier> defaultSetSupplier() {
        return HashSet::new;
    }

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

    private static  BinaryOperator uniqueKeyMerger() {
        return (i1, i2) -> { throw new IllegalStateException("Duplicate key found"); };
    }

    private static  Supplier> defaultMapSupplier() {
        return HashMap::new;
    }

    private static >Function>>, CompletableFuture> toMapStrategy(BinaryOperator duplicateKeyResolutionStrategy, Supplier mapFactory) {
        return result -> result.thenApply(futures -> futures.collect(toMap(Map.Entry::getKey, Map.Entry::getValue, duplicateKeyResolutionStrategy, mapFactory)));
    }

    private static  Function> toEntry(Function keyMapper, Function valueMapper) {
        return entry -> new AbstractMap.SimpleEntry<>(keyMapper.apply(entry), valueMapper.apply(entry));
    }
}