4.x: Streamable, improved range and collect internals + JMH

src/jmh/java/io/reactivex/rxjava4/streamable/StreamableCollectPerf.java

@@ -0,0 +1,134 @@
+/*
+ * Copyright (c) 2016-present, RxJava Contributors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in
+ * compliance with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is
+ * distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See
+ * the License for the specific language governing permissions and limitations under the License.
+ */
+
+package io.reactivex.rxjava4.streamable;
+
+import java.util.*;
+import java.util.concurrent.*;
+import java.util.stream.Collectors;
+
+import org.openjdk.jmh.annotations.*;
+
+import io.reactivex.rxjava4.core.*;
+
+///
+/// Measure how much overhead and optimization is possible with the trampoline-reducer design
+/// of operators where each upstream item normally incurs 2 atomic operation due to needing
+/// a non-reentrant and serialized method [Streamer#next()] and [Streamer#finish()] calls,
+/// plus the cost of [CompletionStage#whenComplete(java.util.function.BiConsumer)] which allocates
+/// a continuation unnecessary for our particular operator.
+///
+/// Specs: i7 13700K, 64GB DDR5 4800MT CL40 RAM, Windows 11 25H2, Adoptium 26.0.1.8
+///
+/// ## 0 - Standard implementation
+///
+/// ```
+/// Benchmark                        (times)   Mode  Cnt         Score        Error  Units
+/// StreamableCollectPerf.benchmark        1  thrpt    5  13655087,463 ┬▒ 101002,376  ops/s
+/// StreamableCollectPerf.benchmark       10  thrpt    5   4345229,158 ┬▒  57122,655  ops/s
+/// StreamableCollectPerf.benchmark      100  thrpt    5    565231,578 ┬▒  10040,906  ops/s
+/// StreamableCollectPerf.benchmark     1000  thrpt    5     57845,486 ┬▒    459,777  ops/s
+/// StreamableCollectPerf.benchmark    10000  thrpt    5      5572,254 ┬▒    100,655  ops/s
+/// StreamableCollectPerf.benchmark   100000  thrpt    5       556,943 ┬▒      5,274  ops/s
+/// StreamableCollectPerf.benchmark  1000000  thrpt    5        56,532 ┬▒      1,475  ops/s
+/// ```
+///
+/// ## 1 - Avoid decrementing the `wip` counter as much as possible
+///
+/// +21% on 1M
+///
+/// ```
+/// Benchmark                        (times)   Mode  Cnt         Score        Error  Units
+/// StreamableCollectPerf.benchmark        1  thrpt    5  14576437,071 ┬▒ 213560,086  ops/s
+/// StreamableCollectPerf.benchmark       10  thrpt    5   5136926,834 ┬▒  75954,060  ops/s
+/// StreamableCollectPerf.benchmark      100  thrpt    5    682086,450 ┬▒  12237,157  ops/s
+/// StreamableCollectPerf.benchmark     1000  thrpt    5     69919,864 ┬▒   1098,636  ops/s
+/// StreamableCollectPerf.benchmark    10000  thrpt    5      6810,584 ┬▒    214,358  ops/s
+/// StreamableCollectPerf.benchmark   100000  thrpt    5       679,188 ┬▒      9,162  ops/s
+/// StreamableCollectPerf.benchmark  1000000  thrpt    5        68,844 ┬▒      6,202  ops/s
+/// ```
+///
+/// ## 2 - Avoid calling whenComplete
+///
+/// + 21% vs optimization 1, + 47% vs original
+///
+/// ```
+/// Benchmark                        (times)   Mode  Cnt         Score       Error  Units
+/// StreamableCollectPerf.benchmark        1  thrpt    5  16142384,643 ┬▒ 35614,613  ops/s
+/// StreamableCollectPerf.benchmark       10  thrpt    5   6145162,207 ┬▒ 81621,754  ops/s
+/// StreamableCollectPerf.benchmark      100  thrpt    5    832527,797 ┬▒ 13667,557  ops/s
+/// StreamableCollectPerf.benchmark     1000  thrpt    5     83811,957 ┬▒  2609,034  ops/s
+/// StreamableCollectPerf.benchmark    10000  thrpt    5      8411,203 ┬▒    31,486  ops/s
+/// StreamableCollectPerf.benchmark   100000  thrpt    5       836,305 ┬▒    14,383  ops/s
+/// StreamableCollectPerf.benchmark  1000000  thrpt    5        83,539 ┬▒     1,681  ops/s
+/// ```
+///
+///  ## 3 - Using a synchronous indexer
+///
+/// +413% vs optimization 2, +659% vs original
+///
+/// ```
+/// Benchmark                        (times)   Mode  Cnt         Score        Error  Units
+/// StreamableCollectPerf.benchmark        1  thrpt    5  16030652,636 ┬▒ 106257,177  ops/s
+/// StreamableCollectPerf.benchmark       10  thrpt    5  56457006,395 ┬▒ 407256,208  ops/s
+/// StreamableCollectPerf.benchmark      100  thrpt    5  22649008,057 ┬▒  87657,793  ops/s
+/// StreamableCollectPerf.benchmark     1000  thrpt    5    449961,073 ┬▒   4267,484  ops/s
+/// StreamableCollectPerf.benchmark    10000  thrpt    5     45892,148 ┬▒    367,249  ops/s
+/// StreamableCollectPerf.benchmark   100000  thrpt    5      4548,625 ┬▒     15,360  ops/s
+/// StreamableCollectPerf.benchmark  1000000  thrpt    5       429,078 ┬▒      6,829  ops/s
+/// ```
+///
+/// ## 4 - Removing the unnecessary volatile and double bookkeeping from [Streamable#range(int, int)]
+///
+/// +105% vs original
+///
+/// ```
+/// Benchmark                              (times)   Mode  Cnt         Score        Error  Units
+/// StreamableCollectPerf.benchmarkHidden        1  thrpt    5  16018548,246 ┬▒ 134658,704  ops/s
+/// StreamableCollectPerf.benchmarkHidden       10  thrpt    5   7977042,864 ┬▒ 126909,304  ops/s
+/// StreamableCollectPerf.benchmarkHidden      100  thrpt    5   1140735,027 ┬▒  14580,951  ops/s
+/// StreamableCollectPerf.benchmarkHidden     1000  thrpt    5    115843,240 ┬▒   1519,896  ops/s
+/// StreamableCollectPerf.benchmarkHidden    10000  thrpt    5     11606,956 ┬▒    181,325  ops/s
+/// StreamableCollectPerf.benchmarkHidden   100000  thrpt    5      1161,570 ┬▒     20,902  ops/s
+/// StreamableCollectPerf.benchmarkHidden  1000000  thrpt    5       116,332 ┬▒      1,313  ops/s
+/// ```
+@BenchmarkMode(Mode.Throughput)
+@Warmup(iterations = 5, time = 1, timeUnit = TimeUnit.SECONDS)
+@Measurement(iterations = 5, time = 1, timeUnit = TimeUnit.SECONDS)
+@OutputTimeUnit(TimeUnit.SECONDS)
+@Fork(value = 1)
+@State(Scope.Thread)
+public class StreamableCollectPerf {
+    @Param({ "1", "10", "100", "1000", "10000", "100000", "1000000" })
+    public int times;
+
+    Streamable<Optional<Integer>> result;
+
+    Streamable<Optional<Integer>> resultHidden;
+
+    @Setup
+    public void setup() {
+        result = Streamable.range(1, times).collect(Collectors.maxBy(Comparator.naturalOrder()));
+        resultHidden = Streamable.range(1, times).hide().collect(Collectors.maxBy(Comparator.naturalOrder()));
+    }
+
+    @Benchmark
+    public Object benchmark() {
+        return result.blockingFirst();
+    }
+
+    @Benchmark
+    public Object benchmarkHidden() {
+        return resultHidden.blockingFirst();
+    }
+}

