Rework Dispatchers.Unconfined to use thread-local event loop

Fixes #704

Author: qwwdfsad

common/kotlinx-coroutines-core-common/src/Dispatchers.common.kt

@@ -42,10 +42,10 @@ public expect object Dispatchers {
 
     /**
      * A coroutine dispatcher that is not confined to any specific thread.
-     * It executes initial continuation of the coroutine _immediately_ in the current call-frame
+     * It executes initial continuation of the coroutine in the current call-frame
      * and lets the coroutine resume in whatever thread that is used by the corresponding suspending function, without
-     * mandating any specific threading policy.
-     * **Note: use with extreme caution, not for general code**.
+     * mandating any specific threading policy. Nested coroutines launched in this dispatcher form an event-loop to avoid
+     * stack overflows.
      *
      * Note, that if you need your coroutine to be confined to a particular thread or a thread-pool after resumption,
      * but still want to execute it in the current call-frame until its first suspension, then you can use

common/kotlinx-coroutines-core-common/src/Unconfined.kt

@@ -4,13 +4,44 @@
 
 package kotlinx.coroutines
 
+import kotlinx.coroutines.internal.*
 import kotlin.coroutines.*
+import kotlin.jvm.*
 
 /**
  * A coroutine dispatcher that is not confined to any specific thread.
  */
+@NativeThreadLocal
 internal object Unconfined : CoroutineDispatcher() {
-    override fun isDispatchNeeded(context: CoroutineContext): Boolean = false
-    override fun dispatch(context: CoroutineContext, block: Runnable) { throw UnsupportedOperationException() }
+    private data class State(@JvmField var isActive: Boolean = false,
+                             @JvmField val threadLocalQueue: ArrayList<Runnable> = ArrayList())
+    private val state = CommonThreadLocal { State() }
+
+    override fun dispatch(context: CoroutineContext, block: Runnable) {
+        // Stack-based event loop on top of thread-local arraylist
+        val state = state.get()
+        if (state.isActive) {
+            state.threadLocalQueue.add(block)
+            return
+        }
+
+        try {
+            state.isActive = true
+            block.run()
+            while (!state.threadLocalQueue.isEmpty()) {
+                val element = state.threadLocalQueue.removeAt(state.threadLocalQueue.lastIndex)
+                element.run()
+            }
+        } catch (e: Throwable) {
+            /*
+             * This exception doesn't happen normally, only if user either submitted throwing runnable
+             * or if we have a bug in implementation. Anyway, reset state of the dispatcher to the initial.
+             */
+            state.threadLocalQueue.clear()
+            throw DispatchException("Unexpected exception in Unconfined loop, clearing pending tasks", e)
+        } finally {
+            state.isActive = false
+        }
+    }
     override fun toString(): String = "Unconfined"
 }

common/kotlinx-coroutines-core-common/src/internal/ThreadLocal.common.kt

@@ -0,0 +1,13 @@
+/*
+ * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.internal
+
+@OptionalExpectation
+@UseExperimental(ExperimentalMultiplatform::class)
+internal expect annotation class NativeThreadLocal()
+
+internal expect class CommonThreadLocal<T>(supplier: () -> T) {
+    fun get(): T
+}

common/kotlinx-coroutines-core-common/test/UnconfinedTest.kt

@@ -0,0 +1,74 @@
+/*
+ * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines
+
+import kotlin.test.*
+
+class UnconfinedTest : TestBase() {
+
+    @Test
+    fun testOrder() = runTest {
+        expect(1)
+        launch(Dispatchers.Unconfined) {
+            expect(2)
+            launch {
+                expect(4)
+                launch {
+                    expect(7)
+                }
+
+                launch {
+                    expect(6)
+                }
+                expect(5)
+            }
+
+            expect(3)
+        }
+
+        finish(8)
+    }
+
+    @Test
+    fun testBlockThrows() = runTest {
+        expect(1)
+        try {
+            withContext(Dispatchers.Unconfined) {
+                expect(2)
+                withContext(Dispatchers.Unconfined + CoroutineName("a")) {
+                    expect(3)
+                }
+
+                expect(4)
+                launch(start = CoroutineStart.ATOMIC) {
+                    expect(5)
+                }
+
+                throw TestException()
+            }
+        } catch (e: TestException) {
+            finish(6)
+        }
+    }
+
+    @Test
+    fun enterMultipleTimes() = runTest {
+        launch(Unconfined) {
+            expect(1)
+        }
+
+        launch(Unconfined) {
+            expect(2)
+        }
+
+        launch(Unconfined) {
+            expect(3)
+        }
+
+        finish(4)
+    }
+
+    class TestException : Throwable()
+}

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

@@ -0,0 +1,13 @@
+/*
+ * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.internal
+
+import java.lang.ThreadLocal
+
+internal actual typealias CommonThreadLocal<T> = ThreadLocalWithInitialValue<T>
+
+internal class ThreadLocalWithInitialValue<T>(private val supplier: () -> T) : ThreadLocal<T>() {
+    override fun initialValue(): T = supplier()
+}

core/kotlinx-coroutines-core/test/UnconfinedConcurrentTest.kt

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines
+
+import org.junit.*
+import org.junit.Test
+import java.util.concurrent.*
+import kotlin.test.*
+
+class UnconfinedConcurrentTest : TestBase() {
+    private val threads = 4
+    private val executor = newFixedThreadPoolContext(threads, "UnconfinedConcurrentTest")
+    private val threadLocal = ThreadLocal<Int>()
+
+    @After
+    fun tearDown() {
+        executor.close()
+    }
+
+    @Test(timeout = 10_000L)
+    fun testConcurrent() = runTest {
+        val iterations = 10_000
+        val startBarrier = CyclicBarrier(threads + 1)
+        val finishLatch = CountDownLatch(threads)
+
+        repeat(threads) { id ->
+            launch(executor) {
+                startBarrier.await()
+                repeat(iterations) {
+                    threadLocal.set(0)
+                    launch(Dispatchers.Unconfined) {
+                        assertEquals(0, threadLocal.get())
+                        launch(Dispatchers.Unconfined) {
+                            assertEquals(id, threadLocal.get())
+                        }
+
+                        threadLocal.set(id)
+                    }
+                }
+
+                finishLatch.countDown()
+            }
+        }
+
+        startBarrier.await()
+        finishLatch.await()
+    }
+}

js/kotlinx-coroutines-core-js/src/internal/ThreadLocal.kt

@@ -0,0 +1,10 @@
+/*
+ * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.internal
+
+internal actual class CommonThreadLocal<T> actual constructor(supplier: () -> T) {
+    private val value = supplier()
+    actual fun get(): T = value
+}

native/kotlinx-coroutines-core-native/src/Dispatchers.kt

@@ -12,7 +12,7 @@ public actual object Dispatchers {
 
     public actual val Main: MainCoroutineDispatcher = NativeMainDispatcher(Default)
 
-    public actual val Unconfined: CoroutineDispatcher = kotlinx.coroutines.Unconfined
+    public actual val Unconfined: CoroutineDispatcher get() = kotlinx.coroutines.Unconfined // Avoid freezing
 }
 
 private class NativeMainDispatcher(val delegate: CoroutineDispatcher) : MainCoroutineDispatcher() {

native/kotlinx-coroutines-core-native/src/internal/ThreadLocal.kt

@@ -0,0 +1,14 @@
+/*
+ * Copyright 2016-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines.internal
+import kotlin.native.concurrent.*
+
+@Suppress("ACTUAL_WITHOUT_EXPECT")
+internal actual typealias NativeThreadLocal = kotlin.native.ThreadLocal
+
+internal actual class CommonThreadLocal<T> actual constructor(supplier: () -> T) {
+    private val value = supplier()
+    actual fun get(): T = value
+}