CurrentThreadScheduler Memory Leak Fixed

- Current/Immediate/NewThread/Executor Schedulers are passing unit tests
- Current/NewThread/Executor Schedulers do not leak memory on the recursion test (Immediate can’t be used for recursion otherwise it stack overflows)

Author: benjchristensen

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

@@ -17,17 +17,20 @@
 
 import java.util.PriorityQueue;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
 
 import rx.Scheduler;
 import rx.Subscription;
+import rx.subscriptions.CompositeSubscription;
+import rx.subscriptions.MultipleAssignmentSubscription;
 import rx.util.functions.Func2;
 
 /**
  * Schedules work on the current thread but does not execute immediately. Work is put in a queue and executed after the current unit of work is completed.
  */
 public class CurrentThreadScheduler extends Scheduler {
     private static final CurrentThreadScheduler INSTANCE = new CurrentThreadScheduler();
+    private static final AtomicLong counter = new AtomicLong(0);
 
     public static CurrentThreadScheduler getInstance() {
         return INSTANCE;
@@ -38,25 +41,27 @@ public static CurrentThreadScheduler getInstance() {
     /* package accessible for unit tests */CurrentThreadScheduler() {
     }
 
-    private final AtomicInteger counter = new AtomicInteger(0);
-
     @Override
     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(discardableAction, now());
-        return discardableAction;
+        enqueue(innerScheduler, discardableAction, 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
+        InnerCurrentThreadScheduler innerScheduler = new InnerCurrentThreadScheduler();
         DiscardableAction<T> discardableAction = new DiscardableAction<T>(state, new SleepingAction<T>(action, this, execTime));
-        enqueue(discardableAction, execTime);
+        enqueue(innerScheduler, discardableAction, execTime);
         return discardableAction;
     }
 
-    private void enqueue(DiscardableAction<?> action, long execTime) {
+    private static void enqueue(Scheduler scheduler, DiscardableAction<?> action, long execTime) {
         PriorityQueue<TimedAction> queue = QUEUE.get();
         boolean exec = queue == null;
 
@@ -69,19 +74,50 @@ private void enqueue(DiscardableAction<?> action, long execTime) {
 
         if (exec) {
             while (!queue.isEmpty()) {
-                queue.poll().action.call(this);
+                queue.poll().action.call(scheduler);
             }
 
             QUEUE.set(null);
         }
     }
 
+    private static class InnerCurrentThreadScheduler extends Scheduler implements Subscription {
+        private final MultipleAssignmentSubscription childSubscription = new MultipleAssignmentSubscription();
+
+        @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());
+            return childSubscription;
+        }
+
+        @Override
+        public <T> Subscription schedule(T state, Func2<? super Scheduler, ? super T, ? extends Subscription> action, long delayTime, TimeUnit unit) {
+            long execTime = now() + unit.toMillis(delayTime);
+
+            DiscardableAction<T> discardableAction = new DiscardableAction<T>(state, action);
+            childSubscription.set(discardableAction);
+            enqueue(this, discardableAction, execTime);
+            return childSubscription;
+        }
+
+        @Override
+        public void unsubscribe() {
+            childSubscription.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 Long execTime;
-        final Integer count; // In case if time between enqueueing took less than 1ms
+        final Long count; // In case if time between enqueueing took less than 1ms
 
-        private TimedAction(DiscardableAction<?> action, Long execTime, Integer count) {
+        private TimedAction(DiscardableAction<?> action, Long execTime, Long count) {
             this.action = action;
             this.execTime = execTime;
             this.count = count;

rxjava-core/src/test/java/rx/schedulers/AbstractSchedulerConcurrencyTests.java

@@ -8,13 +8,15 @@
 
 import org.junit.Test;
 
+import rx.Observable;
+import rx.Observer;
 import rx.Scheduler;
 import rx.Subscription;
 import rx.operators.SafeObservableSubscription;
-import rx.subscriptions.CompositeSubscription;
 import rx.subscriptions.Subscriptions;
 import rx.util.functions.Action0;
 import rx.util.functions.Action1;
+import rx.util.functions.Func1;
 import rx.util.functions.Func2;
 
 /**
@@ -26,6 +28,55 @@
  */
 public abstract class AbstractSchedulerConcurrencyTests extends AbstractSchedulerTests {
 
+    /**
+     * Bug report: https://github.com/Netflix/RxJava/issues/431
+     */
+    @Test
+    public final void testUnSubscribeForScheduler() throws InterruptedException {
+        final AtomicInteger countReceived = new AtomicInteger();
+        final AtomicInteger countGenerated = new AtomicInteger();
+        final SafeObservableSubscription s = new SafeObservableSubscription();
+        final CountDownLatch latch = new CountDownLatch(1);
+
+        s.wrap(Observable.interval(50, TimeUnit.MILLISECONDS)
+                .map(new Func1<Long, Long>() {
+                    @Override
+                    public Long call(Long aLong) {
+                        countGenerated.incrementAndGet();
+                        return aLong;
+                    }
+                })
+                .subscribeOn(getScheduler())
+                .observeOn(getScheduler())
+                .subscribe(new Observer<Long>() {
+                    @Override
+                    public void onCompleted() {
+                        System.out.println("--- completed");
+                    }
+
+                    @Override
+                    public void onError(Throwable e) {
+                        System.out.println("--- onError");
+                    }
+
+                    @Override
+                    public void onNext(Long args) {
+                        if (countReceived.incrementAndGet() == 2) {
+                            s.unsubscribe();
+                            latch.countDown();
+                        }
+                        System.out.println("==> Received " + args);
+                    }
+                }));
+
+        latch.await(1000, TimeUnit.MILLISECONDS);
+
+        System.out.println("----------- it thinks it is finished ------------------ ");
+        Thread.sleep(100);
+
+        assertEquals(2, countGenerated.get());
+    }
+
     @Test
     public void testUnsubscribeRecursiveScheduleWithStateAndFunc2() throws InterruptedException {
         final CountDownLatch latch = new CountDownLatch(1);