Best way to support lazy/late parallel in customised stream method in Java 8

Question

Java implementation uses ReferencePipeline class which supports lazy/late .parallel() by design. That means these each group of lines of code are exactly identical:

// sequential identical lines:
stream.map(mapper).distict().filter(filter).sequential()...
stream.sequential().map(mapper).distict().filter(filter)...

// parallel identical lines:
stream.map(mapper).distict().filter(filter).parallel()...
stream.parallel().map(mapper).distict().filter(filter)...

Lets assume that I want to build a new customised stream method with the following signature:
public <T> static Stream<T> myMethod(Stream<T> stream)

And as requirement, the same .parallel() lazy/late behaviour should be present for that method. That means each group of lines should have exactly identical behaviours:

// sequential identical lines:
myMethod(stream).sequential()...
myMethod(stream.sequential())...

// parallel identical lines:
myMethod(stream).parallel()...
myMethod(stream.parallel())...

How can I do that? A simple example would useful.

public <T> static Stream<T> myMethod(Stream<T> stream) {
    // Any implementation that changes the stream
    // For simplification lets assume that this extensions
    // switches the odd with even positions, or anything else
    // that it's simpler and easy to demonstrate

    // This prevents parallel to be lazy/late!
    Spliterator<T> spliterator = stream.spliterator();
    return StreamSupport.stream(new Spliterator<T>() {
        // My easy implementation
    }, stream.isParallel());
}

Note that the use of .spliterator() and StreamSupport.stream() has implications in parallel processing discussed here: Understanding sequential vs parallel stream spliterators in Java 8 and Java 9

UPDATE: The junit5 test with assertj:

@ParameterizedTest(name="[{index}] {0}/{1}/{2} = {3} --> {4}")
@CsvSource({
        "-,-,-,1!,Sequential: default behaviour",
        "P,-,-,2+,Parallel: set in stage1",
        "-,P,-,2+,Parallel: set in stage2",
        "-,-,P,2+,Parallel: set in stage3",
        "P,S,-,1!,Sequential: set in stage2",
        "P,-,S,1!,Sequential: set in stage3",
        "P,S,P,2+,Parallel: set last in stage3",
        "S,P,S,1!,Sequential: set last in stage3",
})
void myMethodTest(String stage1f, String stage2f, String stage3f, String expThreads, String name) throws Exception {
    Set<String> set1 = new ConcurrentSkipListSet<>();
    Set<String> set2 = new ConcurrentSkipListSet<>();
    Set<String> set3 = new ConcurrentSkipListSet<>();
    int parallelism = 4;
    int minExpected = expThreads.equals("1!") ? 1 : 2;
    int maxExpected = expThreads.equals("1!") ? 1 : parallelism;

    BiFunction<String, Stream<Long>, Stream<Long>> mode = (flag, stream) -> {
        switch (flag) {
            case "P": return stream.parallel();
            case "S": return stream.sequential();
            default: return stream;
        }
    };

    Stream<Long> stage1 = mode.apply(stage1f, LongStream.range(0, 1000_000).boxed())
            .peek(x -> set1.add(Thread.currentThread().getName()));

    Stream<Long> stage2 = mode.apply(stage2f, myMethod(stage1).map(x -> 2*x))
            .peek(x -> set2.add(Thread.currentThread().getName()));

    Stream<Long> stage3 = mode.apply(stage3f, myMethod(stage2).map(x -> 2*x))
            .peek(x -> set3.add(Thread.currentThread().getName()));

    new ForkJoinPool(parallelism).submit(() -> {
        List<Long> list = stage3.collect(Collectors.toList());
        System.out.print("list:" + list.size() + "  threads:" + parallelism + "  flags:" + stage1f + "/" + stage2f + "/" + stage3f + "  ");
    }).get();

    System.out.print("stage1:" + set1.size() + "/" + maxExpected + "  ");
    System.out.print("stage2:" + set2.size() + "/" + maxExpected + "  ");
    System.out.println("stage3:" + set3.size() + "/" + maxExpected + "  ");
    assertThat(set1.size()).isBetween(minExpected, maxExpected);
    assertThat(set2.size()).isBetween(minExpected, maxExpected);
    assertThat(set3.size()).isBetween(minExpected, maxExpected);
}