src/main/java/io/reactivex/rxjava4/internal/operators/streamable/StreamableCollector.java

@@ -20,8 +20,10 @@
 
 import io.reactivex.rxjava4.annotations.NonNull;
 import io.reactivex.rxjava4.core.*;
-import io.reactivex.rxjava4.disposables.StreamerCancellation;
+import io.reactivex.rxjava4.disposables.*;
+import io.reactivex.rxjava4.exceptions.Exceptions;
 import io.reactivex.rxjava4.internal.fuseable.HasUpstreamStreamableSource;
+import io.reactivex.rxjava4.operators.IndexableSource;
 
 public record StreamableCollector<T, A, R>(
         Streamable<T> source,
@@ -34,7 +36,8 @@ public record StreamableCollector<T, A, R>(
                 source.stream(cancellation),
                 collector.supplier().get(),
                 collector.accumulator(),
-                collector.finisher());
+                collector.finisher(),
+                cancellation);
     }
 
     static final class CollectorStreamable<T, A, R>
@@ -54,25 +57,55 @@ static final class CollectorStreamable<T, A, R>
 
         final CompletableFuture<Void> finishReady;
 
+        final StreamerCancellation cancellation;
+
         R current;
 
-        int stage;
+        boolean once;
 
         volatile boolean done;
 
