make sure we're waiting for delayJob properly

Author: Rechee

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

@@ -114,22 +114,23 @@ private class SchedulerChannelTask(
     parentJob: Job
 ) : Disposable {
     private val taskScope = CoroutineScope(dispatcher + Job(parentJob))
-    private val delayJob: Job
+    private val delayResult: Deferred<Unit>
 
     init {
-        delayJob = taskScope.launch {
+        delayResult = taskScope.async {
             delay(delayMillis)
         }
     }
 
     fun execute() {
-        if (delayJob.isCompleted && taskScope.isActive) {
-            block.run()
-        } else {
-            taskScope.launch {
-                delayJob.join()
-                yield()
+        if (taskScope.isActive) {
+            if (delayResult.isCompleted) {
                 block.run()
+            } else {
+                taskScope.launch {
+                    delayResult.await()
+                    block.run()
+                }
             }
         }
     }

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

@@ -42,14 +42,17 @@ class SchedulerStressTest : TestBase() {
         val n = 2000 * stressTestMultiplier
         coroutineScope {
             repeat(n) { i ->
-                val a = ByteArray(1000000) //1MB
-                val disposable = block(Runnable {
-                    runBlocking {
-                        keepMe(a)
-                    }
-                })
-                disposable.dispose()
-                expect(i + 2)
+                launch {
+                    val a = ByteArray(1000000) //1MB
+                    val disposable = block(Runnable {
+                        expectUnreached()
+                        runBlocking {
+                            keepMe(a)
+                        }
+                    })
+                    disposable.dispose()
+                    expect(i + 2)
+                }
                 yield()
             }
         }

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

@@ -313,11 +313,13 @@ class SchedulerTest : TestBase() {
         val iterations = 2
         coroutineScope {
             for (i in (0..iterations)) {
-                suspendCancellableCoroutine<Unit> {
-                    worker.schedule(Runnable {
-                        expect(2 + i)
-                        it.resume(Unit)
-                    })
+                launch {
+                    suspendCancellableCoroutine<Unit> {
+                        worker.schedule(Runnable {
+                            expect(2 + i)
+                            it.resume(Unit)
+                        })
+                    }
                 }
             }
         }