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;
@@ -373,26 +374,46 @@ public void call() {
         }));
         subscriber.setProducer(new Producer() {
 
+            private final AtomicBoolean subscribedToSources = new AtomicBoolean(false);
+            
             @Override
             public void request(long n) {
                 if (selection.choice.get() != null) {
                     // propagate the request to that single Subscriber that won
                     selection.choice.get().requestMore(n);
-                } else {
+                } else if (subscribedToSources.compareAndSet(false, true)){
+                    //need to subscribe to all the sources
                     for (Observable<? extends T> source : sources) {
                         if (subscriber.isUnsubscribed()) {
-                            break;
+                            return;
                         }
                         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)
+                        // 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());
-                            break;
+                            return;
                         }
                         source.unsafeSubscribe(ambSubscriber);
                     }
+                } else {
+                    //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);
+                            }
+                        }
+                    }
                 }
             }
         });

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

@@ -16,12 +16,15 @@
 package rx.internal.operators;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 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.Arrays;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 
@@ -219,4 +222,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 testRequestsPropagatedToChildren() 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();
+    }
+
 }