Merge pull request #1573 from benjchristensen/parallel-backpressure

Backpressure: parallel

Author: benjchristensen

rxjava-core/src/main/java/rx/internal/operators/OperatorParallel.java

@@ -15,14 +15,16 @@
  */
 package rx.internal.operators;
 
-import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
 
 import rx.Observable;
+import rx.Observable.OnSubscribe;
 import rx.Observable.Operator;
+import rx.Producer;
 import rx.Scheduler;
 import rx.Subscriber;
 import rx.functions.Func1;
-import rx.subjects.Subject;
 
 /**
  * Identifies unit of work that can be executed in parallel on a given Scheduler.
@@ -41,85 +43,99 @@ public OperatorParallel(Func1<Observable<T>, Observable<R>> f, Scheduler schedul
 
     @Override
     public Subscriber<? super T> call(final Subscriber<? super R> child) {
-
-        @SuppressWarnings("unchecked")
-        final UnicastPassThruSubject<T>[] subjects = new UnicastPassThruSubject[degreeOfParallelism];
         @SuppressWarnings("unchecked")
         final Observable<R>[] os = new Observable[degreeOfParallelism];
-        for (int i = 0; i < subjects.length; i++) {
-            subjects[i] = UnicastPassThruSubject.<T> create();
-            os[i] = f.call(subjects[i].observeOn(scheduler));
-        }
-
-        // subscribe BEFORE receiving data so everything is hooked up
-        Observable.merge(os).unsafeSubscribe(child);
-
-        return new Subscriber<T>(child) {
-
-            int index = 0; // trust that we receive data synchronously
-
-            @Override
-            public void onCompleted() {
-                for (UnicastPassThruSubject<T> s : subjects) {
-                    s.onCompleted();
-                }
-            }
-
-            @Override
-            public void onError(Throwable e) {
-                // bypass the subjects and immediately terminate
-                child.onError(e);
-            }
+        @SuppressWarnings("unchecked")
+        final Subscriber<? super T>[] ss = new Subscriber[degreeOfParallelism];
+        final ParentSubscriber subscriber = new ParentSubscriber(child, ss);
+        for (int i = 0; i < os.length; i++) {
+            final int index = i;
+            Observable<T> o = Observable.create(new OnSubscribe<T>() {
 
-            @Override
-            public void onNext(T t) {
-                // round-robin subjects
-                subjects[index++].onNext(t);
-                if (index >= degreeOfParallelism) {
-                    index = 0;
+                @Override
+                public void call(Subscriber<? super T> inner) {
+                    ss[index] = inner;
+                    child.add(inner); // unsubscribe chain
+                    inner.setProducer(new Producer() {
+
+                        @Override
+                        public void request(long n) {
+                            // as we receive requests from observeOn propagate upstream to the parent Subscriber
+                            subscriber.requestMore(n);
+                        }
+
+                    });
                 }
-            }
 
-        };
+            });
+            os[i] = f.call(o.observeOn(scheduler));
+        }
 
+        // subscribe BEFORE receiving data so everything is hooked up
+        Observable.merge(os).unsafeSubscribe(child);
+        return subscriber;
     }
 
-    private static class UnicastPassThruSubject<T> extends Subject<T, T> {
-
-        private static <T> UnicastPassThruSubject<T> create() {
-            final AtomicReference<Subscriber<? super T>> subscriber = new AtomicReference<Subscriber<? super T>>();
-            return new UnicastPassThruSubject<T>(subscriber, new OnSubscribe<T>() {
+    private class ParentSubscriber extends Subscriber<T> {
 
-                @Override
-                public void call(Subscriber<? super T> s) {
-                    subscriber.set(s);
-                }
-
-            });
+        final Subscriber<? super R> child;
+        final Subscriber<? super T>[] ss;
+        int index = 0;
+        final AtomicLong initialRequest = new AtomicLong();
+        final AtomicBoolean started = new AtomicBoolean();
 
+        private ParentSubscriber(Subscriber<? super R> child, Subscriber<? super T>[] ss) {
+            super(child);
+            this.child = child;
+            this.ss = ss;
         }
 
-        private final AtomicReference<Subscriber<? super T>> subscriber;
+        public void requestMore(long n) {
+            if (started.get()) {
+                request(n);
+            } else {
+                initialRequest.addAndGet(n);
+            }
+        }
 
-        protected UnicastPassThruSubject(AtomicReference<Subscriber<? super T>> subscriber, OnSubscribe<T> onSubscribe) {
-            super(onSubscribe);
-            this.subscriber = subscriber;
+        @Override
+        public void onStart() {
+            if (started.compareAndSet(false, true)) {
+                // if no request via requestMore has been sent yet, we start with 0 (rather than default Long.MAX_VALUE).
+                request(initialRequest.get());
+            }
         }
 
         @Override
         public void onCompleted() {
-            subscriber.get().onCompleted();
+            for (Subscriber<? super T> s : ss) {
+                s.onCompleted();
+            }
         }
 
         @Override
         public void onError(Throwable e) {
-            subscriber.get().onError(e);
+            child.onError(e);
         }
 
         @Override
         public void onNext(T t) {
-            subscriber.get().onNext(t);
+            /*
+             * There is a possible bug here ... we could get a MissingBackpressureException
+             * if the processing on each of the threads is unbalanced. In other words, if 1 of the
+             * observeOn queues has space, but another is full, this could try emitting to one that
+             * is full and get a MissingBackpressureException.
+             * 
+             * To solve that we'd need to check the outstanding request per Subscriber, which will
+             * need a more complicated mechanism to expose a type that has both the requested + the
+             * Subscriber to emit to.
+             */
+            ss[index++].onNext(t);
+            if (index >= degreeOfParallelism) {
+                index = 0;
+            }
         }
 
