【Java】Stream流水线实现

Stream流水线

流水线的执行

Stream的操作可以分为两类: 中间操作和终端操作, 中间操作只是一种标记, 而只有终端操作才会触发最终的计算. 而中间操作又可以分为有状态操作和无状态操作, 其中无状态操作指元素处理不受前面的元素影响, 而有状态操作会受前面元素的影响. 终端操作又可以分为短路操作和非短路操作, 短路操作指无需处理完全部元素即可返回结果, 而非短路操作必须处理完全部的元素, 如的例子

@Test
public void pipelineLimit() {
    IntStream.range(1,10)
        	.boxed()
            .peek(x -> System.out.printf("A%d%n", x))
            .limit(3)
            .peek(x -> System.out.printf("B%d%n", x))
            .forEach(x -> System.out.printf("C%d%n", x));
}

这段代码输出的是

A1
B1
C1
A2
B2
C2
A3
B3
C3

只输出三组值的原因在于只有终端操作才会触发流水线的执行, 即只有forEach执行才会执行整个流水线; 而当第四次触发流水线的时候, 触发了limit的条件, 即刻终止流, 因此只输出了三组数据.

在流中, skip类似于continue, 而limit类似于break; 它并不会中断整个流的执行, 而只会中断当前流水线的执行, 如下面的例子

@Test
public void pipelineSkip() {
    IntStream.range(1, 10)
        	.boxed()
            .peek(x -> System.out.printf("A%d%n", x))
            .skip(6)
            .peek(x -> System.out.printf("B%d%n", x))
            .forEach(x -> System.out.printf("C%d%n", x));
}

它的输出为:

A1
A2
A3
A4
A5
A6
A7
B7
C7
A8
B8
C8
A9
B9
C9

每一个forEach都触发了流水线的执行, 但是当流水线执行到skip就不会执行后面的内容了, 因此前六个元素只会打印A, 而后三个元素不会被skip, 因此才会打印B/C

对于带有有状态操作的场景, 有状态的操作会执行完该操作前面的所有操作

public void statefulPipeline() {
    Stream.of(1, 6, 2, 5, 4, 3, 9, 8, 7)
            .peek(x -> System.out.printf("A%d%n", x))
            .sorted()
            .peek(x -> System.out.printf("B%d%n", x))
            .forEach(x -> System.out.printf("C%d%n", x));
}

它会执行完标记为A的无序的peek, 然后在逐个执行sorted后面的操作

A3
A9
A8
A7
B1
C1
B2
C2
B3
C3
B4
C4
B5
C5
B6
C6
B7
C7
B8
C8
B9
C9

Stream流的实现过程会修改执行的范围, 如下面的例子中, 因为peek对count的结果没有任何影响, 所以Stream的实现过程中会将peek流程省略掉

	public void shouldOptimizeExecution() {
        List<String> l = Arrays.asList("A", "B", "C", "D");
        long count = l.stream().peek(System.out::println).count();
        System.out.println(count);
    }
// peek不会被执行, 因为中间操作不会影响count()结果

流水线构造

在第一次到达终端操作, 会调用java.util.stream.AbstractPipeline#wrapAndCopyInto构造流水线; 它会先调用java.util.stream.AbstractPipeline#wrapSink将操作包装为链表, 然后再调用java.util.stream.AbstractPipeline#copyInto; 此时若检查到了短路操作, 则会优化执行流程, 否则则调用java.util.stream.AbstractPipeline#copyIntoWithCancel执行正常的执行流程

而对于类似count这样的短路操作,它构建流水线的流程相对于forEach较为简单; 它直接在java.util.stream.AbstractPipeline#exactOutputSizeIfKnown中对每个操作判断是否会导致长度变化, 以此可以跳过某些非必要的操作, 进而达到性能优化的效果

流水线操作的实现

在流水线中, 若是无状态操作的实现则非常简单, 只需要将所有的操作形成一个线性表, 然后按顺序执行就可以了因此不过多介绍; 但是这对于有状态操作无能为力, 有状态操作(如sort)需要等待前面所有的任务都执行完才能开始执行, 倘若每个链路都完全执行的话就会浪费许多计算资源, 而每个链路逐步执行的话有状态操作又会因为状态信息不是最新而产生错误的计算结果. 为了协调好相邻操作的关系, Stream引入了Sink接口完成交互.

interface Sink<T> extends Consumer<T> {
    
    // 开始遍历元素前会调用此方法, 通知sink做好准备
    default void begin(long size) {}

   // 完成遍历后调用此方法, 通知sink没有更多的元素了, 并准备通知下游
    default void end() {}

   // 对于短路操作, 告诉流水线是否可以结束操作, 返回true即可结束操作
    default boolean cancellationRequested() {
        return false;
    }

    // 遍历元素时调用, 对当前的元素进行处理; 
    // 前一个sink会调用当前sink的accept, 而当前sink会调用后一个sink的accept; 此过程会完成装啊提的传递
    // 这个方法继承于Consumer
    // sink中还对基本类型提供了accept, 这里就不展示了
    void accept(T t);
}

