Merge pull request #1533 from benjchristensen/issue-1516

Amb + Backpressure

Author: benjchristensen

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

@@ -16,12 +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.
@@ -262,22 +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);
+        }
+
+        private final void requestMore(long n) {
+            request(n);
         }
 
         @Override
         public void onNext(T args) {
             if (!isSelected()) {
-                unsubscribe();
                 return;
             }
             subscriber.onNext(args);
@@ -286,7 +292,6 @@ public void onNext(T args) {
         @Override
         public void onCompleted() {
             if (!isSelected()) {
-                unsubscribe();
                 return;
             }
             subscriber.onCompleted();
@@ -295,44 +300,102 @@ 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) {
-        AtomicInteger choice = new AtomicInteger(AmbSubscriber.NONE);
-        int index = 0;
-        for (Observable<? extends T> source : sources) {
-            if (subscriber.isUnsubscribed()) {
-                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();
+                }
             }
-            if (choice.get() != AmbSubscriber.NONE) {
-                // Already choose someone, the rest Observables can be skipped.
-                break;
+            
+        }));
+        subscriber.setProducer(new Producer() {
+
+            @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 {
+                    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);
+                    }
+                }
             }
-            AmbSubscriber<T> ambSubscriber = new AmbSubscriber<T>(subscriber, index, choice);
-            subscriber.add(ambSubscriber);
-            source.unsafeSubscribe(ambSubscriber);
-            index++;
-        }
+        });
     }
 
 }

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,13 @@
 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.internal.util.RxRingBuffer;
+import rx.observers.TestSubscriber;
+import rx.schedulers.Schedulers;
 import rx.schedulers.TestScheduler;
 import rx.subscriptions.CompositeSubscription;
 
@@ -157,4 +163,60 @@ 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) {
+                        System.out.println("1-requested: " + 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) {
+                        System.out.println("2-requested: " + n);
+                        requested2.set(n);
+                    }
+
+                });
+            }
+
+        });
+        Observable.amb(o1, o2).subscribe(ts);
+        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());
+    }
 }