CurrentThreadScheduler Delayed Execution Fix

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

@@ -1,12 +1,12 @@
 /**
  * Copyright 2013 Netflix, Inc.
- * 
+ *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
- * 
+ *
  * http://www.apache.org/licenses/LICENSE-2.0
- * 
+ *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -17,9 +17,9 @@
 
 import static org.mockito.Mockito.*;
 
-import java.util.LinkedList;
-import java.util.Queue;
+import java.util.PriorityQueue;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import org.junit.Test;
 import org.mockito.InOrder;
@@ -39,46 +39,70 @@ public static CurrentThreadScheduler getInstance() {
         return INSTANCE;
     }
 
-    private static final ThreadLocal<Queue<DiscardableAction<?>>> QUEUE = new ThreadLocal<Queue<DiscardableAction<?>>>();
+    private static final ThreadLocal<PriorityQueue<TimedAction>> QUEUE = new ThreadLocal<PriorityQueue<TimedAction>>();
 
     private CurrentThreadScheduler() {
     }
 
+    private final AtomicInteger counter = new AtomicInteger(0);
+
     @Override
     public <T> Subscription schedule(T state, Func2<Scheduler, T, Subscription> action) {
         DiscardableAction<T> discardableAction = new DiscardableAction<T>(state, action);
-        enqueue(discardableAction);
+        enqueue(discardableAction, now());
         return discardableAction;
     }
 
     @Override
     public <T> Subscription schedule(T state, Func2<Scheduler, T, Subscription> action, long dueTime, TimeUnit unit) {
-        // since we are executing immediately on this thread we must cause this thread to sleep
-        // TODO right now the 'enqueue' does not take delay into account so if another task is enqueued after this it will 
-        // wait behind the sleeping action ... should that be the case or should it be allowed to proceed ahead of the delayed action?
-        return schedule(state, new SleepingAction<T>(action, this, dueTime, unit));
+        long execTime = now() + unit.toMillis(dueTime);
+
+        DiscardableAction<T> discardableAction = new DiscardableAction<T>(state, new SleepingAction<T>(action, this, execTime));
+        enqueue(discardableAction, execTime);
+        return discardableAction;
     }
 
-    private void enqueue(DiscardableAction<?> action) {
-        Queue<DiscardableAction<?>> queue = QUEUE.get();
+    private void enqueue(DiscardableAction<?> action, long execTime) {
+        PriorityQueue<TimedAction> queue = QUEUE.get();
         boolean exec = queue == null;
 
         if (exec) {
-            queue = new LinkedList<DiscardableAction<?>>();
+            queue = new PriorityQueue<TimedAction>();
             QUEUE.set(queue);
         }
 
-        queue.add(action);
+        queue.add(new TimedAction(action, execTime, counter.incrementAndGet()));
 
         if (exec) {
             while (!queue.isEmpty()) {
-                queue.poll().call(this);
+                queue.poll().action.call(this);
             }
 
             QUEUE.set(null);
         }
     }
 
+    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
+
+        private TimedAction(DiscardableAction<?> action, Long execTime, Integer count) {
+            this.action = action;
+            this.execTime = execTime;
+            this.count = count;
+        }
+
+        @Override
+        public int compareTo(TimedAction that) {
+            int result = execTime.compareTo(that.execTime);
+            if (result == 0) {
+                return count.compareTo(that.count);
+            }
+            return result;
+        }
+    }
+
     public static class UnitTest {
 
         @Test
@@ -146,6 +170,58 @@ public void testSequenceOfActions() {
 
         }
 
+        @Test
+        public void testSequenceOfDelayedActions() {
+            final CurrentThreadScheduler scheduler = new CurrentThreadScheduler();
+
+            final Action0 first = mock(Action0.class);
+            final Action0 second = mock(Action0.class);
+
+            scheduler.schedule(new Action0() {
+                @Override
+                public void call() {
+                    scheduler.schedule(first, 30, TimeUnit.MILLISECONDS);
+                    scheduler.schedule(second, 10, TimeUnit.MILLISECONDS);
+                }
+            });
+
+            InOrder inOrder = inOrder(first, second);
+
+            inOrder.verify(second, times(1)).call();
+            inOrder.verify(first, times(1)).call();
+
+
+        }
+
+        @Test
+        public void testMixOfDelayedAndNonDelayedActions() {
+            final CurrentThreadScheduler scheduler = new CurrentThreadScheduler();
+
+            final Action0 first = mock(Action0.class);
+            final Action0 second = mock(Action0.class);
+            final Action0 third = mock(Action0.class);
+            final Action0 fourth = mock(Action0.class);
+
+            scheduler.schedule(new Action0() {
+                @Override
+                public void call() {
+                    scheduler.schedule(first);
+                    scheduler.schedule(second, 300, TimeUnit.MILLISECONDS);
+                    scheduler.schedule(third, 100, TimeUnit.MILLISECONDS);
+                    scheduler.schedule(fourth);
+                }
+            });
+
+            InOrder inOrder = inOrder(first, second, third, fourth);
+
+            inOrder.verify(first, times(1)).call();
+            inOrder.verify(fourth, times(1)).call();
+            inOrder.verify(third, times(1)).call();
+            inOrder.verify(second, times(1)).call();
+
+
+        }
+
     }
 
 }

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

@@ -48,7 +48,9 @@ public <T> Subscription schedule(T state, Func2<Scheduler, T, Subscription> acti
     @Override
     public <T> Subscription schedule(T state, Func2<Scheduler, T, Subscription> action, long dueTime, TimeUnit unit) {
         // since we are executing immediately on this thread we must cause this thread to sleep
-        return schedule(state, new SleepingAction<T>(action, this, dueTime, unit));
+        long execTime = now() + unit.toMillis(dueTime);
+
+        return schedule(state, new SleepingAction<T>(action, this, execTime));
     }
 
     public static class UnitTest {

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

@@ -1,12 +1,12 @@
 /**
  * Copyright 2013 Netflix, Inc.
- * 
+ *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
- * 
+ *
  * http://www.apache.org/licenses/LICENSE-2.0
- * 
+ *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -26,24 +26,25 @@
     private final Scheduler scheduler;
     private final long execTime;
 
-    public SleepingAction(Func2<Scheduler, T, Subscription> underlying, Scheduler scheduler, long timespan, TimeUnit timeUnit) {
+    public SleepingAction(Func2<Scheduler, T, Subscription> underlying, Scheduler scheduler, long execTime) {
         this.underlying = underlying;
         this.scheduler = scheduler;
-        this.execTime = scheduler.now() + timeUnit.toMillis(timespan);
+        this.execTime = execTime;
     }
 
+
     @Override
     public Subscription call(Scheduler s, T state) {
         if (execTime > scheduler.now()) {
             long delay = execTime - scheduler.now();
             if (delay> 0) {
-            try {
+                try {
                     Thread.sleep(delay);
                 }
-             catch (InterruptedException e) {
-                Thread.currentThread().interrupt();
-                throw new RuntimeException(e);
-            }
+                catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                    throw new RuntimeException(e);
+                }
             }
         }