2.x: Cleanup Observable.flatMap drain logic (#6232)

Author: akarnokd

src/main/java/io/reactivex/internal/operators/observable/ObservableFlatMap.java

@@ -338,23 +338,17 @@ void drainLoop() {
 
                 if (svq != null) {
                     for (;;) {
-                        U o;
-                        for (;;) {
-                            if (checkTerminate()) {
-                                return;
-                            }
-
-                            o = svq.poll();
+                        if (checkTerminate()) {
+                            return;
+                        }
 
-                            if (o == null) {
-                                break;
-                            }
+                        U o = svq.poll();
 
-                            child.onNext(o);
-                        }
                         if (o == null) {
                             break;
                         }
+
+                        child.onNext(o);
                     }
                 }
 
@@ -415,17 +409,10 @@ void drainLoop() {
 
                         @SuppressWarnings("unchecked")
                         InnerObserver<T, U> is = (InnerObserver<T, U>)inner[j];
-
-                        for (;;) {
-                            if (checkTerminate()) {
-                                return;
-                            }
-                            SimpleQueue<U> q = is.queue;
-                            if (q == null) {
-                                break;
-                            }
-                            U o;
+                        SimpleQueue<U> q = is.queue;
+                        if (q != null) {
                             for (;;) {
+                                U o;
                                 try {
                                     o = q.poll();
                                 } catch (Throwable ex) {
@@ -437,7 +424,10 @@ void drainLoop() {
                                     }
                                     removeInner(is);
                                     innerCompleted = true;
-                                    i++;
+                                    j++;
+                                    if (j == n) {
+                                        j = 0;
+                                    }
                                     continue sourceLoop;
                                 }
                                 if (o == null) {
@@ -450,10 +440,8 @@ void drainLoop() {
                                     return;
                                 }
                             }
-                            if (o == null) {
-                                break;
-                            }
                         }
+
                         boolean innerDone = is.done;
                         SimpleQueue<U> innerQueue = is.queue;
                         if (innerDone && (innerQueue == null || innerQueue.isEmpty())) {

src/test/java/io/reactivex/internal/operators/observable/ObservableFlatMapTest.java

@@ -26,14 +26,14 @@
 import io.reactivex.*;
 import io.reactivex.Observable;
 import io.reactivex.Observer;
-import io.reactivex.disposables.Disposable;
+import io.reactivex.disposables.*;
 import io.reactivex.exceptions.*;
 import io.reactivex.functions.*;
 import io.reactivex.internal.functions.Functions;
 import io.reactivex.observers.TestObserver;
 import io.reactivex.plugins.RxJavaPlugins;
 import io.reactivex.schedulers.Schedulers;
-import io.reactivex.subjects.PublishSubject;
+import io.reactivex.subjects.*;
 
 public class ObservableFlatMapTest {
     @Test
@@ -1006,4 +1006,43 @@ public void onNext(Integer t) {
 
         to.assertResult(1, 2, 3, 4, 5, 6, 7, 8, 9, 10);
     }
+
+    @Test
+    public void fusedSourceCrashResumeWithNextSource() {
+        final UnicastSubject<Integer> fusedSource = UnicastSubject.create();
+        TestObserver<Integer> to = new TestObserver<Integer>();
+
+        ObservableFlatMap.MergeObserver<Integer, Integer> merger =
+                new ObservableFlatMap.MergeObserver<Integer, Integer>(to, new Function<Integer, Observable<Integer>>() {
+                    @Override
+                    public Observable<Integer> apply(Integer t)
+                            throws Exception {
+                        if (t == 0) {
+                            return fusedSource
+                                    .map(new Function<Integer, Integer>() {
+                                        @Override
+                                        public Integer apply(Integer v)
+                                                throws Exception { throw new TestException(); }
+                                    })
+                                    .compose(TestHelper.<Integer>observableStripBoundary());
+                        }
+                        return Observable.range(10 * t, 5);
+                    }
+                }, true, Integer.MAX_VALUE, 128);
+
+        merger.onSubscribe(Disposables.empty());
+        merger.getAndIncrement();
+
+        merger.onNext(0);
+        merger.onNext(1);
+        merger.onNext(2);
+
+        assertTrue(fusedSource.hasObservers());
+
+        fusedSource.onNext(-1);
+
+        merger.drainLoop();
+
+        to.assertValuesOnly(10, 11, 12, 13, 14, 20, 21, 22, 23, 24);
+    }
 }