Fix ObserveOn, NewThreadScheduler and ScheduledObserver bugs

@headinthebox and I were working on some code and found differences in behavior between Rx.Net and RxJava with observeOn. This commit should fix that.

Author: benjchristensen

rxjava-core/src/main/java/rx/concurrency/NewThreadScheduler.java

@@ -15,12 +15,15 @@
  */
 package rx.concurrency;
 
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
 
 import rx.Scheduler;
 import rx.Subscription;
-import rx.operators.SafeObservableSubscription;
 import rx.subscriptions.CompositeSubscription;
 import rx.subscriptions.Subscriptions;
 import rx.util.functions.Func2;
@@ -29,27 +32,74 @@
  * Schedules work on a new thread.
  */
 public class NewThreadScheduler extends Scheduler {
-    private static final NewThreadScheduler INSTANCE = new NewThreadScheduler();
+
+    private final static NewThreadScheduler INSTANCE = new NewThreadScheduler();
+    private final static AtomicLong count = new AtomicLong();
 
     public static NewThreadScheduler getInstance() {
         return INSTANCE;
     }
 
-    @Override
-    public <T> Subscription schedule(final T state, final Func2<? super Scheduler, ? super T, ? extends Subscription> action) {
-        final SafeObservableSubscription subscription = new SafeObservableSubscription();
-        final Scheduler _scheduler = this;
+    private NewThreadScheduler() {
 
-        Thread t = new Thread(new Runnable() {
-            @Override
-            public void run() {
-                subscription.wrap(action.call(_scheduler, state));
-            }
-        }, "RxNewThreadScheduler");
+    }
 
-        t.start();
+    private static class EventLoopScheduler extends Scheduler {
+        private final ExecutorService executor;
 
-        return subscription;
+        private EventLoopScheduler() {
+            executor = Executors.newFixedThreadPool(1, new ThreadFactory() {
+
+                @Override
+                public Thread newThread(Runnable r) {
+                    return new Thread(r, "RxNewThreadScheduler-" + count.incrementAndGet());
+                }
+            });
+        }
+
+        @Override
+        public <T> Subscription schedule(final T state, final Func2<? super Scheduler, ? super T, ? extends Subscription> action) {
+            final Scheduler _scheduler = this;
+            return Subscriptions.from(executor.submit(new Runnable() {
+
+                @Override
+                public void run() {
+                    action.call(_scheduler, state);
+                }
+            }));
+        }
+
+        @Override
+        public <T> Subscription schedule(final T state, final Func2<? super Scheduler, ? super T, ? extends Subscription> action, final long delayTime, final TimeUnit unit) {
+            // we will use the system scheduler since it doesn't make sense to launch a new Thread and then sleep
+            // we will instead schedule the event then launch the thread after the delay has passed
+            final Scheduler _scheduler = this;
+            final CompositeSubscription subscription = new CompositeSubscription();
+            ScheduledFuture<?> f = GenericScheduledExecutorService.getInstance().schedule(new Runnable() {
+
+                @Override
+                public void run() {
+                    if (!subscription.isUnsubscribed()) {
+                        // when the delay has passed we now do the work on the actual scheduler
+                        Subscription s = _scheduler.schedule(state, action);
+                        // add the subscription to the CompositeSubscription so it is unsubscribed
+                        subscription.add(s);
+                    }
+                }
+            }, delayTime, unit);
+
+            // add the ScheduledFuture as a subscription so we can cancel the scheduled action if an unsubscribe happens
+            subscription.add(Subscriptions.create(f));
+
+            return subscription;
+        }
+
+    }
+
+    @Override
+    public <T> Subscription schedule(final T state, final Func2<? super Scheduler, ? super T, ? extends Subscription> action) {
+        EventLoopScheduler s = new EventLoopScheduler();
+        return s.schedule(state, action);
     }
 
     @Override

rxjava-core/src/main/java/rx/operators/OperationObserveOn.java

@@ -20,6 +20,7 @@
 
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 
 import org.junit.Test;
 import org.mockito.InOrder;
@@ -33,6 +34,8 @@
 import rx.Subscription;
 import rx.concurrency.ImmediateScheduler;
 import rx.concurrency.Schedulers;
+import rx.subscriptions.CompositeSubscription;
+import rx.util.functions.Func2;
 
 /**
  * Asynchronously notify Observers on the specified Scheduler.
@@ -60,7 +63,9 @@ public Subscription onSubscribe(final Observer<? super T> observer) {
                 // do nothing if we request ImmediateScheduler so we don't invoke overhead
                 return source.subscribe(observer);
             } else {
-                return source.subscribe(new ScheduledObserver<T>(observer, scheduler));
+                CompositeSubscription s = new CompositeSubscription();
+                s.add(source.subscribe(new ScheduledObserver<T>(s, observer, scheduler)));
+                return s;
             }
         }
     }

rxjava-core/src/main/java/rx/operators/OperationRetry.java

@@ -26,11 +26,13 @@
 import rx.Observable;
 import rx.Observable.OnSubscribeFunc;
 import rx.Observer;
+import rx.Scheduler;
 import rx.Subscription;
 import rx.concurrency.Schedulers;
 import rx.subscriptions.CompositeSubscription;
+import rx.subscriptions.MultipleAssignmentSubscription;
 import rx.subscriptions.Subscriptions;
-import rx.util.functions.Action0;
+import rx.util.functions.Func2;
 
 public class OperationRetry {
 
@@ -58,17 +60,19 @@ public Retry(Observable<T> source, int retryCount) {
 
         @Override
         public Subscription onSubscribe(Observer<? super T> observer) {
-            subscription.add(Schedulers.currentThread().schedule(attemptSubscription(observer)));
+            MultipleAssignmentSubscription rescursiveSubscription = new MultipleAssignmentSubscription();
+            subscription.add(Schedulers.currentThread().schedule(rescursiveSubscription, attemptSubscription(observer)));
+            subscription.add(rescursiveSubscription);
             return subscription;
         }
 
-        private Action0 attemptSubscription(final Observer<? super T> observer) {
-            return new Action0() {
+        private Func2<Scheduler, MultipleAssignmentSubscription, Subscription> attemptSubscription(final Observer<? super T> observer) {
+            return new Func2<Scheduler, MultipleAssignmentSubscription, Subscription>() {
 
                 @Override
-                public void call() {
+                public Subscription call(final Scheduler scheduler, final MultipleAssignmentSubscription rescursiveSubscription) {
                     attempts.incrementAndGet();
-                    source.subscribe(new Observer<T>() {
+                    return source.subscribe(new Observer<T>() {
 
                         @Override
                         public void onCompleted() {
@@ -79,10 +83,8 @@ public void onCompleted() {
                         public void onError(Throwable e) {
                             if ((retryCount == INFINITE_RETRY || attempts.get() <= retryCount) && !subscription.isUnsubscribed()) {
                                 // retry again
-                                // remove the last subscription since we have completed (so as we retry we don't build up a huge list)
-                                subscription.removeLast();
-                                // add the new subscription and schedule a retry
-                                subscription.add(Schedulers.currentThread().schedule(attemptSubscription(observer)));
+                                // add the new subscription and schedule a retry recursively
+                                rescursiveSubscription.setSubscription(scheduler.schedule(rescursiveSubscription, attemptSubscription(observer)));
                             } else {
                                 // give up and pass the failure
                                 observer.onError(e);
@@ -96,6 +98,7 @@ public void onNext(T v) {
                     });
 
                 }
+
             };
         }
 
@@ -157,7 +160,7 @@ public void testRetrySuccess() {
             inOrder.verify(observer, times(1)).onCompleted();
             inOrder.verifyNoMoreInteractions();
         }
-        
+
         @Test
         public void testInfiniteRetry() {
             int NUM_FAILURES = 20;

rxjava-core/src/main/java/rx/operators/ScheduledObserver.java

@@ -16,21 +16,28 @@
 package rx.operators;
 
 import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import rx.Notification;
 import rx.Observer;
 import rx.Scheduler;
-import rx.util.functions.Action0;
+import rx.Subscription;
+import rx.subscriptions.CompositeSubscription;
+import rx.subscriptions.MultipleAssignmentSubscription;
+import rx.util.functions.Func2;
 
 /* package */class ScheduledObserver<T> implements Observer<T> {
     private final Observer<? super T> underlying;
     private final Scheduler scheduler;
+    private final CompositeSubscription parentSubscription;
+    private final EventLoop eventLoop = new EventLoop();
 
     private final ConcurrentLinkedQueue<Notification<? extends T>> queue = new ConcurrentLinkedQueue<Notification<? extends T>>();
-    private final AtomicInteger counter = new AtomicInteger(0);
+    private final AtomicBoolean started = new AtomicBoolean();
 
-    public ScheduledObserver(Observer<? super T> underlying, Scheduler scheduler) {
+    public ScheduledObserver(CompositeSubscription s, Observer<? super T> underlying, Scheduler scheduler) {
+        this.parentSubscription = s;
         this.underlying = underlying;
         this.scheduler = scheduler;
     }
@@ -50,46 +57,83 @@ public void onNext(final T args) {
         enqueue(new Notification<T>(args));
     }
 
+    final AtomicInteger counter = new AtomicInteger();
+
     private void enqueue(Notification<? extends T> notification) {
-        // this must happen before 'counter' is used to provide synchronization between threads
+        // this must happen before synchronization between threads
         queue.offer(notification);
 
-        // we now use counter to atomically determine if we need to start processing or not
-        // it will be 0 if it's the first notification or the scheduler has finished processing work
-        // and we need to start doing it again
-        if (counter.getAndIncrement() == 0) {
-            processQueue();
+        /**
+         * If the counter is currently at 0 (before incrementing with this addition)
+         * we will schedule the work.
+         */
+        if (counter.getAndIncrement() <= 0) {
+            if (!started.get() && started.compareAndSet(false, true)) {
+                // first time we use the parent scheduler to start the event loop
+                MultipleAssignmentSubscription recursiveSubscription = new MultipleAssignmentSubscription();
+                parentSubscription.add(scheduler.schedule(recursiveSubscription, eventLoop));
+                parentSubscription.add(recursiveSubscription);
+            } else {
+                // subsequent times we reschedule existing one
+                eventLoop.reschedule();
+            }
         }
     }
 
-    private void processQueue() {
-        scheduler.schedule(new Action0() {
-            @Override
-            public void call() {
-                Notification<? extends T> not = queue.poll();
-
-                switch (not.getKind()) {
-                case OnNext:
-                    underlying.onNext(not.getValue());
-                    break;
-                case OnError:
-                    underlying.onError(not.getThrowable());
-                    break;
-                case OnCompleted:
-                    underlying.onCompleted();
-                    break;
-                default:
-                    throw new IllegalStateException("Unknown kind of notification " + not);
+    private class EventLoop implements Func2<Scheduler, MultipleAssignmentSubscription, Subscription> {
 
-                }
+        volatile Scheduler _recursiveScheduler;
+        volatile MultipleAssignmentSubscription _recursiveSubscription;
 
-                // decrement count and if we still have work to do
-                // recursively schedule ourselves to process again
-                if (counter.decrementAndGet() > 0) {
-                    scheduler.schedule(this);
-                }
+        public void reschedule() {
+            _recursiveSubscription.setSubscription(_recursiveScheduler.schedule(_recursiveSubscription, this));
+        }
 
+        @Override
+        public Subscription call(Scheduler s, MultipleAssignmentSubscription recursiveSubscription) {
+            /*
+             * --------------------------------------------------------------------------------------
+             * Set these the first time through so we can externally trigger recursive execution again
+             */
+            if (_recursiveScheduler == null) {
+                _recursiveScheduler = s;
+            }
+            if (_recursiveSubscription == null) {
+                _recursiveSubscription = recursiveSubscription;
             }
-        });
+            /*
+             * Back to regular flow
+             * --------------------------------------------------------------------------------------
+             */
+
+            do {
+                Notification<? extends T> notification = queue.poll();
+                // if we got a notification, send it
+                if (notification != null) {
+
+                    // if unsubscribed stop working
+                    if (parentSubscription.isUnsubscribed()) {
+                        return parentSubscription;
+                    }
+                    // process notification
+
+                    switch (notification.getKind()) {
+                    case OnNext:
+                        underlying.onNext(notification.getValue());
+                        break;
+                    case OnError:
+                        underlying.onError(notification.getThrowable());
+                        break;
+                    case OnCompleted:
+                        underlying.onCompleted();
+                        break;
+                    default:
+                        throw new IllegalStateException("Unknown kind of notification " + notification);
+                    }
+                }
+            } while (counter.decrementAndGet() > 0);
+
+            return parentSubscription;
+        }
     }
 }