Amb with backpressure support

Author: zsxwing

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

@@ -21,6 +21,7 @@
 
 import rx.Observable;
 import rx.Observable.OnSubscribe;
+import rx.Producer;
 import rx.Subscriber;
 
 /**
@@ -274,6 +275,10 @@ private AmbSubscriber(Subscriber<? super T> subscriber, int index, AtomicInteger
             this.index = index;
         }
 
+        public void requestMore(long n) {
+            request(n);
+        }
+
         @Override
         public void onNext(T args) {
             if (!isSelected()) {
@@ -318,7 +323,8 @@ private OnSubscribeAmb(Iterable<? extends Observable<? extends T>> sources) {
 
     @Override
     public void call(Subscriber<? super T> subscriber) {
-        AtomicInteger choice = new AtomicInteger(AmbSubscriber.NONE);
+        final AtomicInteger choice = new AtomicInteger(AmbSubscriber.NONE);
+        final List<AmbSubscriber<T>> ambSubscribers = new ArrayList<AmbSubscriber<T>>();
         int index = 0;
         for (Observable<? extends T> source : sources) {
             if (subscriber.isUnsubscribed()) {
@@ -329,10 +335,34 @@ public void call(Subscriber<? super T> subscriber) {
                 break;
             }
             AmbSubscriber<T> ambSubscriber = new AmbSubscriber<T>(subscriber, index, choice);
+            ambSubscribers.add(ambSubscriber);
             subscriber.add(ambSubscriber);
             source.unsafeSubscribe(ambSubscriber);
             index++;
         }
+        // setProducer at the end so that `ambSubscribers` can be finalized before `subscriber` calls `request`
+        subscriber.setProducer(new Producer() {
+
+            private volatile AmbSubscriber<T> selectedAmbSubscriber;
+
+            @Override
+            public void request(long n) {
+                if (choice.get() == AmbSubscriber.NONE) {
+                    for (AmbSubscriber<T> ambSubscriber : ambSubscribers) {
+                        // Once one Observable emits a message, `unsubscribe` of other Observables will be called
+                        // and further messages will be dropped. Therefore, requesting all sources won't cause
+                        // the backpressure issue.
+                        ambSubscriber.requestMore(n);
+                    }
+                }
+                else {
+                    if (selectedAmbSubscriber == null) {
+                        selectedAmbSubscriber = ambSubscribers.get(choice.get());
+                    }
+                    selectedAmbSubscriber.requestMore(n);
+                }
+            }
+        });
     }
 
 }

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

@@ -15,13 +15,15 @@
  */
 package rx.internal.operators;
 
+import static org.junit.Assert.assertEquals;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static rx.internal.operators.OnSubscribeAmb.amb;
 
 import java.io.IOException;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -30,9 +32,11 @@
 import rx.Observable;
 import rx.Observable.OnSubscribe;
 import rx.Observer;
+import rx.Producer;
 import rx.Scheduler;
 import rx.Subscriber;
 import rx.functions.Action0;
+import rx.observers.TestSubscriber;
 import rx.schedulers.TestScheduler;
 import rx.subscriptions.CompositeSubscription;
 
@@ -157,4 +161,44 @@ public void testAmb3() {
         inOrder.verifyNoMoreInteractions();
     }
 
+    @Test
+    public void testProducerRequestThroughAmb() {
+        TestSubscriber<Integer> ts = new TestSubscriber<Integer>();
+        ts.requestMore(3);
+        final AtomicLong requested1 = new AtomicLong();
+        final AtomicLong requested2 = new AtomicLong();
+        Observable<Integer> o1 = Observable.create(new OnSubscribe<Integer>() {
+
+            @Override
+            public void call(Subscriber<? super Integer> s) {
+                s.setProducer(new Producer() {
+
+                    @Override
+                    public void request(long n) {
+                        requested1.set(n);
+                    }
+
+                });
+            }
+
+        });
+        Observable<Integer> o2 = Observable.create(new OnSubscribe<Integer>() {
+
+            @Override
+            public void call(Subscriber<? super Integer> s) {
+                s.setProducer(new Producer() {
+
+                    @Override
+                    public void request(long n) {
+                        requested2.set(n);
+                    }
+
+                });
+            }
+
+        });
+        Observable.amb(o1, o2).subscribe(ts);
+        assertEquals(3, requested1.get());
+        assertEquals(3, requested2.get());
+    }
 }