Java8 Stream 并行计算实现的原理

转自：http://lvheyang.com/?p=87

这两天组内的小伙伴在学习Java8，推广在新项目内使用新特性。正好看到了Stream 带来的遍历的多线程并发：
    

Arrays.asList(1, 2, 3, 4, 5, 6, 7, 9, 8, 0, 1)
        .stream()
        .parallel()
        .collect(Collectors.groupingBy(x -> x % 10))
        .forEach((x, y) -> System.out.println(x + ":" + y));

和小伙伴一起试着用各种玩法玩了一下Java8的函数式编程特性之后，感叹到这样子并发计算真的是越来越简单了的。

但是深入思考之后就会很自然的想到一个问题，这个过程中，我们并没有显示的告诉Stream，我们需要多少个线程进行并行计算？我们能否复用之前的线程进行计算？

带着这个问题我们先打开了VisualVM，查看一下我们运行这样一个任务会启动多少个线程？



我们可以看到默认的parallel计算启动了三个线程进行并行。这三个线程是怎么来的呢？抱着这个问题，我们来参考一下Jdk8的源码，来看看它是如何设置这个值的。

我们知道Stream 是一个惰性求值的系统（如何进行惰性求值，我会在另一篇博客中进行分析），那么我们只需要找它最后求值的过程，看它是怎样进行求值的就可以了。在AbstractPipeline 这个类里面我们找到了Stream 计算的最终求值过程的默认实现：

/**
 * Evaluate the pipeline with a terminal operation to produce a result.
 *
 * @param <R> the type of result
 * @param terminalOp the terminal operation to be applied to the pipeline.
 * @return the result
 */
final <R> R evaluate(TerminalOp<E_OUT, R> terminalOp) {
    assert getOutputShape() == terminalOp.inputShape();
    if (linkedOrConsumed)
        throw new IllegalStateException(MSG_STREAM_LINKED);
    linkedOrConsumed = true;

    return isParallel()
           ? terminalOp.evaluateParallel(this, sourceSpliterator(terminalOp.getOpFlags()))
           : terminalOp.evaluateSequential(this, sourceSpliterator(terminalOp.getOpFlags()));
}

在15行，我们可以看到，在求值的时候会检查并行计算的标志位，如果标志了并行计算的话，我们就会并行求值，反之则会串行求值。我们可以进一步进入并行求值的逻辑中，这是一个TerminalOp的默认接口方法，默认实现就是直接调用串行求值，在FindOp、ForEachOp、MatchOp 和 ReduceOp 中得到了覆盖。

@Override
public <P_IN> O evaluateParallel(PipelineHelper<T> helper,
                                 Spliterator<P_IN> spliterator) {
    return new FindTask<>(this, helper, spliterator).invoke();
}

如FindOp的代码示例，这四个操作都是创建一个Task的示例，然后执行invoke方法。这些Task的继承关系如图：



可以看出所有的Task 都继承自Jdk7 中引入的ForkJoin 并行框架的ForkJoinTask。所以我们可以看出Stream 的并行是依赖于ForkJoin 框架的。以AbstractTask 为例我们看看它是如何进行并行计算的：

@Override
public void compute() {
    Spliterator<P_IN> rs = spliterator, ls; // right, left spliterators
    long sizeEstimate = rs.estimateSize();
    long sizeThreshold = getTargetSize(sizeEstimate);
    boolean forkRight = false;
    @SuppressWarnings("unchecked") K task = (K) this;
    while (sizeEstimate > sizeThreshold && (ls = rs.trySplit()) != null) {
        K leftChild, rightChild, taskToFork;
        task.leftChild  = leftChild = task.makeChild(ls);
        task.rightChild = rightChild = task.makeChild(rs);
        task.setPendingCount(1);
        if (forkRight) {
            forkRight = false;
            rs = ls;
            task = leftChild;
            taskToFork = rightChild;
        }
        else {
            forkRight = true;
            task = rightChild;
            taskToFork = leftChild;
        }
        taskToFork.fork();
        sizeEstimate = rs.estimateSize();
    }
    task.setLocalResult(task.doLeaf());
    task.tryComplete();
}

ReduceTask(ReduceTask<P_IN, P_OUT, R, S> parent,
           Spliterator spliterator) {
    super(parent, spliterator);
    this.op = parent.op;
}

