2.x: allow subscribeOn to work with blocking create (#4770)

Author: akarnokd

src/main/java/io/reactivex/Flowable.java

@@ -12251,7 +12251,7 @@ public final <E extends Subscriber<? super T>> E subscribeWith(E subscriber) {
     @SchedulerSupport(SchedulerSupport.CUSTOM)
     public final Flowable<T> subscribeOn(Scheduler scheduler) {
         ObjectHelper.requireNonNull(scheduler, "scheduler is null");
-        return RxJavaPlugins.onAssembly(new FlowableSubscribeOn<T>(this, scheduler));
+        return RxJavaPlugins.onAssembly(new FlowableSubscribeOn<T>(this, scheduler, this instanceof FlowableCreate));
     }
 
     /**

src/main/java/io/reactivex/internal/operators/flowable/FlowableSubscribeOn.java

@@ -21,18 +21,28 @@
 import io.reactivex.internal.subscriptions.SubscriptionHelper;
 import io.reactivex.internal.util.BackpressureHelper;
 
+/**
+ * Subscribes to the source Flowable on the specified Scheduler and makes
+ * sure downstream requests are scheduled there as well.
+ *
+ * @param <T> the value type emitted
+ */
 public final class FlowableSubscribeOn<T> extends AbstractFlowableWithUpstream<T , T> {
+
     final Scheduler scheduler;
 
-    public FlowableSubscribeOn(Publisher<T> source, Scheduler scheduler) {
+    final boolean nonScheduledRequests;
+
+    public FlowableSubscribeOn(Publisher<T> source, Scheduler scheduler, boolean nonScheduledRequests) {
         super(source);
         this.scheduler = scheduler;
+        this.nonScheduledRequests = nonScheduledRequests;
     }
 
     @Override
     public void subscribeActual(final Subscriber<? super T> s) {
         Scheduler.Worker w = scheduler.createWorker();
-        final SubscribeOnSubscriber<T> sos = new SubscribeOnSubscriber<T>(s, w, source);
+        final SubscribeOnSubscriber<T> sos = new SubscribeOnSubscriber<T>(s, w, source, nonScheduledRequests);
         s.onSubscribe(sos);
 
         w.schedule(sos);
@@ -42,21 +52,26 @@ static final class SubscribeOnSubscriber<T> extends AtomicReference<Thread>
     implements Subscriber<T>, Subscription, Runnable {
 
         private static final long serialVersionUID = 8094547886072529208L;
+
         final Subscriber<? super T> actual;
+
         final Scheduler.Worker worker;
 
         final AtomicReference<Subscription> s;
 
         final AtomicLong requested;
 
+        final boolean nonScheduledRequests;
+
         Publisher<T> source;
 
-        SubscribeOnSubscriber(Subscriber<? super T> actual, Scheduler.Worker worker, Publisher<T> source) {
+        SubscribeOnSubscriber(Subscriber<? super T> actual, Scheduler.Worker worker, Publisher<T> source, boolean nonScheduledRequests) {
             this.actual = actual;
             this.worker = worker;
             this.source = source;
             this.s = new AtomicReference<Subscription>();
             this.requested = new AtomicLong();
+            this.nonScheduledRequests = nonScheduledRequests;
         }
 
         @Override
@@ -114,7 +129,7 @@ public void request(final long n) {
         }
 
         void requestUpstream(final long n, final Subscription s) {
-            if (Thread.currentThread() == get()) {
+            if (nonScheduledRequests || Thread.currentThread() == get()) {
                 s.request(n);
             } else {
                 worker.schedule(new Runnable() {

src/test/java/io/reactivex/internal/operators/flowable/FlowableSubscribeOnTest.java

@@ -24,6 +24,7 @@
 import io.reactivex.*;
 import io.reactivex.Scheduler.Worker;
 import io.reactivex.disposables.Disposable;
+import io.reactivex.internal.functions.Functions;
 import io.reactivex.internal.operators.flowable.FlowableSubscribeOn.SubscribeOnSubscriber;
 import io.reactivex.internal.subscriptions.BooleanSubscription;
 import io.reactivex.schedulers.*;
@@ -295,7 +296,7 @@ public void deferredRequestRace() {
 
             Worker w = Schedulers.computation().createWorker();
 
-            final SubscribeOnSubscriber<Integer> so = new SubscribeOnSubscriber<Integer>(ts, w, Flowable.<Integer>never());
+            final SubscribeOnSubscriber<Integer> so = new SubscribeOnSubscriber<Integer>(ts, w, Flowable.<Integer>never(), true);
             ts.onSubscribe(so);
 
             final BooleanSubscription bs = new BooleanSubscription();
@@ -321,4 +322,50 @@ public void run() {
             }
         }
     }
+
+    @Test
+    public void nonScheduledRequests() {
+        TestSubscriber<Object> ts = Flowable.create(new FlowableOnSubscribe<Object>() {
+            @Override
+            public void subscribe(FlowableEmitter<Object> s) throws Exception {
+                for (int i = 1; i < 1001; i++) {
+                    s.onNext(i);
+                    Thread.sleep(1);
+                }
+                s.onComplete();
+            }
+        }, BackpressureStrategy.DROP)
+        .subscribeOn(Schedulers.single())
+        .observeOn(Schedulers.computation())
+        .test()
+        .awaitDone(5, TimeUnit.SECONDS)
+        .assertNoErrors()
+        .assertComplete();
+
+        int c = ts.valueCount();
+
+        assertTrue("" + c, c > Flowable.bufferSize());
+    }
+
+    @Test
+    public void scheduledRequests() {
+        Flowable.create(new FlowableOnSubscribe<Object>() {
+            @Override
+            public void subscribe(FlowableEmitter<Object> s) throws Exception {
+                for (int i = 1; i < 1001; i++) {
+                    s.onNext(i);
+                    Thread.sleep(1);
+                }
+                s.onComplete();
+            }
+        }, BackpressureStrategy.DROP)
+        .map(Functions.identity())
+        .subscribeOn(Schedulers.single())
+        .observeOn(Schedulers.computation())
+        .test()
+        .awaitDone(5, TimeUnit.SECONDS)
+        .assertValueCount(Flowable.bufferSize())
+        .assertNoErrors()
+        .assertComplete();
+    }
 }