Fix shutdown sequence for CoroutineScheduler's views

Fixes #678

Author: qwwdfsad

core/kotlinx-coroutines-core/src/scheduling/CoroutineScheduler.kt

@@ -237,6 +237,7 @@ internal class CoroutineScheduler(
     private inline fun createdWorkers(state: Long): Int = (state and CREATED_MASK).toInt()
     private inline fun blockingWorkers(state: Long): Int = (state and BLOCKING_MASK shr BLOCKING_SHIFT).toInt()
 
+    // Guarded by synchronization
     private inline fun incrementCreatedWorkers(): Int = createdWorkers(controlState.incrementAndGet())
     private inline fun decrementCreatedWorkers(): Int = createdWorkers(controlState.getAndDecrement())
 
@@ -340,7 +341,7 @@ internal class CoroutineScheduler(
     fun dispatch(block: Runnable, taskContext: TaskContext = NonBlockingContext, fair: Boolean = false) {
         timeSource.trackTask() // this is needed for virtual time support
         // TODO at some point make DispatchTask extend Task and make its field settable to save an allocation
-        val task = Task(block, schedulerTimeSource.nanoTime(), taskContext)
+        val task = createTask(block, taskContext)
         // try to submit the task to the local queue and act depending on the result
         when (submitToLocalQueue(task, fair)) {
             ADDED -> return
@@ -356,6 +357,8 @@ internal class CoroutineScheduler(
         }
     }
 
+    internal fun createTask(block: Runnable, taskContext: TaskContext) = Task(block, schedulerTimeSource.nanoTime(), taskContext)
+
     /**
      * Unparks or creates a new [Worker] for executing non-blocking tasks if there are idle cores
      */
@@ -764,8 +767,12 @@ internal class CoroutineScheduler(
         private fun afterTask(task: Task) {
             if (task.mode != TaskMode.NON_BLOCKING) {
                 decrementBlockingWorkers()
-                assert(state == WorkerState.BLOCKING) { "Expected BLOCKING state, but has $state" }
-                state = WorkerState.RETIRING
+                val currentState = state
+                // Shutdown sequence of blocking dispatcher
+                if (currentState !== WorkerState.TERMINATED) {
+                    assert(currentState == WorkerState.BLOCKING) { "Expected BLOCKING state, but has $currentState" }
+                    state = WorkerState.RETIRING
+                }
             }
         }
 

core/kotlinx-coroutines-core/src/scheduling/Dispatcher.kt

@@ -101,7 +101,8 @@ open class ExperimentalCoroutineDispatcher(
         try {
             coroutineScheduler.dispatch(block, context, fair)
         } catch (e: RejectedExecutionException) {
-            DefaultExecutor.execute(block)
+            // Context shouldn't be lost here to properly invoke before/after task
+            DefaultExecutor.execute(coroutineScheduler.createTask(block, context))
         }
 
     private fun createScheduler() = CoroutineScheduler(corePoolSize, maxPoolSize, idleWorkerKeepAliveNs)

core/kotlinx-coroutines-core/test/scheduling/CoroutineSchedulerCloseStressTest.kt

@@ -7,16 +7,29 @@ package kotlinx.coroutines.experimental.scheduling
 import kotlinx.atomicfu.*
 import kotlinx.coroutines.experimental.*
 import org.junit.Test
+import org.junit.runner.*
+import org.junit.runners.*
 import java.util.*
+import java.util.concurrent.*
 import kotlin.test.*
 
-class CoroutineSchedulerCloseStressTest : TestBase() {
+@RunWith(Parameterized::class)
+class CoroutineSchedulerCloseStressTest(private val mode: Mode) : TestBase() {
+    enum class Mode { CPU, BLOCKING, CPU_LIMITED }
+
+    companion object {
+        @Parameterized.Parameters(name = "mode={0}")
+        @JvmStatic
+        fun params(): Collection<Array<Any>> = Mode.values().map { arrayOf<Any>(it) }
+    }
+
     private val N_REPEAT = 2 * stressTestMultiplier
     private val MAX_LEVEL = 5
     private val N_COROS = (1 shl (MAX_LEVEL + 1)) - 1
     private val N_THREADS = 4
     private val rnd = Random()
 
+    private lateinit var closeableDispatcher: ExperimentalCoroutineDispatcher
     private lateinit var dispatcher: ExecutorCoroutineDispatcher
     private var closeIndex = -1
 
@@ -28,7 +41,7 @@ class CoroutineSchedulerCloseStressTest : TestBase() {
         try {
             launchCoroutines()
         } finally {
-            dispatcher.close()
+            closeableDispatcher.close()
         }
     }
 
@@ -41,7 +54,12 @@ class CoroutineSchedulerCloseStressTest : TestBase() {
     }
 
     private fun launchCoroutines() = runBlocking {
-        dispatcher = ExperimentalCoroutineDispatcher(N_THREADS)
+        closeableDispatcher = ExperimentalCoroutineDispatcher(N_THREADS)
+        dispatcher = when (mode) {
+            Mode.CPU -> closeableDispatcher
+            Mode.CPU_LIMITED -> closeableDispatcher.limited(N_THREADS) as ExecutorCoroutineDispatcher
+            Mode.BLOCKING -> closeableDispatcher.blocking(N_THREADS) as ExecutorCoroutineDispatcher
+        }
         started.value = 0
         finished.value = 0
         withContext(dispatcher) {
@@ -54,15 +72,19 @@ class CoroutineSchedulerCloseStressTest : TestBase() {
     private fun CoroutineScope.launchChild(index: Int, level: Int): Job = launch(start = CoroutineStart.ATOMIC) {
         started.incrementAndGet()
         try {
-            if (index == closeIndex) dispatcher.close()
+            if (index == closeIndex) closeableDispatcher.close()
             if (level < MAX_LEVEL) {
                 launchChild(2 * index + 1, level + 1)
                 launchChild(2 * index + 2, level + 1)
             } else {
-                delay(1000)
+                if (rnd.nextBoolean()) {
+                    delay(1000)
+                } else {
+                    yield()
+                }
             }
         } finally {
             finished.incrementAndGet()
         }
     }
-}
+}