Add way to convert CoroutineDispatcher to Scheduler

reactive/kotlinx-coroutines-rx2/api/kotlinx-coroutines-rx2.api

@@ -67,7 +67,9 @@ public final class kotlinx/coroutines/rx2/RxObservableKt {
 }
 
 public final class kotlinx/coroutines/rx2/RxSchedulerKt {
-	public static final fun asCoroutineDispatcher (Lio/reactivex/Scheduler;)Lkotlinx/coroutines/rx2/SchedulerCoroutineDispatcher;
+	public static final fun asCoroutineDispatcher (Lio/reactivex/Scheduler;)Lkotlinx/coroutines/CoroutineDispatcher;
+	public static final synthetic fun asCoroutineDispatcher (Lio/reactivex/Scheduler;)Lkotlinx/coroutines/rx2/SchedulerCoroutineDispatcher;
+	public static final fun asScheduler (Lkotlinx/coroutines/CoroutineDispatcher;)Lio/reactivex/Scheduler;
 }
 
 public final class kotlinx/coroutines/rx2/RxSingleKt {

reactive/kotlinx-coroutines-rx2/build.gradle

@@ -5,6 +5,7 @@
 dependencies {
     compile project(':kotlinx-coroutines-reactive')
     testCompile project(':kotlinx-coroutines-reactive').sourceSets.test.output
+    testCompile project(':kotlinx-coroutines-test')
     testCompile "org.reactivestreams:reactive-streams-tck:$reactive_streams_version"
     compile "io.reactivex.rxjava2:rxjava:$rxjava2_version"
 }

reactive/kotlinx-coroutines-rx2/src/RxScheduler.kt

@@ -4,16 +4,136 @@
 
 package kotlinx.coroutines.rx2
 
-import io.reactivex.Scheduler
+import io.reactivex.*
+import io.reactivex.disposables.*
+import io.reactivex.plugins.*
 import kotlinx.coroutines.*
-import java.util.concurrent.TimeUnit
-import kotlin.coroutines.CoroutineContext
+import kotlinx.coroutines.channels.*
+import java.util.concurrent.*
+import kotlin.coroutines.*
 
 /**
  * Converts an instance of [Scheduler] to an implementation of [CoroutineDispatcher]
  * and provides native support of [delay] and [withTimeout].
  */
-public fun Scheduler.asCoroutineDispatcher(): SchedulerCoroutineDispatcher = SchedulerCoroutineDispatcher(this)
+public fun Scheduler.asCoroutineDispatcher(): CoroutineDispatcher =
+    if (this is DispatcherScheduler) {
+        dispatcher
+    } else {
+        SchedulerCoroutineDispatcher(this)
+    }
+
+@Deprecated(level = DeprecationLevel.HIDDEN, message = "Since 1.4.2, binary compatibility with earlier versions")
+@JvmName("asCoroutineDispatcher")
+public fun Scheduler.asCoroutineDispatcher0(): SchedulerCoroutineDispatcher =
+    SchedulerCoroutineDispatcher(this)
+
+/**
+ * Converts an instance of [CoroutineDispatcher] to an implementation of [Scheduler].
+ */
+public fun CoroutineDispatcher.asScheduler(): Scheduler =
+    if (this is SchedulerCoroutineDispatcher) {
+        scheduler
+    } else {
+        DispatcherScheduler(this)
+    }
+
+private class DispatcherScheduler(val dispatcher: CoroutineDispatcher) : Scheduler() {
+
+    private val schedulerJob = SupervisorJob()
+
+    private val scope = CoroutineScope(schedulerJob + dispatcher + CoroutineExceptionHandler { _, throwable ->
+        RxJavaPlugins.onError(throwable)
+    })
+
+    override fun scheduleDirect(block: Runnable): Disposable =
+        scheduleDirect(block, 0, TimeUnit.MILLISECONDS)
+
+    override fun scheduleDirect(block: Runnable, delay: Long, unit: TimeUnit): Disposable {
+        if (!scope.isActive) return Disposables.disposed()
+        val newBlock = RxJavaPlugins.onSchedule(block)
+        return scope.launch {
+            delay(unit.toMillis(delay))
+            newBlock.run()
+        }.asDisposable()
+    }
+
+    override fun createWorker(): Worker =
+        DispatcherWorker(dispatcher, schedulerJob).also {
+            it.start()
+        }
+
+    override fun shutdown() {
+        scope.cancel()
+    }
+
+    private class DispatcherWorker(dispatcher: CoroutineDispatcher, parentJob: Job) : Worker() {
+
+        private val workerJob = SupervisorJob(parentJob)
+        private val workerScope = CoroutineScope(workerJob + dispatcher)
+        private val blockChannel = Channel<SchedulerChannelTask>(Channel.UNLIMITED)
+
+        fun start() {
+            workerScope.launch {
+                while (isActive) {
+                    val task = blockChannel.receive()
+                    task.execute()
+                }
+            }
+        }
+
+        override fun isDisposed(): Boolean = !workerScope.isActive
+
+        override fun schedule(block: Runnable): Disposable =
+            schedule(block, 0, TimeUnit.MILLISECONDS)
+
+        override fun schedule(block: Runnable, delay: Long, unit: TimeUnit): Disposable {
+            if (!workerScope.isActive) return Disposables.disposed()
+
+            val newBlock = RxJavaPlugins.onSchedule(block)
+
+            val taskJob = Job(workerJob)
+            val task = SchedulerChannelTask(newBlock, taskJob)
+
+            if (delay <= 0L) {
+                blockChannel.offer(task)
+            } else {
+                // Use `taskJob` as the parent here so the delay will also get cancelled if the Disposable
+                // is disposed.
+                workerScope.launch(taskJob) {
+                    // Delay *before* enqueuing the task, so other tasks (e.g. via schedule without delay)
+                    // aren't blocked by the delay.
+                    delay(unit.toMillis(delay))
+                    // Once the task is ready to run, it still needs to be executed via the queue to comply
+                    // with the Scheduler contract of running all worker tasks in a non-overlapping manner.
+                    blockChannel.offer(task)
+                }
+            }
+
+            return taskJob.asDisposable()
+        }
+
+        override fun dispose() {
+            workerScope.cancel()
+        }
+    }
+}
+
+/**
+ * Represents a task to be queued sequentially on a [Channel] for a [Scheduler.Worker].
+ *
+ * Delayed tasks do not block [Channel] from processing other tasks
+ */
+private class SchedulerChannelTask(
+    private val block: Runnable,
+    job: Job
+) : JobDisposable(job) {
+    fun execute() {
+        if (job.isActive) {
+            block.run()
+        }
+    }
+}
 
 /**
  * Implements [CoroutineDispatcher] on top of an arbitrary [Scheduler].
@@ -50,3 +170,13 @@ public class SchedulerCoroutineDispatcher(
     /** @suppress */
     override fun hashCode(): Int = System.identityHashCode(scheduler)
 }
