Implemented the AsyncOnSubscribe

Author: Aaron Tull

src/main/java/rx/observables/AsyncOnSubscribe.java

@@ -0,0 +1,270 @@
+/**
+ * Copyright 2014 Netflix, Inc.
+ * 
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * 
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * 
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package rx.observables;
+
+import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
+
+import rx.Observable;
+import rx.Observable.OnSubscribe;
+import rx.Observer;
+import rx.Producer;
+import rx.Subscriber;
+import rx.Subscription;
+import rx.functions.Action0;
+import rx.internal.operators.BackpressureUtils;
+import rx.internal.operators.BufferUntilSubscriber;
+import rx.observers.SerializedObserver;
+import rx.plugins.RxJavaPlugins;
+import rx.subscriptions.BooleanSubscription;
+
+public abstract class AsyncOnSubscribe<S, T> implements OnSubscribe<T> {
+    
+    /**
+     * Executed once when subscribed to by a subscriber (via {@link OnSubscribe#call(Subscriber)})
+     * to produce a state value. This value is passed into {@link #next(Object, Observer) next(S
+     * state, Observer <T> observer)} on the first iteration. Subsequent iterations of {@code next}
+     * will receive the state returned by the previous invocation of {@code next}.
+     * 
+     * @return the initial state value
+     */
+    protected abstract S generateState();
+
+    /**
+     * Called to produce data to the downstream subscribers. To emit data to a downstream subscriber
+     * call {@code observer.onNext(t)}. To signal an error condition call
+     * {@code observer.onError(throwable)} or throw an Exception. To signal the end of a data stream
+     * call {@code 
+     * observer.onCompleted()}. Implementations of this method must follow the following rules.
+     * 
+     * <ul>
+     * <li>Must not call {@code observer.onNext(t)} more than 1 time per invocation.</li>
+     * <li>Must not call {@code observer.onNext(t)} concurrently.</li>
+     * </ul>
+     * 
+     * The value returned from an invocation of this method will be passed in as the {@code state}
+     * argument of the next invocation of this method.
+     * 
+     * @param state
+     *            the state value (from {@link #generateState()} on the first invocation or the
+     *            previous invocation of this method.
+     * @param observer
+     *            the observer of data emitted by
+     * @return the next iteration's state value
+     */
+    protected abstract S next(S state, long requested, Observer<Observable<T>> observer);
+    
+    /**
+     * Clean up behavior that is executed after the downstream subscriber's subscription is
+     * unsubscribed. This method will be invoked exactly once.
+     * 
+     * @param state
+     *            the last state value prior from {@link #generateState()} or
+     *            {@link #next(Object, Observer)} before unsubscribe.
+     */
+    protected void onUnsubscribe(S state) {
+
+    }
+    
+    @Override
+    public final void call(Subscriber<? super T> actualSubscriber) {
+        S state = generateState();
+        UnicastSubject<Observable<T>> unicast = UnicastSubject.<Observable<T>>create();
+        Observable<T> concat = Observable.concat(unicast);
+        AsyncOuterSubscriber<S, T> outerSubscriberProducer = new AsyncOuterSubscriber<S, T>(this, state, unicast);
+        actualSubscriber.add(outerSubscriberProducer);
+        concat.subscribe(actualSubscriber);
+        actualSubscriber.setProducer(outerSubscriberProducer);
+    }
+    
+    private static class AsyncOuterSubscriber<S, T> 
+            extends AtomicLong 
+            implements Producer, Subscription, Observer<Observable<T>> {
+        /** */
+        private static final long serialVersionUID = -7884904861928856832L;
+        
+        private volatile int isUnsubscribed;
+        @SuppressWarnings("rawtypes")
+        private static final AtomicIntegerFieldUpdater<AsyncOuterSubscriber> IS_UNSUBSCRIBED =
+            AtomicIntegerFieldUpdater.newUpdater(AsyncOuterSubscriber.class, "isUnsubscribed");
+
+        private final AsyncOnSubscribe<S, T> parent;
+        private final SerializedObserver<Observable<T>> serializedSubscriber;
+        private final Set<Subscription> subscriptions = new HashSet<Subscription>();
+        
+        private boolean hasTerminated = false;
+        private boolean onNextCalled = false;
+        
+        private S state;
+
+        private final UnicastSubject<Observable<T>> merger;
+
+        public AsyncOuterSubscriber(AsyncOnSubscribe<S, T> parent, S initialState, UnicastSubject<Observable<T>> merger) {
+            this.parent = parent;
+            this.serializedSubscriber = new SerializedObserver<Observable<T>>(this);
+            this.state = initialState;
+            this.merger = merger;
+        }
+
+        @Override
+        public void unsubscribe() {
+            System.out.println("unsub main " + isUnsubscribed);
+            if (IS_UNSUBSCRIBED.compareAndSet(this, 0, 1)) {
+                // it's safe to process terminal behavior
+                System.out.println("subscriptions: " + subscriptions);
+                parent.onUnsubscribe(state);
+                for(Subscription s : subscriptions) {
+                    System.out.println("unsubing: " + s);
+                    if (!s.isUnsubscribed()) {
+                        s.unsubscribe();
+                    }
+                }
+            }
+        }
+
+        @Override
+        public boolean isUnsubscribed() {
+            return isUnsubscribed != 0;
+        }
+
+        public S nextIteration(long n) {
+            return parent.next(state, n, serializedSubscriber);
+        }
+        
+        @Override
+        public void request(long n) {
+            long previousRequestAmount = BackpressureUtils.getAndAddRequest(this, n);
+            if (n > 0 && previousRequestAmount == 0L) {
+                do {
+                    // check if unsubscribed before doing any work
+                    if (isUnsubscribed()) {
+                        unsubscribe();
+                        return;
+                    }
+                    // otherwise try one iteration for a request of `numRequested` elements
+                    try {
+                        onNextCalled = false;
+                        state = nextIteration(n);
+                    } catch (Throwable ex) {
+                        handleThrownError(parent, state, ex);
+                        return;
+                    }
+                } while (!hasTerminated);
+            }
+        }
+
+        private void handleThrownError(final AsyncOnSubscribe<S, T> p, S st, Throwable ex) {
+            if (hasTerminated) {
+                RxJavaPlugins.getInstance().getErrorHandler().handleError(ex);
+            } else {
+                hasTerminated = true;
+                merger.onError(ex);
+                unsubscribe();
+            }
+        }
+        
+        @Override
+        public void onCompleted() {
+            if (hasTerminated) {
+                throw new IllegalStateException("Terminal event already emitted.");
+            }
+            hasTerminated = true;
+            merger.onCompleted();
+        }
+
+        @Override
+        public void onError(Throwable e) {
+            if (hasTerminated) {
+                throw new IllegalStateException("Terminal event already emitted.");
+            }
+            hasTerminated = true;
+            merger.onError(e);
+        }
+        
+        private static Subscription SUBSCRIPTION_SENTINAL = new BooleanSubscription();
+
+        @Override
+        public void onNext(final Observable<T> t) {
+            if (onNextCalled) {
+                throw new IllegalStateException("onNext called multiple times!");
+            }
+            onNextCalled = true;
+            if (hasTerminated)
+                return;
+            BufferUntilSubscriber<T> buffer = BufferUntilSubscriber.<T>create();
+            final AtomicReference<Subscription> holder = new AtomicReference<Subscription>(SUBSCRIPTION_SENTINAL);
+            final Subscription innerSubscription = t.doOnTerminate(new Action0(){
+                @Override
+                public void call() {
+                    if (!holder.compareAndSet(SUBSCRIPTION_SENTINAL, null)) {
+                        Subscription h = holder.get();
+                        System.out.println("Removing subscription " + h);
+                        subscriptions.remove(h);
+                    }
+                }
+            }).subscribe(buffer);
+            if (holder.compareAndSet(SUBSCRIPTION_SENTINAL, innerSubscription)) {
+                System.out.println("t: " + t + " s: "+ innerSubscription);
+                subscriptions.add(innerSubscription);
+            }
+            merger.onNext(buffer);
+        }
+    }
+    
+    private static final class UnicastSubject<T> extends Observable<T> implements Observer<T>  {
+        public static <T> UnicastSubject<T> create() {
+            return new UnicastSubject<T>(new State<T>());
+        }
+
+        private State<T> state;
+
+        protected UnicastSubject(final State<T> state) {
+            super(new OnSubscribe<T>() {
+                @Override
+                public void call(Subscriber<? super T> s) {
+                    if (state.subscriber != null) {
+                        s.onError(new IllegalStateException("Only 1 Subscriber permitted"));
+                    } else {
+                        state.subscriber = s;
+                    }
+                }});
+            this.state = state;
+        }
+
+        @Override
+        public void onCompleted() {
+            state.subscriber.onCompleted();
+        }
+
+        @Override
+        public void onError(Throwable e) {
+            state.subscriber.onError(e);
+        }
+
+        @Override
+        public void onNext(T t) {
+            state.subscriber.onNext(t);
+        }
+        
+        private static class State<T> {
+            private Subscriber<? super T> subscriber;
+        }
+    }    
+}