-        CollectorStreamable(Streamer<T> upstream, A storage, BiConsumer<A, T> accumulator, Function<A, R> finisher) {
+        CollectorStreamable(Streamer<T> upstream, A storage,
+                BiConsumer<A, T> accumulator,
+                Function<A, R> finisher,
+                StreamerCancellation cancellation) {
             this.upstream = upstream;
             this.wip = new AtomicInteger();
             this.storage = storage;
             this.accumulator = accumulator;
             this.finisher = finisher;
             this.nextReady = new CompletableFuture<>();
             this.finishReady = new CompletableFuture<>();
+            this.cancellation = cancellation;
         }
 
+        @SuppressWarnings("unchecked")
         @Override
         public @NonNull CompletionStage<Boolean> next() {
-            if (stage++ == 0) {
+            if (!once) {
+                once = true;
+
+                if (upstream instanceof IndexableSource<?> isrc) {
+                    long max = isrc.limit();
+                    for (long index = 0; index < max; index++) {
+                        if (cancellation.isDisposed()) {
+                            return CompletableFuture.failedFuture(new CancellationException());
+                        }
+
+                        T value;
+
+                        try {
+                            value = (T)isrc.elementAt(index);
+                        } catch (Throwable ex) {
+                            Exceptions.throwIfFatal(ex);
+                            return CompletableFuture.failedFuture(ex);
+                        }
+                        accumulator.accept(storage, value);
+                    }
+                    current = finisher.apply(storage);
+                    return NEXT_TRUE;
+                }
+
                 drain();
                 return nextReady;
             }
@@ -96,14 +129,28 @@ void drain() {
                 return;
             }
 
+            int wipMax = 1;
+            int wipIndex = 0;
             do {
                 if (done) {
-                    upstream.finish().whenComplete(this);
+                    StreamableHelper.whenComplete(upstream.finish(), this);
                     break;
                 } else {
-                    upstream.next().whenComplete(this);
+                    StreamableHelper.whenComplete(upstream.next(), this);
+                }
+                if (++wipIndex == wipMax) {
+                    var newWip = wip.get();
+                    if (newWip != wipMax) {
+                        wipMax = newWip;
+                    } else {
+                        wipMax = wip.addAndGet(-wipMax);
+                        if (wipMax == 0) {
+                            break;
+                        }
+                        wipIndex = 0;
+                    }
                 }
-            } while (wip.decrementAndGet() != 0);
+            } while (true);
         }
 
         @Override

src/main/java/io/reactivex/rxjava4/internal/operators/streamable/StreamableHelper.java

@@ -401,4 +401,23 @@ public static <T> StreamableInterceptConfig<T> createOnError(Consumer<? super Th
             });
             return cf;
         }, v -> v, (_, v) -> v);
