CurrentThreadScheduler Fixes

benjchristensen · benjchristensen · commit 43e3d777bb03 · 2013-12-31T10:37:41.000-08:00
- outer/inner scheduling so nested order is correct while not deadlocking on certain nested use cases as found in previous testing
diff --git a/rxjava-core/src/main/java/rx/schedulers/CurrentThreadScheduler.java b/rxjava-core/src/main/java/rx/schedulers/CurrentThreadScheduler.java
@@ -21,8 +21,8 @@
 
 import rx.Scheduler;
 import rx.Subscription;
-import rx.subscriptions.CompositeSubscription;
 import rx.subscriptions.MultipleAssignmentSubscription;
+import rx.util.functions.Func1;
 import rx.util.functions.Func2;
 
 /**
@@ -36,7 +36,17 @@ public static CurrentThreadScheduler getInstance() {
         return INSTANCE;
     }
 
-    private static final ThreadLocal<PriorityQueue<TimedAction>> QUEUE = new ThreadLocal<PriorityQueue<TimedAction>>();
+    private static final ThreadLocal<PriorityQueue<TimedAction>> QUEUE = new ThreadLocal<PriorityQueue<TimedAction>>() {
+        protected java.util.PriorityQueue<TimedAction> initialValue() {
+            return new PriorityQueue<TimedAction>();
+        };
+    };
+
+    private static final ThreadLocal<Boolean> PROCESSING = new ThreadLocal<Boolean>() {
+        protected Boolean initialValue() {
+            return Boolean.FALSE;
+        };
+    };
 
     /* package accessible for unit tests */CurrentThreadScheduler() {
     }
@@ -45,50 +55,57 @@ public static CurrentThreadScheduler getInstance() {
     public <T> Subscription schedule(T state, Func2<? super Scheduler, ? super T, ? extends Subscription> action) {
         // immediately move to the InnerCurrentThreadScheduler
         InnerCurrentThreadScheduler innerScheduler = new InnerCurrentThreadScheduler();
-        DiscardableAction<T> discardableAction = new DiscardableAction<T>(state, action);
-        enqueue(innerScheduler, discardableAction, now());
+        innerScheduler.enqueue(new DiscardableAction<T>(state, action), now());
+        enqueueFromOuter(innerScheduler, now());
         return innerScheduler;
     }
 
     @Override
     public <T> Subscription schedule(T state, Func2<? super Scheduler, ? super T, ? extends Subscription> action, long dueTime, TimeUnit unit) {
         long execTime = now() + unit.toMillis(dueTime);
 
-        // immediately move to the InnerCurrentThreadScheduler
+        // create an inner scheduler and queue it for execution
         InnerCurrentThreadScheduler innerScheduler = new InnerCurrentThreadScheduler();
-        DiscardableAction<T> discardableAction = new DiscardableAction<T>(state, new SleepingAction<T>(action, this, execTime));
-        enqueue(innerScheduler, discardableAction, execTime);
-        return discardableAction;
+        innerScheduler.enqueue(new DiscardableAction<T>(state, new SleepingAction<T>(action, this, execTime)), execTime);
+        enqueueFromOuter(innerScheduler, execTime);
+        return innerScheduler;
     }
 
-    private static void enqueue(Scheduler scheduler, DiscardableAction<?> action, long execTime) {
+    /*
+     * This will accept InnerCurrentThreadScheduler instances and execute them in order they are received
+     * and on each of them will loop internally until each is complete.
+     */
+    private void enqueueFromOuter(final InnerCurrentThreadScheduler innerScheduler, long execTime) {
+        // Note that everything here is single-threaded so we won't have race conditions
         PriorityQueue<TimedAction> queue = QUEUE.get();
-        boolean exec = queue == null;
+        queue.add(new TimedAction(new Func1<Scheduler, Subscription>() {
 
-        if (exec) {
-            queue = new PriorityQueue<TimedAction>();
-            QUEUE.set(queue);
-        }
-
-        queue.add(new TimedAction(action, execTime, counter.incrementAndGet()));
+            @Override
+            public Subscription call(Scheduler _) {
+                // when the InnerCurrentThreadScheduler gets scheduled we want to process its tasks
+                return innerScheduler.startProcessing();
+            }
+        }, execTime, counter.incrementAndGet()));
 
-        if (exec) {
+        // first time through starts the loop
+        if (!PROCESSING.get()) {
+            PROCESSING.set(Boolean.TRUE);
             while (!queue.isEmpty()) {
-                queue.poll().action.call(scheduler);
+                queue.poll().action.call(innerScheduler);
             }
-
-            QUEUE.set(null);
+            PROCESSING.set(Boolean.FALSE);
         }
     }
 
     private static class InnerCurrentThreadScheduler extends Scheduler implements Subscription {
         private final MultipleAssignmentSubscription childSubscription = new MultipleAssignmentSubscription();
+        private final PriorityQueue<TimedAction> innerQueue = new PriorityQueue<TimedAction>();
 
         @Override
         public <T> Subscription schedule(T state, Func2<? super Scheduler, ? super T, ? extends Subscription> action) {
             DiscardableAction<T> discardableAction = new DiscardableAction<T>(state, action);
             childSubscription.set(discardableAction);
-            enqueue(this, discardableAction, now());
+            enqueue(discardableAction, now());
             return childSubscription;
         }
 
@@ -98,10 +115,21 @@ public <T> Subscription schedule(T state, Func2<? super Scheduler, ? super T, ?
 
             DiscardableAction<T> discardableAction = new DiscardableAction<T>(state, action);
             childSubscription.set(discardableAction);
-            enqueue(this, discardableAction, execTime);
+            enqueue(discardableAction, execTime);
             return childSubscription;
         }
 
+        private void enqueue(Func1<Scheduler, Subscription> action, long execTime) {
+            innerQueue.add(new TimedAction(action, execTime, counter.incrementAndGet()));
+        }
+
+        private Subscription startProcessing() {
+            while (!innerQueue.isEmpty()) {
+                innerQueue.poll().action.call(this);
+            }
+            return this;
+        }
+
         @Override
         public void unsubscribe() {
             childSubscription.unsubscribe();
@@ -113,11 +141,11 @@ public void unsubscribe() {
      * Use time to sort items so delayed actions are sorted to their appropriate position in the queue.
      */
     private static class TimedAction implements Comparable<TimedAction> {
-        final DiscardableAction<?> action;
+        final Func1<Scheduler, Subscription> action;
         final Long execTime;
         final Long count; // In case if time between enqueueing took less than 1ms
 
-        private TimedAction(DiscardableAction<?> action, Long execTime, Long count) {
+        private TimedAction(Func1<Scheduler, Subscription> action, Long execTime, Long count) {
             this.action = action;
             this.execTime = execTime;
             this.count = count;