-    }
+    };
+
 }

rxjava-core/src/test/java/rx/internal/operators/OperatorParallelTest.java

@@ -21,7 +21,6 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
-import org.junit.Ignore;
 import org.junit.Test;
 
 import rx.Observable;
@@ -57,7 +56,7 @@ public Integer[] call(Integer t) {
                                     // TODO why is this exception not being thrown?
                                     throw new RuntimeException(e);
                                 }
-                                //                                                                System.out.println("V: " + t  + " Thread: " + Thread.currentThread());
+                                //                                                                                                System.out.println("V: " + t  + " Thread: " + Thread.currentThread());
                                 innerCount.incrementAndGet();
                                 return new Integer[] { t, t * 99 };
                             }
@@ -112,8 +111,6 @@ public void call(String v) {
         assertEquals(NUM, count.get());
     }
 
-    // parallel does not support backpressure right now
-    @Ignore
     @Test
     public void testBackpressureViaOuterObserveOn() {
         final AtomicInteger emitted = new AtomicInteger();
@@ -148,12 +145,10 @@ public String call(Integer t) {
         ts.awaitTerminalEvent();
         ts.assertNoErrors();
         System.out.println("testBackpressureViaObserveOn emitted => " + emitted.get());
-        assertTrue(emitted.get() < 2000 + RxRingBuffer.SIZE); // should have no more than the buffer size beyond the 2000 in take
-        assertEquals(2000, ts.getOnNextEvents().size());
+        assertTrue(emitted.get() < 20000 + (RxRingBuffer.SIZE * Schedulers.computation().parallelism())); // should have no more than the buffer size beyond the 20000 in take
+        assertEquals(20000, ts.getOnNextEvents().size());
     }
 
-    // parallel does not support backpressure right now
-    @Ignore
     @Test
     public void testBackpressureOnInnerObserveOn() {
         final AtomicInteger emitted = new AtomicInteger();
@@ -188,8 +183,8 @@ public String call(Integer t) {
         ts.awaitTerminalEvent();
         ts.assertNoErrors();
         System.out.println("testBackpressureViaObserveOn emitted => " + emitted.get());
-        assertTrue(emitted.get() < 20000 + RxRingBuffer.SIZE); // should have no more than the buffer size beyond the 2000 in take
-        assertEquals(2000, ts.getOnNextEvents().size());
+        assertTrue(emitted.get() < 20000 + (RxRingBuffer.SIZE * Schedulers.computation().parallelism())); // should have no more than the buffer size beyond the 20000 in take
+        assertEquals(20000, ts.getOnNextEvents().size());
     }
 
     @Test(timeout = 10000)
@@ -226,7 +221,8 @@ public String call(Integer t) {
         ts.awaitTerminalEvent();
         ts.assertNoErrors();
         System.out.println("emitted: " + emitted.get());
-        assertEquals(2000, emitted.get()); // no async, so should be perfect
+        // we allow buffering inside each parallel Observable
+        assertEquals(RxRingBuffer.SIZE * Schedulers.computation().parallelism(), emitted.get()); // no async, so should be perfect
         assertEquals(2000, ts.getOnNextEvents().size());
     }
 }