+
+private open class JobDisposable(protected val job: Job) : Disposable {
+    override fun isDisposed(): Boolean = !job.isActive
+
+    override fun dispose() {
+        job.cancel()
+    }
+}
+
+private fun Job.asDisposable(): Disposable = JobDisposable(this)

reactive/kotlinx-coroutines-rx2/test/SchedulerStressTest.kt

@@ -0,0 +1,114 @@
+/*
+ * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.rx2
+
+import kotlinx.coroutines.*
+import kotlinx.coroutines.test.*
+import org.junit.*
+import java.util.concurrent.*
+
+class SchedulerStressTest : TestBase() {
+    @Before
+    fun setup() {
+        ignoreLostThreads("RxCachedThreadScheduler-", "RxCachedWorkerPoolEvictor-", "RxSchedulerPurge-")
+    }
+
+    /**
+     * Test that we don't get an OOM if we schedule many jobs at once. It's expected that if you don't dispose that you'd
+     * see a OOM error.
+     */
+    @Test
+    fun testSchedulerDisposed(): Unit = runTest {
+        val dispatcher = currentDispatcher() as CoroutineDispatcher
+        val scheduler = dispatcher.asScheduler()
+        testRunnableDisposed(scheduler::scheduleDirect)
+    }
+
+    @Test
+    fun testSchedulerWorkerDisposed(): Unit = runTest {
+        val dispatcher = currentDispatcher() as CoroutineDispatcher
+        val scheduler = dispatcher.asScheduler()
+        testRunnableDisposed(scheduler.createWorker()::schedule)
+    }
+
+    private suspend fun testRunnableDisposed(block: RxSchedulerBlockNoDelay) {
+        expect(1)
+
+        val dispatcher = currentDispatcher() as CoroutineDispatcher
+        val scheduler = dispatcher.asScheduler()
+
+        val n = 2000 * stressTestMultiplier
+        coroutineScope {
+            repeat(n) { i ->
+                launch {
+                    val a = ByteArray(1000000) //1MB
+                    val disposable = block(Runnable {
+                        expectUnreached()
+                        runBlocking {
+                            keepMe(a)
+                        }
+                    })
+                    disposable.dispose()
+                    expect(i + 2)
+                }
+                yield()
+            }
+        }
+
+        scheduler.shutdown()
+        finish(n + 2)
+    }
+
+    /**
+     * Test function that holds a reference. Used for testing OOM situations
+     */
+    private suspend fun keepMe(a: ByteArray) {
+        delay(10)
+    }
+
+    /**
+     * Test that we don't get an OOM if we schedule many delayed jobs at once. It's expected that if you don't dispose that you'd
+     * see a OOM error.
+     */
+    @Test
+    fun testSchedulerDisposedDuringDelay(): Unit = runBlockingTest {
+        val dispatcher = currentDispatcher() as CoroutineDispatcher
+        val scheduler = dispatcher.asScheduler()
+        testRunnableDisposedDuringDelay(scheduler::scheduleDirect)
+    }
+
+    @Test
+    fun testSchedulerWorkerDisposedDuringDelay(): Unit = runBlockingTest {
+        val dispatcher = currentDispatcher() as CoroutineDispatcher
+        val scheduler = dispatcher.asScheduler()
+        testRunnableDisposedDuringDelay(scheduler.createWorker()::schedule)
+    }
+
+    private suspend fun TestCoroutineScope.testRunnableDisposedDuringDelay(block: RxSchedulerBlockWithDelay) {
+        expect(1)
+
+        val dispatcher = currentDispatcher() as CoroutineDispatcher
+        val scheduler = dispatcher.asScheduler()
+
+        val n = 2000 * stressTestMultiplier
+        coroutineScope {
+            repeat(n) { i ->
+                val a = ByteArray(1000000) //1MB
+                val delayMillis: Long = 10
+                val disposable = block(Runnable {
+                    runBlocking {
+                        keepMe(a)
+                    }
+                }, delayMillis, TimeUnit.MILLISECONDS)
+                disposable.dispose()
+                advanceTimeBy(delayMillis)
+                yield()
+            }
+        }
+
+        scheduler.shutdown()
+        finish(2)
+    }
+}