Propagate request(n) to each AmbSubscriber

... and then keep propagating to the winner.

Author: benjchristensen

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

@@ -16,13 +16,17 @@
 package rx.internal.operators;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
-import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicReference;
 
 import rx.Observable;
 import rx.Observable.OnSubscribe;
 import rx.Producer;
 import rx.Subscriber;
+import rx.functions.Action0;
+import rx.subscriptions.Subscriptions;
 
 /**
  * Given multiple {@link Observable}s, propagates the one that first emits an item.
@@ -263,26 +267,23 @@ public static <T> OnSubscribe<T> amb(final Iterable<? extends Observable<? exten
 
     private static final class AmbSubscriber<T> extends Subscriber<T> {
 
-        private static final int NONE = -1;
-
         private final Subscriber<? super T> subscriber;
-        private final int index;
-        private final AtomicInteger choice;
+        private final Selection<T> selection;
 
-        private AmbSubscriber(Subscriber<? super T> subscriber, int index, AtomicInteger choice) {
+        private AmbSubscriber(long requested, Subscriber<? super T> subscriber, Selection<T> selection) {
             this.subscriber = subscriber;
-            this.choice = choice;
-            this.index = index;
+            this.selection = selection;
+            // initial request
+            request(requested);
         }
 
-        public void requestMore(long n) {
+        private final void requestMore(long n) {
             request(n);
         }
 
         @Override
         public void onNext(T args) {
             if (!isSelected()) {
-                unsubscribe();
                 return;
             }
             subscriber.onNext(args);
@@ -291,7 +292,6 @@ public void onNext(T args) {
         @Override
         public void onCompleted() {
             if (!isSelected()) {
-                unsubscribe();
                 return;
             }
             subscriber.onCompleted();
@@ -300,66 +300,99 @@ public void onCompleted() {
         @Override
         public void onError(Throwable e) {
             if (!isSelected()) {
-                unsubscribe();
                 return;
             }
             subscriber.onError(e);
         }
 
         private boolean isSelected() {
-            int ch = choice.get();
-            if (ch == NONE) {
-                return choice.compareAndSet(NONE, index);
+            if (selection.choice.get() == this) {
+                // fast-path
+                return true;
+            } else {
+                if (selection.choice.compareAndSet(null, this)) {
+                    selection.unsubscribeOthers(this);
+                    return true;
+                } else {
+                    // we lost so unsubscribe ... and force cleanup again due to possible race conditions
+                    selection.unsubscribeLosers();
+                    return false;
+                }
+            }
+        }
+    }
+
+    private static class Selection<T> {
+        final AtomicReference<AmbSubscriber<T>> choice = new AtomicReference<AmbSubscriber<T>>();
+        final Collection<AmbSubscriber<T>> ambSubscribers = new ConcurrentLinkedQueue<AmbSubscriber<T>>();
+
+        public void unsubscribeLosers() {
+            AmbSubscriber<T> winner = choice.get();
+            if(winner != null) {
+                unsubscribeOthers(winner);
+            }
+        }
+        
+        public void unsubscribeOthers(AmbSubscriber<T> notThis) {
+            for (AmbSubscriber<T> other : ambSubscribers) {
+                if (other != notThis) {
+                    other.unsubscribe();
+                }
             }
-            return ch == index;
+            ambSubscribers.clear();
         }
+
     }
 
     private final Iterable<? extends Observable<? extends T>> sources;
+    private final Selection<T> selection = new Selection<T>();
 
     private OnSubscribeAmb(Iterable<? extends Observable<? extends T>> sources) {
         this.sources = sources;
     }
 
     @Override
-    public void call(Subscriber<? super T> subscriber) {
-        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()) {
-                break;
-            }
-            if (choice.get() != AmbSubscriber.NONE) {
-                // Already choose someone, the rest Observables can be skipped.
-                break;
+    public void call(final Subscriber<? super T> subscriber) {
+        subscriber.add(Subscriptions.create(new Action0() {
+
+            @Override
+            public void call() {
+                if (selection.choice.get() != null) {
+                    // there is a single winner so we unsubscribe it
+                    selection.choice.get().unsubscribe();
+                } 
+                // if we are racing with others still existing, we'll also unsubscribe them
+                if(!selection.ambSubscribers.isEmpty()) {
+                    for (AmbSubscriber<T> other : selection.ambSubscribers) {
+                        other.unsubscribe();
+                    }
+                    selection.ambSubscribers.clear();
+                }
             }
-            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());
+                if (selection.choice.get() != null) {
+                    // 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;
+                        }
+                        source.unsafeSubscribe(ambSubscriber);
                     }
-                    selectedAmbSubscriber.requestMore(n);
                 }
             }
         });

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

@@ -36,7 +36,9 @@
 import rx.Scheduler;
 import rx.Subscriber;
 import rx.functions.Action0;
+import rx.internal.util.RxRingBuffer;
 import rx.observers.TestSubscriber;
+import rx.schedulers.Schedulers;
 import rx.schedulers.TestScheduler;
 import rx.subscriptions.CompositeSubscription;
 
@@ -175,6 +177,7 @@ public void call(Subscriber<? super Integer> s) {
 
                     @Override
                     public void request(long n) {
+                        System.out.println("1-requested: " + n);
                         requested1.set(n);
                     }
 
@@ -190,6 +193,7 @@ public void call(Subscriber<? super Integer> s) {
 
                     @Override
                     public void request(long n) {
+                        System.out.println("2-requested: " + n);
                         requested2.set(n);
                     }
 
@@ -201,4 +205,18 @@ public void request(long n) {
         assertEquals(3, requested1.get());
         assertEquals(3, requested2.get());
     }
+
+    @Test
+    public void testBackpressure() {
+        TestSubscriber<Integer> ts = new TestSubscriber<Integer>();
+        Observable.range(0, RxRingBuffer.SIZE * 2)
+                .ambWith(Observable.range(0, RxRingBuffer.SIZE * 2))
+                .observeOn(Schedulers.computation()) // observeOn has a backpressured RxRingBuffer
+                .delay(1, TimeUnit.MICROSECONDS) // make it a slightly slow consumer
+                .subscribe(ts);
+
+        ts.awaitTerminalEvent();
+        ts.assertNoErrors();
+        assertEquals(RxRingBuffer.SIZE * 2, ts.getOnNextEvents().size());
+    }
 }