Never loose coroutines when CoroutineScheduler is closed

* CoroutineScheduler close/shutdown sequence now guaranteed that a
  submitted task is either executed by the scheduler that is being
  concurrently shutdown or RejectedExecutionException is thrown.
* RejectedExecutionException is caught by the coroutine dispatcher
  implementation and coroutines are rescheduled into DefaultExecutor
  if a (custom) instance  of CoroutineScheduler is being shutdown.
* Added stress test for coroutine scheduler closing.

Author: elizarov

core/kotlinx-coroutines-core/src/internal/LockFreeMPMCQueue.kt

@@ -9,7 +9,9 @@ import kotlinx.atomicfu.*
 internal open class LockFreeMPMCQueueNode<T> {
     val next = atomic<T?>(null)
 
+    // internal declarations for inline functions
     @PublishedApi internal val nextValue: T? get() = next.value
+    @PublishedApi internal fun nextCas(expect: T?, update: T?) = next.compareAndSet(expect, update)
 }
 
 /*
@@ -23,9 +25,14 @@ internal open class LockFreeMPMCQueue<T : LockFreeMPMCQueueNode<T>> {
         atomic(LockFreeMPMCQueueNode<T>() as T) // sentinel
 
     private val tail = atomic(head.value)
-    internal val headValue: T get() = head.value
 
-    public fun addLast(node: T): Boolean {
+    // internal declarations for inline functions
+    @PublishedApi internal val headValue: T get() = head.value
+    @PublishedApi internal val tailValue: T get() = tail.value
+    @PublishedApi internal fun headCas(curHead: T, update: T) = head.compareAndSet(curHead, update)
+    @PublishedApi internal fun tailCas(curTail: T, update: T) = tail.compareAndSet(curTail, update)
+
+    public fun addLast(node: T) {
         tail.loop { curTail ->
             val curNext = curTail.next.value
             if (curNext != null) {
@@ -34,6 +41,22 @@ internal open class LockFreeMPMCQueue<T : LockFreeMPMCQueueNode<T>> {
             }
             if (curTail.next.compareAndSet(null, node)) {
                 tail.compareAndSet(curTail, node)
+                return
+            }
+        }
+    }
+
+    public inline fun addLastIfPrev(node: T, predicate: (prev: Any) -> Boolean): Boolean {
+        while(true) {
+            val curTail = tailValue
+            val curNext = curTail.nextValue
+            if (curNext != null) {
+                tailCas(curTail, curNext)
+                continue // retry
+            }
+            if (!predicate(curTail)) return false
+            if (curTail.nextCas(null, node)) {
+                tailCas(curTail, node)
                 return true
             }
         }
@@ -48,9 +71,7 @@ internal open class LockFreeMPMCQueue<T : LockFreeMPMCQueueNode<T>> {
         }
     }
 
-    fun headCas(curHead: T, update: T) = head.compareAndSet(curHead, update)
-
-    public inline fun removeFirstOrNullIf(predicate: (T) -> Boolean): T? {
+    public inline fun removeFirstOrNullIf(predicate: (first: T) -> Boolean): T? {
         while (true) {
             val curHead = headValue
             val next = curHead.nextValue ?: return null

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

@@ -301,23 +301,29 @@ internal class CoroutineScheduler(
         val currentWorker = Thread.currentThread() as? Worker
         // Capture # of created workers that cannot change anymore (mind the synchronized block!)
         val created = synchronized(workers) { createdWorkers }
+        // Shutdown all workers with the only exception of the current thread
         for (i in 1..created) {
             val worker = workers[i]!!
-            if (worker.isAlive && worker !== currentWorker) {
-                LockSupport.unpark(worker)
-                worker.join(timeout)
+            if (worker !== currentWorker) {
+                while (worker.isAlive) {
+                    LockSupport.unpark(worker)
+                    worker.join(timeout)
+                }
+                val state = worker.state
+                check(state === WorkerState.TERMINATED) { "Expected TERMINATED state, but found $state"}
                 worker.localQueue.offloadAllWork(globalQueue)
             }
-
         }
+        // Make sure no more work is added to GlobalQueue from anywhere
+        check(globalQueue.add(CLOSED_TASK)) { "GlobalQueue could not be closed yet" }
         // Finish processing tasks from globalQueue and/or from this worker's local queue
         while (true) {
-            val task = currentWorker?.findTask() ?: globalQueue.removeFirstOrNull() ?: break
+            val task = currentWorker?.findTask() ?: globalQueue.removeFirstIfNotClosed() ?: break
             runSafely(task)
         }
         // Shutdown current thread
         currentWorker?.tryReleaseCpu(WorkerState.TERMINATED)
-        // cleanup state to make sure that tryUnpark tries to create new threads and fails because isTerminated
+        // check & cleanup state
         assert(cpuPermits.availablePermits() == corePoolSize)
         parkedWorkersStack.value = 0L
         controlState.value = 0L
@@ -339,8 +345,12 @@ internal class CoroutineScheduler(
         when (submitToLocalQueue(task, fair)) {
             ADDED -> return
             NOT_ADDED -> {
-                globalQueue.addLast(task) // offload task to local queue
-                requestCpuWorker() // ask for help
+                // try to offload task to global queue
+                if (!globalQueue.add(task)) {
+                    // Global queue is closed in the last step of close/shutdown -- no more tasks should be accepted
+                    throw RejectedExecutionException("$schedulerName was terminated")
+                }
+                requestCpuWorker()
             }
             else -> requestCpuWorker() // ask for help
         }
@@ -439,7 +449,7 @@ internal class CoroutineScheduler(
     private fun createNewWorker(): Int {
         synchronized(workers) {
             // Make sure we're not trying to resurrect terminated scheduler
-            if (isTerminated) throw RejectedExecutionException("$schedulerName was terminated")
+            if (isTerminated) return -1
             val state = controlState.value
             val created = createdWorkers(state)
             val blocking = blockingWorkers(state)
@@ -464,6 +474,12 @@ internal class CoroutineScheduler(
             ?: return NOT_ADDED
         if (worker.scheduler !== this) return NOT_ADDED // different scheduler's worker (!!!)
 
+        /*
+         * This worker could have been already terminated from this thread by close/shutdown and it should not
+         * accept any more tasks into its local queue.
+         */
+        if (worker.state === WorkerState.TERMINATED) return NOT_ADDED
+
         var result = ADDED
         if (task.mode == TaskMode.NON_BLOCKING) {
             /*
@@ -923,9 +939,9 @@ internal class CoroutineScheduler(
              * once per two core pool size iterations
              */
             val globalFirst = nextInt(2 * corePoolSize) == 0
-            if (globalFirst) globalQueue.removeFirstOrNull()?.let { return it }
+            if (globalFirst) globalQueue.removeFirstIfNotClosed()?.let { return it }
             localQueue.poll()?.let { return it }
-            if (!globalFirst) globalQueue.removeFirstOrNull()?.let { return it }
+            if (!globalFirst) globalQueue.removeFirstIfNotClosed()?.let { return it }
             return trySteal()
         }
 

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

