fix Amb backpressure problems

Author: davidmoten

src/main/java/rx/internal/operators/OnSubscribeAmb.java

@@ -19,6 +19,7 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 
 import rx.Observable;
@@ -371,6 +372,21 @@ public void call() {
             }
 
         }));
+        //need to subscribe to all the sources
+        for (Observable<? extends T> source : sources) {
+            if (subscriber.isUnsubscribed()) {
+                return;
+            }
+            AmbSubscriber<T> ambSubscriber = new AmbSubscriber<T>(0, subscriber, selection);
+            selection.ambSubscribers.add(ambSubscriber);
+            // check again if choice has been made so can stop subscribing
+            if (selection.choice.get() != null) {
+                // Already chose one, the rest can be skipped and we can clean up
+                selection.unsubscribeOthers(selection.choice.get());
+                return;
+            }
+            source.unsafeSubscribe(ambSubscriber);
+        }
         subscriber.setProducer(new Producer() {
 
             @Override
@@ -379,19 +395,21 @@ public void request(long n) {
                     // propagate the request to that single Subscriber that won
                     selection.choice.get().requestMore(n);
                 } else {
-                    for (Observable<? extends T> source : sources) {
-                        if (subscriber.isUnsubscribed()) {
-                            break;
-                        }
-                        AmbSubscriber<T> ambSubscriber = new AmbSubscriber<T>(n, subscriber, selection);
-                        selection.ambSubscribers.add(ambSubscriber);
-                        // possible race condition in previous lines ... a choice may have been made so double check (instead of synchronizing)
-                        if (selection.choice.get() != null) {
-                            // Already chose one, the rest can be skipped and we can clean up
-                            selection.unsubscribeOthers(selection.choice.get());
-                            break;
+                    //subscriptions already happened so propagate the request to all the  
+                    //amb subscribers
+                    for (AmbSubscriber<T> ambSubscriber: selection.ambSubscribers) {
+                        if (!ambSubscriber.isUnsubscribed()) {
+                            // make a best endeavours check to not waste requests 
+                            // if first emission has already occurred
+                            if (selection.choice.get() == ambSubscriber) {
+                                ambSubscriber.requestMore(n);
+                                //don't need to request from other subscribers because choice has been made
+                                //and request has gone to choice 
+                                return;
+                            } else {
+                                ambSubscriber.requestMore(n);
+                            }
                         }
-                        source.unsafeSubscribe(ambSubscriber);
                     }
                 }
             }

src/test/java/rx/internal/operators/OnSubscribeAmbTest.java

@@ -22,6 +22,7 @@
 import static rx.internal.operators.OnSubscribeAmb.amb;
 
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 
@@ -219,4 +220,49 @@ public void testBackpressure() {
         ts.assertNoErrors();
         assertEquals(RxRingBuffer.SIZE * 2, ts.getOnNextEvents().size());
     }
+    
+    
+    @Test
+    public void testSubscriptionOnlyHappensOnce() throws InterruptedException {
+        final AtomicLong count = new AtomicLong();
+        Action0 incrementer = new Action0() {
+            @Override
+            public void call() {
+                count.incrementAndGet();
+            }
+        };
+        //this aync stream should emit first
+        Observable<Integer> o1 = Observable.just(1).doOnSubscribe(incrementer)
+                .delay(100, TimeUnit.MILLISECONDS).subscribeOn(Schedulers.computation());
+        //this stream emits second
+        Observable<Integer> o2 = Observable.just(1).doOnSubscribe(incrementer)
+                .delay(100, TimeUnit.MILLISECONDS).subscribeOn(Schedulers.computation());
+        TestSubscriber<Integer> ts = new TestSubscriber<Integer>();
+        Observable.amb(o1, o2).subscribe(ts);
+        ts.requestMore(1);
+        ts.awaitTerminalEvent(5, TimeUnit.SECONDS);
+        ts.assertNoErrors();
+        assertEquals(2, count.get());
+    }
+    
+    @Test
+    public void testSecondaryRequestsPropagatedToChildren() throws InterruptedException {
+        //this aync stream should emit first
+        Observable<Integer> o1 = Observable.from(Arrays.asList(1, 2, 3))
+                .delay(100, TimeUnit.MILLISECONDS).subscribeOn(Schedulers.computation());
+        //this stream emits second
+        Observable<Integer> o2 = Observable.from(Arrays.asList(4, 5, 6))
+                .delay(200, TimeUnit.MILLISECONDS).subscribeOn(Schedulers.computation());
+        TestSubscriber<Integer> ts = new TestSubscriber<Integer>();
+        //before subscription request 1
+        ts.requestMore(1);
+        Observable.amb(o1, o2).subscribe(ts);
+        // before first emission request 20 more
+        // this request should suffice to emit all
+        ts.requestMore(20);
+        //ensure stream does not hang
+        ts.awaitTerminalEvent(5, TimeUnit.SECONDS);
+        ts.assertNoErrors();
+    }
+
 }