以下通过map和peek介绍无状态中间操作, 他们都是简单地调用方法然后通知下游即可

再通过sorted有状态非短路操作, 它会在开始执行前将所有的数据都准备好, 然后调用Arrays.sort进行排序

之后通过skip和limit介绍有状态短路操作, 之所以把他们放在一起的原因在于它们是通过一个类SliceOps实现的

最后通过forEach介绍终端操作, 它也是简单地对元素进行操作

map和peek的实现

以map为例, 这是一个简单的无状态的中间操作, 它直接调用mapper中的方法, 然后传递给下游的Sink

public final <R> Stream<R> map(Function<? super P_OUT, ? extends R> mapper) {
    Objects.requireNonNull(mapper);
    return new StatelessOp<P_OUT, R>(this, StreamShape.REFERENCE,
                                 StreamOpFlag.NOT_SORTED | StreamOpFlag.NOT_DISTINCT) {
        @Override
        Sink<P_OUT> opWrapSink(int flags, Sink<R> sink) {
            return new Sink.ChainedReference<P_OUT, R>(sink) {
                @Override
                public void accept(P_OUT u) {
                    // 这里直接通知下游接收当前操作的结果
                    downstream.accept(mapper.apply(u));
                }
            };
        }
    };
}

与map类似, peek通过action修改元素值, 然后再传递给下游

@Override
public final Stream<P_OUT> peek(Consumer<? super P_OUT> action) {
    Objects.requireNonNull(action);
    return new StatelessOp<P_OUT, P_OUT>(this, StreamShape.REFERENCE,
                                 0) {
        @Override
        Sink<P_OUT> opWrapSink(int flags, Sink<P_OUT> sink) {
            return new Sink.ChainedReference<P_OUT, P_OUT>(sink) {
                @Override
                public void accept(P_OUT u) {
                    // 与map十分相似, 只是从Function变为了Consumer
                    // 使得流水线可以对元素的值修改生效
                    action.accept(u);
                    downstream.accept(u);
                }
            };
        }
    };
}

Sorted实现

而sorted就是一个有状态的中间操作, 它需要所有元素都完成转换后才能够进行排序; 通过断点调试, 对于自定义Comparator的sorted操作, 调用的是SizedRefSortingSink, 它可以用于排序有大小的容器. 它本质上就是先将所有的数据都存放到一个集合中, 然后在调用Arrays.sort, 调用完成后再根据是否是短路操作判断是否要通知下游

private static final class SizedRefSortingSink<T> extends AbstractRefSortingSink<T> {
    private T[] array;
    private int offset;

    SizedRefSortingSink(Sink<? super T> sink, Comparator<? super T> comparator) {
        super(sink, comparator);
    }

    @Override
    @SuppressWarnings("unchecked")
    public void begin(long size) {
        if (size >= Nodes.MAX_ARRAY_SIZE)
            throw new IllegalArgumentException(Nodes.BAD_SIZE);
        // 1. 创建一个存放待排序元素的列表
        array = (T[]) new Object[(int) size];
    }

    @Override
    public void end() {
        // 3. 存放完成后开始排序
        Arrays.sort(array, 0, offset, comparator);
        // 排序完成后调用下游操作, 且只将非短路操作下发
        downstream.begin(offset);
        if (!cancellationRequestedCalled) {
            for (int i = 0; i < offset; i++)
                downstream.accept(array[i]);
        }
        else {
            for (int i = 0; i < offset && !downstream.cancellationRequested(); i++)
                downstream.accept(array[i]);
        }
        downstream.end();
        array = null;
    }

    // 2. 将所有的元素都存放到临时列表当中
    @Override
    public void accept(T t) {
        array[offset++] = t;
    }
}

limit/skip实现

limit的实现同前两者也类似, 它维护了一个计数器保存limit的参数值, 只有小于这个参数才会执行downstream

@Override
Sink<T> opWrapSink(int flags, Sink<T> sink) {
    return new Sink.ChainedReference<>(sink) {
        long n = skip;
        long m = normalizedLimit;

        @Override
        public void begin(long size) {
            // 多个limit/slice之间通过调用链传递最终的长度
            downstream.begin(calcSize(size, skip, m));
        }

        @Override
        public void accept(T t) {
            // n==0, 说明是limit模式, 根据m进行判断
            if (n == 0) {
                // 根据计数器来判断是否要执行下一步
                if (m > 0) {
                    m--;
                    downstream.accept(t);
                }
            }
            // 若是skip模式, 当skip计数器未将为0时, 跳过当前元素
            else {
                n--;
            }
        }

        @Override
        public boolean cancellationRequested() {
            // limit是短路操作, 当m为0的时候短路整个流水线
            return m == 0 || downstream.cancellationRequested();
        }
    };
}

forEach实现

static final class OfRef<T> extends ForEachOp<T> {
    final Consumer<? super T> consumer;

    OfRef(Consumer<? super T> consumer, boolean ordered) {
        super(ordered);
        this.consumer = consumer;
    }

    @Override
    public void accept(T t) {
        consumer.accept(t);
    }
}