@@ -47,10 +47,18 @@ open class ExperimentalCoroutineDispatcher(
     private var coroutineScheduler = createScheduler()
 
     override fun dispatch(context: CoroutineContext, block: Runnable): Unit =
-        coroutineScheduler.dispatch(block)
+        try {
+            coroutineScheduler.dispatch(block)
+        } catch (e: RejectedExecutionException) {
+            DefaultExecutor.dispatch(context, block)
+        }
 
     override fun dispatchYield(context: CoroutineContext, block: Runnable): Unit =
-        coroutineScheduler.dispatch(block, fair = true)
+        try {
+            coroutineScheduler.dispatch(block, fair = true)
+        } catch (e: RejectedExecutionException) {
+            DefaultExecutor.dispatchYield(context, block)
+        }
 
     override fun close() = coroutineScheduler.close()
 
@@ -84,7 +92,11 @@ open class ExperimentalCoroutineDispatcher(
     }
 
     internal fun dispatchWithContext(block: Runnable, context: TaskContext, fair: Boolean): Unit =
-        coroutineScheduler.dispatch(block, context, fair)
+        try {
+            coroutineScheduler.dispatch(block, context, fair)
+        } catch (e: RejectedExecutionException) {
+            DefaultExecutor.execute(block)
+        }
 
     private fun createScheduler() = CoroutineScheduler(corePoolSize, maxPoolSize, idleWorkerKeepAliveNs)
 

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

@@ -101,8 +101,19 @@ internal class Task(
         "Task[${block.classSimpleName}@${block.hexAddress}, $submissionTime, $taskContext]"
 }
 
+private val EMPTY_RUNNABLE = Runnable {}
+internal val CLOSED_TASK = Task(EMPTY_RUNNABLE, 0, NonBlockingContext)
+
 // Open for tests
 internal open class GlobalQueue : LockFreeMPMCQueue<Task>() {
+    // Returns false when GlobalQueue was was already closed
+    public fun add(task: Task): Boolean =
+        addLastIfPrev(task) { prev -> prev !== CLOSED_TASK }
+
+    // Returns null when GlobalQueue was was already closed
+    public fun removeFirstIfNotClosed(): Task? =
+        removeFirstOrNullIf { first -> first !== CLOSED_TASK }
+
     // Open for tests
     public open fun removeFirstBlockingModeOrNull(): Task? =
         removeFirstOrNullIf { it.mode == TaskMode.PROBABLY_BLOCKING }

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

@@ -140,14 +140,23 @@ internal class WorkQueue {
     private fun offloadWork(globalQueue: GlobalQueue) {
         repeat((bufferSize / 2).coerceAtLeast(1)) {
             val task = pollExternal() ?: return
-            globalQueue.addLast(task)
+            addToGlobalQueue(globalQueue, task)
         }
     }
 
+    private fun addToGlobalQueue(globalQueue: GlobalQueue, task: Task) {
+        /*
+         * globalQueue is closed as the very last step in the shutdown sequence when all worker threads had
+         * been already shutdown (with the only exception of the last worker thread that might be performing
+         * shutdown procedure itself). As a consistency check we do a [cheap!] check that it is not closed here yet.
+         */
+        check(globalQueue.add(task)) { "GlobalQueue could not be closed yet" }
+    }
+
     internal fun offloadAllWork(globalQueue: GlobalQueue) {
+        lastScheduledTask.getAndSet(null)?.let { addToGlobalQueue(globalQueue, it) }
         while (true) {
-            val task = pollExternal() ?: return
-            globalQueue.addLast(task)
+            addToGlobalQueue(globalQueue, pollExternal() ?: return)
         }
     }
 

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

@@ -0,0 +1,68 @@
+/*
+ * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.experimental.scheduling
+
+import kotlinx.atomicfu.*
+import kotlinx.coroutines.experimental.*
+import org.junit.Test
+import java.util.*
+import kotlin.test.*
+
+class CoroutineSchedulerCloseStressTest : TestBase() {
+    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 dispatcher: ExecutorCoroutineDispatcher
+    private var closeIndex = -1
+
+    private val started = atomic(0)
+    private val finished = atomic(0)
+
+    @Test
+    fun testNormalClose() {
+        try {
+            launchCoroutines()
+        } finally {
+            dispatcher.close()
+        }
+    }
+
+    @Test
+    fun testRacingClose() {
+        repeat(N_REPEAT) {
+            closeIndex = rnd.nextInt(N_COROS)
+            launchCoroutines()
+        }
+    }
+
+    private fun launchCoroutines() = runBlocking {
+        dispatcher = ExperimentalCoroutineDispatcher(N_THREADS)
+        started.value = 0
+        finished.value = 0
+        withContext(dispatcher) {
+            launchChild(0, 0)
+        }
+        assertEquals(N_COROS, started.value)
+        assertEquals(N_COROS, finished.value)
+    }
+
+    private fun CoroutineScope.launchChild(index: Int, level: Int): Job = launch(start = CoroutineStart.ATOMIC) {
+        started.incrementAndGet()
+        try {
+            if (index == closeIndex) dispatcher.close()
+            if (level < MAX_LEVEL) {
+                launchChild(2 * index + 1, level + 1)
+                launchChild(2 * index + 2, level + 1)
+            } else {
+                delay(1000)
+            }
+        } finally {
+            finished.incrementAndGet()
+        }
+    }
+}