@Override
protected ReduceTask<P_IN, P_OUT, R, S> makeChild(Spliterator spliterator) {
    return new ReduceTask<>(this, spliterator);
}

这里面的主要逻辑就是

    先调用当前splititerator 方法的estimateSize 方法，预估这个分片中的数据量
    根据预估的数据量获取最小处理单元的大小阈值，即当数据量已经小于这个阈值的时候进行计算，否则进行fork 将任务划分成更小的数据块，进行求解。这里值得注意的是，getTargetSize 在第一次调用的时候会设置：
        预测数据量大小 / (默认并发度 * 4) 的结果作为最小执行单元的数量（配置的默认值是cpu 数 – 1，可以通过java.util.concurrent.ForkJoinPool.common.parallelism设置）
    如果当前分片大小仍然大于处理数据单元的阈值，且分片继续尝试切分成功，那么就继续切分，分别将左右分片的任务创建为新的Task，并且将当前的任务关联为两个新任务的父级任务（逻辑在makeChild 里面）
    先后对左右子节点的任务进行fork，对另外的分区进行分解。同时设定pending 为1，这代表一个task 实际上只会有一个等待的子节点（被fork）。
    当任务已经分解到足够小的时候退出循环，尝试进行结束。调用子类实现的doLeaf方法，完成最小计算单元的计算任务，并设置到当前任务的localResult中
    调用tryComplete 方法进行最终任务的扫尾工作，如果该任务pending 值不等于0，则原子的减1，如果已经等于0，说明任务都已经完成，则调用onCompletion 回调，如果该任务是叶子任务，则直接销毁中间数据结束；如果是中间节点会将左右子节点的结果进行合并
    检查如果这个任务已经没有父级任务了，则将该任务置为正常结束，如果还有则尝试递归的去调用父级节点的onCompletion回调，逐级进行任务的合并。

public final void tryComplete() {
    CountedCompleter<?> a = this, s = a;
    for (int c;;) {
        if ((c = a.pending) == 0) {
            a.onCompletion(s);
            if ((a = (s = a).completer) == null) {
                s.quietlyComplete();
                return;
            }
        }
        else if (U.compareAndSwapInt(a, PENDING, c, c - 1))
            return;
    }
}

说了这么多，大家也基本理解了Stream 的实现原理了。其实本质上就是在ForkJoin上进行了一层封装，将Stream 不断尝试分解成更小的split，然后使用fork/join 框架分而治之。

所以我们以往关于Fork/Join 的经验也都可以派上用场，可以解答之前我们的几个疑问：

    我在visualvm 中看到的 parallize 的3个线程是怎么来的？
        答：由于 taskToFork.fork() 调用，parallize使用了默认的ForkJoinPool.common 默认的一个静态线程池，这个线程池的默认线程个数是cpu 数量-1。由于我的代码是运行在四个逻辑内核的MacBook 上，所以这里的线程个数为3。如下面代码和注释所示：

if (parallelism < 0 && // default 1 less than #cores
    (parallelism = Runtime.getRuntime().availableProcessors() - 1) <= 0)
    parallelism = 1;
if (parallelism > MAX_CAP)
    parallelism = MAX_CAP;

    如何控制parallize 的线程数？
        答：我们可以自己构建一个ForkJoinPool，向其中提交一个parallize 任务，可以做到控制并发度。如以下示例代码：我们将之前的stream 的过程构造成一个runnable 的lambda 匿名函数 ()-> {…}。提交至线程池中，就可以按照我们想要的并发度进行计算了。

ForkJoinPool pool = new ForkJoinPool(2);
ret = pool.submit(() -> {
    return LongStream.range(1, 50 * 1024 * 1024).boxed().collect(Collectors.toList())
            .stream()
            .parallel()
            .map(x -> x * 2)
            .filter(x -> x < 1500)
            .reduce((x,y) -> x+y)
            .get();
}).get();

接下来打算继续深入这两个很有意思的问题：

    深入介绍一下Stream的惰性求值过程，最好能跟Scala 的Stream 实现进行比较：Java8 Stream 惰性求值实现分析 – 驴和羊

    深入介绍ForkJoin 的底层实现，包括它是如何进行线程调度和cache line sharing 优化的

参考文献：

1. http://blog.krecan.net/2014/03/18/how-to-specify-thread-pool-for-java-8-parallel-streams/