running blocks more directly inside DispatcherScheduler. also fixing stress tests

Author: Rechee

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

@@ -35,38 +35,35 @@ public fun CoroutineDispatcher.asScheduler(): Scheduler =
 
 private class DispatcherScheduler(internal val dispatcher: CoroutineDispatcher) : Scheduler() {
 
-    val job = SupervisorJob()
-    private val scope = CoroutineScope(job)
+    private val job = SupervisorJob()
+    private val scope = CoroutineScope(job + dispatcher)
 
-    override fun scheduleDirect(run: Runnable): Disposable {
+    override fun scheduleDirect(block: Runnable): Disposable {
         if (scope.isActive) {
             return scope.launch {
-                dispatchBlock(run)
+                block.runWithRx()
             }.asDisposable()
         }
 
         return Disposables.disposed()
     }
 
-    override fun scheduleDirect(run: Runnable, delay: Long, unit: TimeUnit): Disposable {
+    override fun scheduleDirect(block: Runnable, delay: Long, unit: TimeUnit): Disposable {
         if (delay <= 0) {
-            return scheduleDirect(run)
+            return scheduleDirect(block)
         }
 
         if (scope.isActive) {
             return scope.launch {
                 delay(unit.toMillis(delay))
-                dispatchBlock(run)
+                block.runWithRx()
             }.asDisposable()
         }
 
         return Disposables.disposed()
     }
 
-    private fun dispatchBlock(block: Runnable) {
-        val decoratedRun = RxJavaPlugins.onSchedule(block)
-        dispatcher.dispatch(EmptyCoroutineContext, decoratedRun)
-    }
+    private fun Runnable.runWithRx() = RxJavaPlugins.onSchedule(this).run()
 
     override fun createWorker(): Worker =
         DispatcherWorker(dispatcher, job)
@@ -86,7 +83,7 @@ private class DispatcherScheduler(internal val dispatcher: CoroutineDispatcher)
         override fun schedule(block: Runnable): Disposable {
             startProcessingQueue()
             if (workerScope.isActive) {
-                val job = workerScope.launch(Dispatchers.Unconfined, CoroutineStart.LAZY) {
+                val job = workerScope.launch(start = CoroutineStart.LAZY) {
                     block.run()
                 }
                 blockChannel.offer(job)
@@ -169,7 +166,12 @@ public class SchedulerCoroutineDispatcher(
     override fun hashCode(): Int = System.identityHashCode(scheduler)
 }
 
-private fun Job.asDisposable(): Disposable = object : Disposable {
-    override fun isDisposed(): Boolean = !isActive
-    override fun dispose() = cancel()
+private class JobDisposable(private val job: Job) : Disposable {
+    override fun isDisposed(): Boolean = !job.isActive
+
+    override fun dispose() {
+        job.cancel()
+    }
 }
+
+public fun Job.asDisposable(): Disposable = JobDisposable(this)

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

@@ -7,14 +7,17 @@ package kotlinx.coroutines.rx2
 import kotlinx.coroutines.*
 import org.junit.*
 import java.util.concurrent.*
-import kotlin.coroutines.*
 
 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 testScheduleDirectDisposed(): Unit = runTest {
         expect(1)
@@ -29,20 +32,14 @@ class SchedulerStressTest : TestBase() {
         val n = 2000 * stressTestMultiplier
         coroutineScope {
             repeat(n) { i ->
-                launch {
-                    val a = ByteArray(1000000) //1MB
-                    suspendCancellableCoroutine<Unit> {
-                        val disposable = scheduler.scheduleDirect {
-                            runBlocking {
-                                keepMe(a)
-                                it.resume(Unit)
-                            }
-                        }
-                        expect(i + 2)
-                        disposable.dispose()
-                        it.resume(Unit)
+                val a = ByteArray(1000000) //1MB
+                val disposable = scheduler.scheduleDirect {
+                    runBlocking {
+                        keepMe(a)
                     }
                 }
+                disposable.dispose()
+                expect(i + 2)
                 yield()
             }
         }
@@ -51,34 +48,37 @@ class SchedulerStressTest : TestBase() {
         finish(n + 2)
     }
 
-    @Ignore
+    /**
+     * 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 testScheduleDirectDisposedDuringDelay(): Unit = runTest {
         expect(1)
 
-        fun keepMe(a: ByteArray) {
-            // does nothing, makes sure the variable is kept in state-machine
+        suspend fun keepMe(a: ByteArray) {
+            delay(10)
         }
 
         val dispatcher = currentDispatcher() as CoroutineDispatcher
         val scheduler = dispatcher.asScheduler()
 
         val n = 2000 * stressTestMultiplier
         coroutineScope {
-            repeat(n) {
-                launch {
-                    val a = ByteArray(1000000) //1MB
-                    val disposable = scheduler.scheduleDirect({
-                        expectUnreached()
-                    }, 50, TimeUnit.MILLISECONDS)
-                    disposable.dispose()
-                    check(disposable.isDisposed)
-                    delay(60)
-                    keepMe(a)
-                }
+            repeat(n) { i ->
+                val a = ByteArray(1000000) //1MB
+                val disposable = scheduler.scheduleDirect({
+                    runBlocking {
+                        keepMe(a)
+                    }
+                }, 10, TimeUnit.MILLISECONDS)
+                disposable.dispose()
+                expect(i + 2)
+                yield()
             }
         }
 
-        finish(2)
+        scheduler.shutdown()
+        finish(n + 2)
     }
 }