Experimenting with Schedulers.currentScheduler()

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

@@ -99,11 +99,13 @@ public void schedule(final Action1<Scheduler.Inner> action) {
                 @Override
                 public void run() {
                     try {
+                        Schedulers.setCurrentScheduler(_inner);
                         if (innerSubscription.isUnsubscribed()) {
                             return;
                         }
                         action.call(_inner);
                     } finally {
+                        Schedulers.setCurrentScheduler(null);
                         // remove the subscription now that we're completed
                         Subscription s = sf.get();
                         if (s != null) {

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

@@ -60,6 +60,16 @@ private Schedulers() {
 
     }
 
+    private static ThreadLocal<Scheduler.Inner> currentScheduler = new ThreadLocal<Scheduler.Inner>();
+
+    /* package */static void setCurrentScheduler(Scheduler.Inner scheduler) {
+        currentScheduler.set(scheduler);
+    }
+
+    public static Scheduler.Inner current() {
+        return currentScheduler.get();
+    }
+
     /**
      * {@link Scheduler} that executes work immediately on the current thread.
      * 

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

@@ -16,7 +16,22 @@
 
 package rx.schedulers;
 
+import static org.junit.Assert.*;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.junit.Test;
+
+import rx.Observable;
+import rx.Observable.OnSubscribe;
 import rx.Scheduler;
+import rx.Scheduler.Inner;
+import rx.Subscriber;
+import rx.functions.Action1;
+import rx.observers.TestSubscriber;
 
 public class NewThreadSchedulerTest extends AbstractSchedulerConcurrencyTests {
 
@@ -25,4 +40,77 @@ protected Scheduler getScheduler() {
         return NewThreadScheduler.getInstance();
     }
 
+    @Test
+    public void testCurrentScheduler() throws InterruptedException {
+
+        final AtomicReference<Scheduler.Inner> capturedCurrent = new AtomicReference<Scheduler.Inner>();
+        final AtomicReference<Scheduler.Inner> capturedInner = new AtomicReference<Scheduler.Inner>();
+        final AtomicReference<Thread> first = new AtomicReference<Thread>();
+        final AtomicReference<Thread> second = new AtomicReference<Thread>();
+        final AtomicInteger c = new AtomicInteger();
+        final CountDownLatch latch = new CountDownLatch(1);
+        Schedulers.newThread().schedule(new Action1<Inner>() {
+
+            @Override
+            public void call(Inner inner) {
+                c.incrementAndGet();
+                capturedInner.set(inner);
+                first.set(Thread.currentThread());
+                capturedCurrent.set(Schedulers.current());
+                Schedulers.current().schedule(new Action1<Inner>() {
+
+                    @Override
+                    public void call(Inner t1) {
+                        c.incrementAndGet();
+                        second.set(Thread.currentThread());
+                        latch.countDown();
+                    }
+
+                });
+            }
+
+        });
+        latch.await(1000, TimeUnit.MILLISECONDS);
+        assertNull(Schedulers.current());
+        assertEquals(2, c.get());
+        assertEquals(first.get(), second.get());
+        assertEquals(capturedCurrent.get(), capturedInner.get());
+    }
+
+    @Test
+    public void testCurrentSchedulerViaSubscribeOn() throws InterruptedException {
+
+        final AtomicReference<Scheduler.Inner> capturedCurrent = new AtomicReference<Scheduler.Inner>();
+        final AtomicReference<Thread> first = new AtomicReference<Thread>();
+        final AtomicReference<Thread> second = new AtomicReference<Thread>();
+        final AtomicInteger c = new AtomicInteger();
+
+        TestSubscriber<Integer> ts = new TestSubscriber<Integer>();
+        Observable.create(new OnSubscribe<Integer>() {
+
+            @Override
+            public void call(final Subscriber<? super Integer> s) {
+                c.incrementAndGet();
+                first.set(Thread.currentThread());
+                capturedCurrent.set(Schedulers.current());
+                Schedulers.current().schedule(new Action1<Inner>() {
+
+                    @Override
+                    public void call(Inner t1) {
+                        c.incrementAndGet();
+                        second.set(Thread.currentThread());
+                        s.onCompleted();
+                    }
+
+                });
+            }
+
+        }).subscribeOn(Schedulers.newThread()).subscribe(ts);
+
+        ts.awaitTerminalEvent();
+        assertNotNull(capturedCurrent.get());
+        assertNull(Schedulers.current());
+        assertEquals(2, c.get());
+        assertEquals(first.get(), second.get());
+    }
 }