3.x: Add onSubscribe hook to ParallelFlowable operators (#7191)

akarnokd · web-flow · commit 22c5e0bfe0ca · 2021-02-24T17:40:46.000+01:00
diff --git a/src/main/java/io/reactivex/rxjava3/core/Scheduler.java b/src/main/java/io/reactivex/rxjava3/core/Scheduler.java
@@ -91,7 +91,7 @@
 public abstract class Scheduler {
     /**
      * Value representing whether to use {@link System#nanoTime()}, or default as clock for {@link #now(TimeUnit)}
-     * and {@link Scheduler.Worker#now(TimeUnit)}
+     * and {@link Scheduler.Worker#now(TimeUnit)}.
      * <p>
      * Associated system parameter:
      * <ul>
@@ -111,7 +111,7 @@ public abstract class Scheduler {
      * @throws NullPointerException if {@code unit} is {@code null}
      */
     static long computeNow(TimeUnit unit) {
-        if(!IS_DRIFT_USE_NANOTIME) {
+        if (!IS_DRIFT_USE_NANOTIME) {
             return unit.convert(System.currentTimeMillis(), TimeUnit.MILLISECONDS);
         }
         return unit.convert(System.nanoTime(), TimeUnit.NANOSECONDS);
diff --git a/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelCollect.java b/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelCollect.java
@@ -47,6 +47,8 @@ public ParallelCollect(ParallelFlowable<? extends T> source,
 
     @Override
     public void subscribe(Subscriber<? super C>[] subscribers) {
+        subscribers = RxJavaPlugins.onSubscribe(this, subscribers);
+
         if (!validate(subscribers)) {
             return;
         }
diff --git a/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelConcatMap.java b/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelConcatMap.java
@@ -19,6 +19,7 @@
 import io.reactivex.rxjava3.internal.operators.flowable.FlowableConcatMap;
 import io.reactivex.rxjava3.internal.util.ErrorMode;
 import io.reactivex.rxjava3.parallel.ParallelFlowable;
+import io.reactivex.rxjava3.plugins.RxJavaPlugins;
 
 import java.util.Objects;
 
@@ -55,6 +56,8 @@ public int parallelism() {
 
     @Override
     public void subscribe(Subscriber<? super R>[] subscribers) {
+        subscribers = RxJavaPlugins.onSubscribe(this, subscribers);
+
         if (!validate(subscribers)) {
             return;
         }
diff --git a/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelDoOnNextTry.java b/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelDoOnNextTry.java
@@ -48,6 +48,8 @@ public ParallelDoOnNextTry(ParallelFlowable<T> source, Consumer<? super T> onNex
 
     @Override
     public void subscribe(Subscriber<? super T>[] subscribers) {
+        subscribers = RxJavaPlugins.onSubscribe(this, subscribers);
+
         if (!validate(subscribers)) {
             return;
         }
diff --git a/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelFilter.java b/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelFilter.java
@@ -40,6 +40,8 @@ public ParallelFilter(ParallelFlowable<T> source, Predicate<? super T> predicate
 
     @Override
     public void subscribe(Subscriber<? super T>[] subscribers) {
+        subscribers = RxJavaPlugins.onSubscribe(this, subscribers);
+
         if (!validate(subscribers)) {
             return;
         }
diff --git a/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelFilterTry.java b/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelFilterTry.java
@@ -46,6 +46,8 @@ public ParallelFilterTry(ParallelFlowable<T> source, Predicate<? super T> predic
 
     @Override
     public void subscribe(Subscriber<? super T>[] subscribers) {
+        subscribers = RxJavaPlugins.onSubscribe(this, subscribers);
+
         if (!validate(subscribers)) {
             return;
         }
diff --git a/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelFlatMap.java b/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelFlatMap.java
@@ -18,6 +18,7 @@
 import io.reactivex.rxjava3.functions.Function;
 import io.reactivex.rxjava3.internal.operators.flowable.FlowableFlatMap;
 import io.reactivex.rxjava3.parallel.ParallelFlowable;
+import io.reactivex.rxjava3.plugins.RxJavaPlugins;
 
 /**
  * Flattens the generated Publishers on each rail.
@@ -57,6 +58,8 @@ public int parallelism() {
 
     @Override
     public void subscribe(Subscriber<? super R>[] subscribers) {
+        subscribers = RxJavaPlugins.onSubscribe(this, subscribers);
+
         if (!validate(subscribers)) {
             return;
         }
diff --git a/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelFlatMapIterable.java b/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelFlatMapIterable.java
@@ -18,6 +18,7 @@
 import io.reactivex.rxjava3.functions.Function;
 import io.reactivex.rxjava3.internal.operators.flowable.FlowableFlattenIterable;
 import io.reactivex.rxjava3.parallel.ParallelFlowable;
+import io.reactivex.rxjava3.plugins.RxJavaPlugins;
 
 /**
  * Flattens the generated {@link Iterable}s on each rail.
@@ -50,6 +51,8 @@ public int parallelism() {
 
     @Override
     public void subscribe(Subscriber<? super R>[] subscribers) {
+        subscribers = RxJavaPlugins.onSubscribe(this, subscribers);
+
         if (!validate(subscribers)) {
             return;
         }
diff --git a/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelFromArray.java b/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelFromArray.java
@@ -16,6 +16,7 @@
 import org.reactivestreams.*;
 
 import io.reactivex.rxjava3.parallel.ParallelFlowable;
+import io.reactivex.rxjava3.plugins.RxJavaPlugins;
 
 /**
  * Wraps multiple Publishers into a ParallelFlowable which runs them
@@ -37,6 +38,8 @@ public int parallelism() {
 
     @Override
     public void subscribe(Subscriber<? super T>[] subscribers) {
+        subscribers = RxJavaPlugins.onSubscribe(this, subscribers);
+
         if (!validate(subscribers)) {
             return;
         }
diff --git a/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelFromPublisher.java b/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelFromPublisher.java
@@ -24,6 +24,7 @@
 import io.reactivex.rxjava3.internal.subscriptions.SubscriptionHelper;
 import io.reactivex.rxjava3.internal.util.BackpressureHelper;
 import io.reactivex.rxjava3.parallel.ParallelFlowable;
+import io.reactivex.rxjava3.plugins.RxJavaPlugins;
 
 /**
  * Dispatches the values from upstream in a round robin fashion to subscribers which are
@@ -51,6 +52,8 @@ public int parallelism() {
 
     @Override
     public void subscribe(Subscriber<? super T>[] subscribers) {
+        subscribers = RxJavaPlugins.onSubscribe(this, subscribers);
+
         if (!validate(subscribers)) {
             return;
         }
diff --git a/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelMap.java b/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelMap.java
@@ -44,6 +44,8 @@ public ParallelMap(ParallelFlowable<T> source, Function<? super T, ? extends R>
 
     @Override
     public void subscribe(Subscriber<? super R>[] subscribers) {
+        subscribers = RxJavaPlugins.onSubscribe(this, subscribers);
+
         if (!validate(subscribers)) {
             return;
         }
diff --git a/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelMapTry.java b/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelMapTry.java
@@ -49,6 +49,8 @@ public ParallelMapTry(ParallelFlowable<T> source, Function<? super T, ? extends
 
     @Override
     public void subscribe(Subscriber<? super R>[] subscribers) {
+        subscribers = RxJavaPlugins.onSubscribe(this, subscribers);
+
         if (!validate(subscribers)) {
             return;
         }
diff --git a/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelPeek.java b/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelPeek.java
@@ -66,6 +66,8 @@ public ParallelPeek(ParallelFlowable<T> source,
 
     @Override
     public void subscribe(Subscriber<? super T>[] subscribers) {
+        subscribers = RxJavaPlugins.onSubscribe(this, subscribers);
+
         if (!validate(subscribers)) {
             return;
         }
diff --git a/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelReduce.java b/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelReduce.java
@@ -46,6 +46,8 @@ public ParallelReduce(ParallelFlowable<? extends T> source, Supplier<R> initialS
 
     @Override
     public void subscribe(Subscriber<? super R>[] subscribers) {
+        subscribers = RxJavaPlugins.onSubscribe(this, subscribers);
+
         if (!validate(subscribers)) {
             return;
         }
diff --git a/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelRunOn.java b/src/main/java/io/reactivex/rxjava3/internal/operators/parallel/ParallelRunOn.java
@@ -49,7 +49,9 @@ public ParallelRunOn(ParallelFlowable<? extends T> parent,
     }
 
     @Override
-    public void subscribe(final Subscriber<? super T>[] subscribers) {
+    public void subscribe(Subscriber<? super T>[] subscribers) {
+        subscribers = RxJavaPlugins.onSubscribe(this, subscribers);
+
         if (!validate(subscribers)) {
             return;
         }
diff --git a/src/main/java/io/reactivex/rxjava3/plugins/RxJavaPlugins.java b/src/main/java/io/reactivex/rxjava3/plugins/RxJavaPlugins.java
@@ -112,6 +112,10 @@ public final class RxJavaPlugins {
     @Nullable
     static volatile BiFunction<? super Completable, ? super CompletableObserver, ? extends CompletableObserver> onCompletableSubscribe;
 
+    @SuppressWarnings("rawtypes")
+    @Nullable
+    static volatile BiFunction<? super ParallelFlowable, ? super Subscriber[], ? extends Subscriber[]> onParallelSubscribe;
+
     @Nullable
     static volatile BooleanSupplier onBeforeBlocking;
 
@@ -525,6 +529,7 @@ public static void reset() {
         setOnMaybeSubscribe(null);
 
         setOnParallelAssembly(null);
+        setOnParallelSubscribe(null);
 
         setFailOnNonBlockingScheduler(false);
         setOnBeforeBlocking(null);
@@ -992,6 +997,23 @@ public static <T> MaybeObserver<? super T> onSubscribe(@NonNull Maybe<T> source,
         return observer;
     }
 
+    /**
+     * Calls the associated hook function.
+     * @param <T> the value type
+     * @param source the hook's input value
+     * @param subscribers the array of subscribers
+     * @return the value returned by the hook
+     */
+    @SuppressWarnings({ "rawtypes" })
+    @NonNull
+    public static <T> Subscriber<@NonNull ? super T>[] onSubscribe(@NonNull ParallelFlowable<T> source, @NonNull Subscriber<@NonNull ? super T>[] subscribers) {
+        BiFunction<? super ParallelFlowable, ? super Subscriber[], ? extends Subscriber[]> f = onParallelSubscribe;
+        if (f != null) {
+            return apply(f, source, subscribers);
+        }
+        return subscribers;
+    }
+
     /**
      * Calls the associated hook function.
      * @param <T> the value type
@@ -1128,6 +1150,32 @@ public static void setOnParallelAssembly(@Nullable Function<? super ParallelFlow
         return onParallelAssembly;
     }
 
+    /**
+     * Sets the specific hook function.
+     * @param handler the hook function to set, null allowed
+     * @since 3.0.11 - experimental
+     */
+    @SuppressWarnings("rawtypes")
+    @Experimental
+    public static void setOnParallelSubscribe(@Nullable BiFunction<? super ParallelFlowable, ? super Subscriber[], ? extends Subscriber[]> handler) {
+        if (lockdown) {
+            throw new IllegalStateException("Plugins can't be changed anymore");
+        }
+        onParallelSubscribe = handler;
+    }
+
+    /**
+     * Returns the current hook function.
+     * @return the hook function, may be null
+     * @since 3.0.11 - experimental
+     */
+    @SuppressWarnings("rawtypes")
+    @Experimental
+    @Nullable
+    public static BiFunction<? super ParallelFlowable, ? super Subscriber[], ? extends Subscriber[]> getOnParallelSubscribe() {
+        return onParallelSubscribe;
+    }
+
     /**
      * Calls the associated hook function.
      * <p>History: 2.0.6 - experimental; 2.1 - beta
diff --git a/src/test/java/io/reactivex/rxjava3/plugins/RxJavaPluginsTest.java b/src/test/java/io/reactivex/rxjava3/plugins/RxJavaPluginsTest.java
@@ -592,6 +592,61 @@ public void onComplete() {
         .assertComplete();
     }
 
+    @SuppressWarnings("rawtypes")
+    @Test
+    public void parallelFlowableStart() {
+        try {
+            RxJavaPlugins.setOnParallelSubscribe(new BiFunction<ParallelFlowable, Subscriber[], Subscriber[]>() {
+                @Override
+                public Subscriber[] apply(ParallelFlowable f, final Subscriber[] t) {
+                    return new Subscriber[] { new Subscriber() {
+
+                            @Override
+                            public void onSubscribe(Subscription s) {
+                                t[0].onSubscribe(s);
+                            }
+
+                            @SuppressWarnings("unchecked")
+                            @Override
+                            public void onNext(Object value) {
+                                t[0].onNext((Integer)value - 9);
+                            }
+
+                            @Override
+                            public void onError(Throwable e) {
+                                t[0].onError(e);
+                            }
+
+                            @Override
+                            public void onComplete() {
+                                t[0].onComplete();
+                            }
+
+                        }
+                    };
+                }
+            });
+
+            Flowable.range(10, 3)
+            .parallel(1)
+            .sequential()
+            .test()
+            .assertValues(1, 2, 3)
+            .assertNoErrors()
+            .assertComplete();
+        } finally {
+            RxJavaPlugins.reset();
+        }
+        // make sure the reset worked
+        Flowable.range(10, 3)
+        .parallel(1)
+        .sequential()
+        .test()
+        .assertValues(10, 11, 12)
+        .assertNoErrors()
+        .assertComplete();
+    }
+
     @SuppressWarnings("rawtypes")
     @Test
     public void singleCreate() {
@@ -1176,6 +1231,7 @@ public void onComplete() {
             }
 
             AllSubscriber all = new AllSubscriber();
+            Subscriber[] allArray = { all };
 
             assertNull(RxJavaPlugins.onSubscribe(Observable.never(), null));
 
@@ -1197,6 +1253,10 @@ public void onComplete() {
 
             assertSame(all, RxJavaPlugins.onSubscribe(Maybe.never(), all));
 
+            assertNull(RxJavaPlugins.onSubscribe(Flowable.never().parallel(), null));
+
+            assertSame(allArray, RxJavaPlugins.onSubscribe(Flowable.never().parallel(), allArray));
+
             final Scheduler s = ImmediateThinScheduler.INSTANCE;
             Supplier<Scheduler> c = new Supplier<Scheduler>() {
                 @Override

Original file line number	Diff line number	Diff line change
`@@ -47,6 +47,8 @@ public ParallelCollect(ParallelFlowable<? extends T> source,`
`47`	`47`
`48`	`48`	`@Override`
`49`	`49`	`public void subscribe(Subscriber<? super C>[] subscribers) {`
	`50`	`+ subscribers = RxJavaPlugins.onSubscribe(this, subscribers);`
	`51`	`+`
`50`	`52`	`if (!validate(subscribers)) {`
`51`	`53`	`return;`
`52`	`54`	`}`
Original file line number	Diff line number	Diff line change
`@@ -48,6 +48,8 @@ public ParallelDoOnNextTry(ParallelFlowable<T> source, Consumer<? super T> onNex`
`48`	`48`
`49`	`49`	`@Override`
`50`	`50`	`public void subscribe(Subscriber<? super T>[] subscribers) {`
	`51`	`+ subscribers = RxJavaPlugins.onSubscribe(this, subscribers);`
	`52`	`+`
`51`	`53`	`if (!validate(subscribers)) {`
`52`	`54`	`return;`
`53`	`55`	`}`
Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,8 @@ public ParallelFilter(ParallelFlowable<T> source, Predicate<? super T> predicate`
`40`	`40`
`41`	`41`	`@Override`
`42`	`42`	`public void subscribe(Subscriber<? super T>[] subscribers) {`
	`43`	`+ subscribers = RxJavaPlugins.onSubscribe(this, subscribers);`
	`44`	`+`
`43`	`45`	`if (!validate(subscribers)) {`
`44`	`46`	`return;`
`45`	`47`	`}`
Original file line number	Diff line number	Diff line change
`@@ -46,6 +46,8 @@ public ParallelFilterTry(ParallelFlowable<T> source, Predicate<? super T> predic`
`46`	`46`
`47`	`47`	`@Override`
`48`	`48`	`public void subscribe(Subscriber<? super T>[] subscribers) {`
	`49`	`+ subscribers = RxJavaPlugins.onSubscribe(this, subscribers);`
	`50`	`+`
`49`	`51`	`if (!validate(subscribers)) {`
`50`	`52`	`return;`
`51`	`53`	`}`