src/test/java/rx/observables/AsyncOnSubscribeTest.java

@@ -0,0 +1,117 @@
+package rx.observables;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import rx.Observable;
+import rx.Observable.OnSubscribe;
+import rx.Observer;
+import rx.functions.Action0;
+
+public class AsyncOnSubscribeTest {
+
+    @Test
+    public void test() throws InterruptedException {
+        final AtomicBoolean unsubbedFrom1 = new AtomicBoolean(false);
+        final AtomicBoolean unsubbedFrom2 = new AtomicBoolean(false);
+        final AtomicBoolean unsubbedFrom3 = new AtomicBoolean(false);
+        final AtomicBoolean unsubbedFrom4 = new AtomicBoolean(false);
+        final CountDownLatch l1 = new CountDownLatch(1);
+        final CountDownLatch l2 = new CountDownLatch(1);
+        final CountDownLatch l3 = new CountDownLatch(1);
+        final CountDownLatch l4 = new CountDownLatch(1);
+        final AtomicReference<String> failed1 = new AtomicReference<String>();
+        final AtomicReference<String> failed2 = new AtomicReference<String>();
+        final AtomicReference<String> failed3 = new AtomicReference<String>();
+        final AtomicReference<String> failed4 = new AtomicReference<String>();
+        OnSubscribe<Integer> os = new AsyncOnSubscribe<Integer, Integer>() {
+            @Override
+            protected Integer generateState() {
+                return 1;
+            }
+
+            @Override
+            protected Integer next(Integer state, long requested, Observer<Observable<Integer>> observer) {
+                if (state == 1) {
+                    Observable<Integer> o1 = Observable.just(1)
+                            .delay(100, TimeUnit.MILLISECONDS)
+                            .doOnUnsubscribe(new Action0(){
+                                @Override
+                                public void call() {
+                                    unsubbedFrom1.set(true);
+                                    if (l1.getCount() <= 0)
+                                        failed1.set("1st Observable that emits terminal unsubscribed twice");
+                                    l1.countDown();
+                                }});
+                    observer.onNext(o1);
+                } else if (state == 2) {
+                    Observable<Integer> o = Observable.just(2)
+                            .delay(50, TimeUnit.MILLISECONDS)
+                            .doOnUnsubscribe(new Action0(){
+                                @Override
+                                public void call() {
+                                    unsubbedFrom2.set(true);
+                                    if (l2.getCount() <= 0)
+                                        failed2.set("2nd Observable that emits terminal unsubscribed twice");
+                                    l2.countDown();
+                                }});
+                    observer.onNext(o);
+                } else if (state == 3) {
+                    Observable<Integer> o = Observable.<Integer>never()
+                            .doOnUnsubscribe(new Action0(){
+                                @Override
+                                public void call() {
+                                    unsubbedFrom3.set(true);
+                                    if (l3.getCount() <= 0)
+                                        failed3.set("Observable.never() unsubscribed twice");
+                                    l3.countDown();
+                                }});
+                    observer.onNext(o);
+                } else if (state == 4) {
+                    Observable<Integer> o = Observable.just(4)
+                            .delay(2000, TimeUnit.MILLISECONDS)
+                            .doOnUnsubscribe(new Action0(){
+                                @Override
+                                public void call() {
+                                    unsubbedFrom4.set(true);
+                                    if (l4.getCount() <= 0)
+                                        failed4.set("Observable delayed until after outer unsubscribe was unsubscribed twice");
+                                    l4.countDown();
+                                }});
+                    observer.onNext(o);
+                }
+                else 
+                    observer.onCompleted();
+                return state + 1;
+            }
+        };
+        Observable.create(os).take(2).toBlocking().last();
+
+        latch(l1);
+        latch(l2);
+        latch(l3);
+        latch(l4);
+        assertNullOrMessage(failed1.get());
+        assertNullOrMessage(failed2.get());
+        assertNullOrMessage(failed3.get());
+        assertNullOrMessage(failed4.get());
+        Assert.assertTrue("did not unsub from first observable after terminal", unsubbedFrom1.get());
+        Assert.assertTrue("did not unsub from second observable after terminal", unsubbedFrom2.get());
+        Assert.assertTrue("did not unsub from Observable.never() inner obs", unsubbedFrom3.get());
+        Assert.assertTrue("did not unsub from observable unsubscribed from outer", unsubbedFrom4.get());
+    }
+
+    private void assertNullOrMessage(String msg) {
+        Assert.assertEquals(msg, null, msg);
+    }
+
+    private void latch(CountDownLatch latch) throws InterruptedException {
+        if (!latch.await(2, TimeUnit.SECONDS))
+            Assert.fail("count down didn't succeed");
+    }
+}