-    }}
+    }
+
+    /**
+     * Shortcuts the stage to call the consumer directly if it is detected as completed.
+     * @param <T> the signal type of the {@link CompletionStage}
+     * @param stage the stage to handle
+     * @param consumer the consumer called with the stage's value or its exception
+     */
+    public static <T> void whenComplete(CompletionStage<T> stage, java.util.function.BiConsumer<? super T, ? super Throwable> consumer) {
+        if (stage instanceof CompletableFuture<T> cf && cf.isDone()) {
+            if (cf.isCompletedExceptionally()) {
+                consumer.accept(null, cf.exceptionNow());
+            } else {
+                consumer.accept(cf.getNow(null), null);
+            }
+        } else {
+            stage.whenComplete(consumer);
+        }
+    }
+}

src/main/java/io/reactivex/rxjava4/internal/operators/streamable/StreamableRange.java

@@ -18,6 +18,7 @@
 import io.reactivex.rxjava4.annotations.NonNull;
 import io.reactivex.rxjava4.core.*;
 import io.reactivex.rxjava4.disposables.StreamerCancellation;
+import io.reactivex.rxjava4.operators.IndexableSource;
 
 public record StreamableRange(int start, int count) implements Streamable<Integer> {
 
@@ -26,27 +27,25 @@ public record StreamableRange(int start, int count) implements Streamable<Intege
         return new RangeStreamer<>(start, start + count);
     }
 
-    static final class RangeStreamer<T> implements Streamer<Integer> {
+    static final class RangeStreamer<T> implements Streamer<Integer>, IndexableSource<Integer> {
 
-        final int end;
+        final int start;
 
-        volatile int current;
+        final int end;
 
-        volatile int index;
+        int current;
 
         RangeStreamer(int start, int end) {
-            index = start;
+            this.start = start;
+            this.current = start - 1;
             this.end = end;
         }
 
         @Override
         public @NonNull CompletionStage<Boolean> next() {
-            int i = index;
-            if (i >= end) {
+            if (++current >= end) {
                 return NEXT_FALSE;
             }
-            current = i;
-            index = i + 1;
             return NEXT_TRUE;
         }
 
@@ -57,9 +56,18 @@ static final class RangeStreamer<T> implements Streamer<Integer> {
 
         @Override
         public @NonNull CompletionStage<Void> finish() {
-            index = end;
             current = end;
             return FINISHED;
         }
+
+        @Override
+        public Integer elementAt(long index) {
+            return (int)(start + index);
+        }
+
+        @Override
+        public long limit() {
+            return end - start;
+        }
     }
 }

src/main/java/io/reactivex/rxjava4/operators/IndexableSource.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2016-present, RxJava Contributors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in
+ * compliance with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is
+ * distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See
+ * the License for the specific language governing permissions and limitations under the License.
+ */
+
+package io.reactivex.rxjava4.operators;
+
+import io.reactivex.rxjava4.annotations.NonNull;
+import io.reactivex.rxjava4.core.Streamer;
+
+/// Represents a source which can be accessed via a zero-based index synchronously,
+/// without going through the usual [Streamer#next()] calls to obtain the next item.
+/// @param <T> the element type of the source
+/// @since 4.0.0
+public interface IndexableSource<T> {
+
+    /**
+     * Obtain an element from the given index.
+     * Make sure you read only up to {@link #limit()}
+     * @param index the index
+     * @return the element at the specified index
+     * @throws Throwable if the indexed access involves computation that can throw
+     */
+    @NonNull
+    T elementAt(long index) throws Throwable;
+
+    /**
+     * Returns the limit of how many items can be obtained via [{@link #elementAt(long)}.
+     * @return the index limit, exclusive
+     */
+    long limit();
+}