Allow use of the returned subscription to cancel periodic scheduling

The documentation for schedulePeriodically indicates that the returned
subscription can be used to unsubscribe from the periodic action, or to
unschedule it if it has not been scheduled yet. That was the case only
before the first action took place, and it was then impossible to
unsubscribe using the given subscription, although unsubscribing the
worker did work.

Author: samueltardieu

rxjava-core/src/main/java/rx/Scheduler.java

@@ -19,6 +19,7 @@
 
 import rx.functions.Action0;
 import rx.schedulers.Schedulers;
+import rx.subscriptions.MultipleAssignmentSubscription;
 
 /**
  * Represents an object that schedules units of work.
@@ -108,18 +109,20 @@ public Subscription schedulePeriodically(final Action0 action, long initialDelay
             final long periodInNanos = unit.toNanos(period);
             final long startInNanos = TimeUnit.MILLISECONDS.toNanos(now()) + unit.toNanos(initialDelay);
 
+            final MultipleAssignmentSubscription mas = new MultipleAssignmentSubscription();
             final Action0 recursiveAction = new Action0() {
                 long count = 0;
                 @Override
                 public void call() {
-                    if (!isUnsubscribed()) {
+                    if (!mas.isUnsubscribed()) {
                         action.call();
                         long nextTick = startInNanos + (++count * periodInNanos);
-                        schedule(this, nextTick - TimeUnit.MILLISECONDS.toNanos(now()), TimeUnit.NANOSECONDS);
+                        mas.set(schedule(this, nextTick - TimeUnit.MILLISECONDS.toNanos(now()), TimeUnit.NANOSECONDS));
                     }
                 }
             };
-            return schedule(recursiveAction, initialDelay, unit);
+            mas.set(schedule(recursiveAction, initialDelay, unit));
+            return mas;
         }
 
         /**

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

@@ -30,6 +30,7 @@
 import org.mockito.Mockito;
 
 import rx.Scheduler;
+import rx.Subscription;
 import rx.functions.Action0;
 import rx.functions.Func1;
 
@@ -77,6 +78,48 @@ public void call() {
         inOrder.verify(calledOp, never()).call(anyLong());
     }
 
+    @SuppressWarnings("unchecked")
+    // mocking is unchecked, unfortunately
+    @Test
+    public final void testPeriodicSchedulingUnsubscription() {
+        final Func1<Long, Void> calledOp = mock(Func1.class);
+
+        final TestScheduler scheduler = new TestScheduler();
+        final Scheduler.Worker inner = scheduler.createWorker();
+
+        final Subscription subscription = inner.schedulePeriodically(new Action0() {
+            @Override
+            public void call() {
+                System.out.println(scheduler.now());
+                calledOp.call(scheduler.now());
+            }
+        }, 1, 2, TimeUnit.SECONDS);
+
+        verify(calledOp, never()).call(anyLong());
+
+        InOrder inOrder = Mockito.inOrder(calledOp);
+
+        scheduler.advanceTimeBy(999L, TimeUnit.MILLISECONDS);
+        inOrder.verify(calledOp, never()).call(anyLong());
+
+        scheduler.advanceTimeBy(1L, TimeUnit.MILLISECONDS);
+        inOrder.verify(calledOp, times(1)).call(1000L);
+
+        scheduler.advanceTimeBy(1999L, TimeUnit.MILLISECONDS);
+        inOrder.verify(calledOp, never()).call(3000L);
+
+        scheduler.advanceTimeBy(1L, TimeUnit.MILLISECONDS);
+        inOrder.verify(calledOp, times(1)).call(3000L);
+
+        scheduler.advanceTimeBy(5L, TimeUnit.SECONDS);
+        inOrder.verify(calledOp, times(1)).call(5000L);
+        inOrder.verify(calledOp, times(1)).call(7000L);
+
+        subscription.unsubscribe();
+        scheduler.advanceTimeBy(11L, TimeUnit.SECONDS);
+        inOrder.verify(calledOp, never()).call(anyLong());
+    }
+
     @Test
     public final void testImmediateUnsubscribes() {
         TestScheduler s = new TestScheduler();
@@ -97,4 +140,24 @@ public void call() {
         assertEquals(0, counter.get());
     }
 
+    @Test
+    public final void testImmediateUnsubscribes2() {
+        TestScheduler s = new TestScheduler();
+        final Scheduler.Worker inner = s.createWorker();
+        final AtomicInteger counter = new AtomicInteger(0);
+
+        final Subscription subscription = inner.schedule(new Action0() {
+
+            @Override
+            public void call() {
+                counter.incrementAndGet();
+                System.out.println("counter: " + counter.get());
+                inner.schedule(this);
+            }
+
+        });
+        subscription.unsubscribe();
+        assertEquals(0, counter.get());
+    }
+
 }