Answer 1

This is tricky and the following solution may perhaps not work in all scenarios:

public static <T> Stream<T> myMethod(Stream<T> stream) {
    // for simplification, this does not alter anything from the source stream
    class CustomSpliterator implements Spliterator<T> {
        Spliterator<T> source;
        CustomSpliterator(Spliterator<T> src) {
            source = src;
        }
        @Override public boolean tryAdvance(Consumer<? super T> action) {
            return source.tryAdvance(action);
        }
        @Override public void forEachRemaining(Consumer<? super T> action) {
            source.forEachRemaining(action);
        }
        @Override public Spliterator<T> trySplit() {
            Spliterator<T> sSp = source.trySplit();
            return sSp == null? null: new CustomSpliterator(sSp);
        }
        @Override public long estimateSize() {
            return source.estimateSize();
        }
        @Override public int characteristics() {
            return source.characteristics();
        }
    }
    class MySpSupplier implements Supplier<Spliterator<T>> {
        Stream<T> downStream;
        @Override public Spliterator<T> get() {
            System.out.println("MySpSupplier.get(): nailing down behavior");
            Stream<T> src=downStream.isParallel()? stream.parallel(): stream.sequential();
            return new CustomSpliterator(src.spliterator());
        }
    }
    MySpSupplier sup = new MySpSupplier();
    Stream<T> s = StreamSupport.stream(sup, Spliterator.ORDERED, stream.isParallel());
    sup.downStream = s;
    return s;
}

But it fulfills the goal, e.g.

Stream<String> s = Stream.of("foo", "bar", "baz");
System.out.println("started with "+(s.isParallel()? "parallel": "sequential")+" stream");
s = s.peek(x -> System.out.println("upstream "+x+": "+Thread.currentThread()));
s = myMethod(s.map(String::toUpperCase));
System.out.println("chaining another ops");
s = s.peek(x -> System.out.println("downstream "+x+": "+Thread.currentThread()));
s = s.filter(str -> str.startsWith("B"));
System.out.println("turning to parallel");
s = s.parallel();
System.out.println("commencing terminal operation");
List<String> l = s.collect(Collectors.toList());
System.out.println("result: "+l);

started with sequential stream
chaining another ops
turning to parallel
commencing terminal operation
MySpSupplier.get(): nailing down behavior
upstream bar: Thread[main,5,main]
upstream foo: Thread[ForkJoinPool.commonPool-worker-1,5,main]
downstream BAR: Thread[main,5,main]
upstream baz: Thread[ForkJoinPool.commonPool-worker-2,5,main]
downstream FOO: Thread[ForkJoinPool.commonPool-worker-1,5,main]
downstream BAZ: Thread[ForkJoinPool.commonPool-worker-2,5,main]
result: [BAR, BAZ]

or

Stream<String> s = Stream.of("foo", "bar", "baz").parallel();
System.out.println("started with "+(s.isParallel()? "parallel": "sequential")+" stream");
s = s.peek(x -> System.out.println("upstream "+x+": "+Thread.currentThread()));
s = myMethod(s.map(String::toUpperCase));
System.out.println("chaining another ops");
s = s.peek(x -> System.out.println("downstream "+x+": "+Thread.currentThread()));
s = s.filter(str -> str.startsWith("B"));
System.out.println("turning to sequential");
s = s.sequential();
System.out.println("commencing terminal operation");
List<String> l = s.collect(Collectors.toList());
System.out.println("result: "+l);

started with parallel stream
chaining another ops
turning to sequential
commencing terminal operation
MySpSupplier.get(): nailing down behavior
upstream foo: Thread[main,5,main]
downstream FOO: Thread[main,5,main]
upstream bar: Thread[main,5,main]
downstream BAR: Thread[main,5,main]
upstream baz: Thread[main,5,main]
downstream BAZ: Thread[main,5,main]
result: [BAR, BAZ]