Merge pull request #1158 from akarnokd/SchedulerFixes0506

Scheduler correctness improvements.

Author: benjchristensen

rxjava-core/src/main/java/rx/schedulers/EventLoopsScheduler.java

@@ -22,14 +22,14 @@
 import rx.Scheduler;
 import rx.Subscription;
 import rx.functions.Action0;
-import rx.schedulers.NewThreadScheduler.OnActionComplete;
+import rx.schedulers.NewThreadScheduler.NewThreadWorker.ScheduledAction;
 import rx.subscriptions.CompositeSubscription;
 import rx.subscriptions.Subscriptions;
 
 /* package */class EventLoopsScheduler extends Scheduler {
-
-    private static class ComputationSchedulerPool {
-        final int cores = Runtime.getRuntime().availableProcessors();
+    /** Manages a fixed number of workers. */
+    static final class FixedSchedulerPool {
+        final int cores;
         final ThreadFactory factory = new ThreadFactory() {
             final AtomicInteger counter = new AtomicInteger();
 
@@ -41,47 +41,46 @@ public Thread newThread(Runnable r) {
             }
         };
 
-        final EventLoopScheduler[] eventLoops;
+        final PoolWorker[] eventLoops;
+        long n;
 
-        ComputationSchedulerPool() {
+        FixedSchedulerPool() {
             // initialize event loops
-            eventLoops = new EventLoopScheduler[cores];
+            this.cores = Runtime.getRuntime().availableProcessors();
+            this.eventLoops = new PoolWorker[cores];
             for (int i = 0; i < cores; i++) {
-                eventLoops[i] = new EventLoopScheduler(factory);
+                this.eventLoops[i] = new PoolWorker(factory);
             }
         }
 
-        private static ComputationSchedulerPool INSTANCE = new ComputationSchedulerPool();
-
-        long n = 0;
-
-        public EventLoopScheduler getEventLoop() {
-            // round-robin selection (improvements to come)
-            return eventLoops[(int) (n++ % cores)];
+        public PoolWorker getEventLoop() {
+            // simple round robin, improvements to come
+            return eventLoops[(int)(n++ % cores)];
         }
-
     }
 
+    final FixedSchedulerPool pool;
+    
+    /**
+     * Create a scheduler with pool size equal to the available processor
+     * count and using least-recent worker selection policy.
+     */
+    EventLoopsScheduler() {
+        pool = new FixedSchedulerPool();
+    }
+    
     @Override
     public Worker createWorker() {
-        return new EventLoop();
+        return new EventLoopWorker(pool.getEventLoop());
     }
 
-    private static class EventLoop extends Scheduler.Worker {
+    private static class EventLoopWorker extends Scheduler.Worker {
         private final CompositeSubscription innerSubscription = new CompositeSubscription();
-        private final EventLoopScheduler pooledEventLoop;
-        private final OnActionComplete onComplete;
-
-        EventLoop() {
-            pooledEventLoop = ComputationSchedulerPool.INSTANCE.getEventLoop();
-            onComplete = new OnActionComplete() {
+        private final PoolWorker poolWorker;
 
-                @Override
-                public void complete(Subscription s) {
-                    innerSubscription.remove(s);
-                }
-
-            };
+        EventLoopWorker(PoolWorker poolWorker) {
+            this.poolWorker = poolWorker;
+            
         }
 
         @Override
@@ -96,29 +95,25 @@ public boolean isUnsubscribed() {
 
         @Override
         public Subscription schedule(Action0 action) {
-            if (innerSubscription.isUnsubscribed()) {
-                // don't schedule, we are unsubscribed
-                return Subscriptions.empty();
-            }
-            return pooledEventLoop.schedule(action, onComplete);
+            return schedule(action, 0, null);
         }
-
         @Override
         public Subscription schedule(Action0 action, long delayTime, TimeUnit unit) {
             if (innerSubscription.isUnsubscribed()) {
                 // don't schedule, we are unsubscribed
                 return Subscriptions.empty();
             }
 
-            return pooledEventLoop.schedule(action, delayTime, unit, onComplete);
+            ScheduledAction s = poolWorker.scheduleActual(action, delayTime, unit);
+            innerSubscription.add(s);
+            s.addParent(innerSubscription);
+            return s;
         }
-
     }
-
-    private static class EventLoopScheduler extends NewThreadScheduler.EventLoopScheduler {
-        EventLoopScheduler(ThreadFactory threadFactory) {
+    
+    private static final class PoolWorker extends NewThreadScheduler.NewThreadWorker {
+        PoolWorker(ThreadFactory threadFactory) {
             super(threadFactory);
         }
     }
-
 }

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

@@ -15,13 +15,13 @@
  */
 package rx.schedulers;
 
-import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
-import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.Future;
+import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
-import java.util.concurrent.atomic.AtomicReference;
 
 import rx.Scheduler;
 import rx.Subscription;
@@ -56,97 +56,122 @@ private NewThreadScheduler() {
 
     @Override
     public Worker createWorker() {
-        return new EventLoopScheduler(THREAD_FACTORY);
+        return new NewThreadWorker(THREAD_FACTORY);
     }
 
-    /* package */static class EventLoopScheduler extends Scheduler.Worker implements Subscription {
+    /* package */static class NewThreadWorker extends Scheduler.Worker implements Subscription {
         private final CompositeSubscription innerSubscription = new CompositeSubscription();
-        private final ExecutorService executor;
+        private final ScheduledExecutorService executor;
 
-        /* package */EventLoopScheduler(ThreadFactory threadFactory) {
-            executor = Executors.newSingleThreadExecutor(threadFactory);
+        /* package */NewThreadWorker(ThreadFactory threadFactory) {
+            executor = Executors.newScheduledThreadPool(1, threadFactory);
         }
 
         @Override
         public Subscription schedule(final Action0 action) {
-            return schedule(action, null);
+            return schedule(action, 0, null);
         }
 
-        /* package */Subscription schedule(final Action0 action, final OnActionComplete onComplete) {
+        @Override
+        public Subscription schedule(final Action0 action, long delayTime, TimeUnit unit) {
             if (innerSubscription.isUnsubscribed()) {
-                // don't schedule, we are unsubscribed
                 return Subscriptions.empty();
             }
-
-            final AtomicReference<Subscription> sf = new AtomicReference<Subscription>();
-            Subscription s = Subscriptions.from(executor.submit(new Runnable() {
-
-                @Override
-                public void run() {
-                    try {
-                        if (innerSubscription.isUnsubscribed()) {
-                            return;
-                        }
-                        action.call();
-                    } finally {
-                        // remove the subscription now that we're completed
-                        Subscription s = sf.get();
-                        if (s != null) {
-                            innerSubscription.remove(s);
-                        }
-                        if (onComplete != null) {
-                            onComplete.complete(s);
-                        }
-                    }
-                }
-            }));
-
-            sf.set(s);
-            innerSubscription.add(s);
-            return s;
+            return scheduleActual(action, delayTime, unit);
         }
 
-        @Override
-        public Subscription schedule(final Action0 action, long delayTime, TimeUnit unit) {
-            return schedule(action, delayTime, unit, null);
+        /* package */ScheduledAction scheduleActual(final Action0 action, long delayTime, TimeUnit unit) {
+            ScheduledAction run = new ScheduledAction(action, innerSubscription);
+            Future<?> f;
+            if (delayTime <= 0) {
+                f = executor.submit(run);
+            } else {
+                f = executor.schedule(run, delayTime, unit);
+            }
+            run.add(Subscriptions.from(f));
+            
+            return run;
+        }
+        
+        /** Remove a child subscription from a composite when unsubscribing. */
+        private static final class Remover implements Subscription {
+            final Subscription s;
+            final CompositeSubscription parent;
+            final AtomicBoolean once;
+            
+            public Remover(Subscription s, CompositeSubscription parent) {
+                this.s = s;
+                this.parent = parent;
+                this.once = new AtomicBoolean();
+            }
+            
+            @Override
+            public boolean isUnsubscribed() {
+                return s.isUnsubscribed();
+            }
+            
+            @Override
+            public void unsubscribe() {
+                if (once.compareAndSet(false, true)) {
+                    parent.remove(s);
+                }
+            }
+            
         }
+        /** 
+         * A runnable that executes an Action0 and can be cancelled
+         * The analogue is the Subscriber in respect of an Observer.
+         */
+        public static final class ScheduledAction implements Runnable, Subscription {
+            final CompositeSubscription cancel;
+            final Action0 action;
+            final CompositeSubscription parent;
+            final AtomicBoolean once;
+
+            public ScheduledAction(Action0 action, CompositeSubscription parent) {
+                this.action = action;
+                this.parent = parent;
+                this.cancel = new CompositeSubscription();
+                this.once = new AtomicBoolean();
+            }
 
-        /* package */Subscription schedule(final Action0 action, long delayTime, TimeUnit unit, final OnActionComplete onComplete) {
-            final AtomicReference<Subscription> sf = new AtomicReference<Subscription>();
-            // 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
-            ScheduledFuture<?> f = GenericScheduledExecutorService.getInstance().schedule(new Runnable() {
-
-                @Override
-                public void run() {
-                    try {
-                        if (innerSubscription.isUnsubscribed()) {
-                            return;
-                        }
-                        // now that the delay is past schedule the work to be done for real on the UI thread
-                        schedule(action);
-                    } finally {
-                        // remove the subscription now that we're completed
-                        Subscription s = sf.get();
-                        if (s != null) {
-                            innerSubscription.remove(s);
-                        }
-                        if (onComplete != null) {
-                            onComplete.complete(s);
-                        }
-                    }
+            @Override
+            public void run() {
+                try {
+                    action.call();
+                } finally {
+                    unsubscribe();
                 }
-            }, delayTime, unit);
+            }
 
-            // add the ScheduledFuture as a subscription so we can cancel the scheduled action if an unsubscribe happens
-            Subscription s = Subscriptions.from(f);
-            sf.set(s);
-            innerSubscription.add(s);
-            return s;
+            @Override
+            public boolean isUnsubscribed() {
+                return cancel.isUnsubscribed();
+            }
+            
+            @Override
+            public void unsubscribe() {
+                if (once.compareAndSet(false, true)) {
+                    cancel.unsubscribe();
+                    parent.remove(this);
+                }
+            }
+            public void add(Subscription s) {
+                cancel.add(s);
+            }
+            /** 
+             * Adds a parent to this ScheduledAction so when it is 
+             * cancelled or terminates, it can remove itself from this parent.
+             * @param parent 
+             */
+            public void addParent(CompositeSubscription parent) {
+                cancel.add(new Remover(this, parent));
+            } 
         }
 
         @Override
         public void unsubscribe() {
+            executor.shutdown();
             innerSubscription.unsubscribe();
         }
 
@@ -156,11 +181,4 @@ public boolean isUnsubscribed() {
         }
 
     }
-
-    /* package */static interface OnActionComplete {
-
-        public void complete(Subscription s);
-
-    }
-
 }