Restore thread context elements when directly resuming to parent (#1577)

This fix solves the problem of restoring thread-context when
returning to another context in an undispatched way.

It impacts suspend/resume performance of coroutines that use ThreadContextElement and undispatched coroutines.

The kotlinx.coroutines code poisons the context with special 'UndispatchedMarker' element and linear lookup is performed only when the marker is present. The code also contains a description of an alternative approach in order to save a linear lookup in complex coroutines hierarchies.

Fast-path of coroutine resumption is slowed down by a single context lookup.

Fixes #985

Co-authored-by: Roman Elizarov <elizarov@gmail.com>
Co-authored-by: Vsevolod Tolstopyatov <qwwdfsad@gmail.com>

12 files changed, 421 insertions, 23 deletions

diff --git a/benchmarks/src/jmh/kotlin/benchmarks/ChannelSinkBenchmark.kt b/benchmarks/src/jmh/kotlin/benchmarks/ChannelSinkBenchmark.kt
index 6c5b6231..f718758d 100644
--- a/benchmarks/src/jmh/kotlin/benchmarks/ChannelSinkBenchmark.kt
+++ b/benchmarks/src/jmh/kotlin/benchmarks/ChannelSinkBenchmark.kt
@@ -10,7 +10,7 @@ import org.openjdk.jmh.annotations.*
 import java.util.concurrent.*
 import kotlin.coroutines.*
 
-@Warmup(iterations = 5, time = 1)
+@Warmup(iterations = 7, time = 1)
 @Measurement(iterations = 5, time = 1)
 @BenchmarkMode(Mode.AverageTime)
 @OutputTimeUnit(TimeUnit.MILLISECONDS)
@@ -41,7 +41,7 @@ open class ChannelSinkBenchmark {
 
     private suspend inline fun run(context: CoroutineContext): Int {
         return Channel
-            .range(1, 1_000_000, context)
+            .range(1, 10_000, context)
             .filter(context) { it % 4 == 0 }
             .fold(0) { a, b -> a + b }
     }
diff --git a/benchmarks/src/jmh/kotlin/benchmarks/ChannelSinkDepthBenchmark.kt b/benchmarks/src/jmh/kotlin/benchmarks/ChannelSinkDepthBenchmark.kt
new file mode 100644
index 00000000..b8a970c8
--- /dev/null
+++ b/benchmarks/src/jmh/kotlin/benchmarks/ChannelSinkDepthBenchmark.kt
@@ -0,0 +1,91 @@
+/*
+ * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package benchmarks
+
+import kotlinx.coroutines.*
+import kotlinx.coroutines.channels.*
+import org.openjdk.jmh.annotations.*
+import java.util.concurrent.*
+import kotlin.coroutines.*
+
+@Warmup(iterations = 7, time = 1)
+@Measurement(iterations = 5, time = 1)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+@State(Scope.Benchmark)
+@Fork(2)
+open class ChannelSinkDepthBenchmark {
+    private val tl = ThreadLocal.withInitial({ 42 })
+
+    private val unconfinedOneElement = Dispatchers.Unconfined + tl.asContextElement()
+
+    @Benchmark
+    fun depth1(): Int = runBlocking {
+        run(1, unconfinedOneElement)
+    }
+
+    @Benchmark
+    fun depth10(): Int = runBlocking {
+        run(10, unconfinedOneElement)
+    }
+
+    @Benchmark
+    fun depth100(): Int = runBlocking {
+        run(100, unconfinedOneElement)
+    }
+
+    @Benchmark
+    fun depth1000(): Int = runBlocking {
+        run(1000, unconfinedOneElement)
+    }
+
+    private suspend inline fun run(callTraceDepth: Int, context: CoroutineContext): Int {
+        return Channel
+            .range(1, 10_000, context)
+            .filter(callTraceDepth, context) { it % 4 == 0 }
+            .fold(0) { a, b -> a + b }
+    }
+
+    private fun Channel.Factory.range(start: Int, count: Int, context: CoroutineContext) =
+        GlobalScope.produce(context) {
+            for (i in start until (start + count))
+                send(i)
+        }
+
+    // Migrated from deprecated operators, are good only for stressing channels
+
+    private fun ReceiveChannel<Int>.filter(
+        callTraceDepth: Int,
+        context: CoroutineContext = Dispatchers.Unconfined,
+        predicate: suspend (Int) -> Boolean
+    ): ReceiveChannel<Int> =
+        GlobalScope.produce(context, onCompletion = { cancel() }) {
+            deeplyNestedFilter(this, callTraceDepth, predicate)
+        }
+
+    private suspend fun ReceiveChannel<Int>.deeplyNestedFilter(
+        sink: ProducerScope<Int>,
+        depth: Int,
+        predicate: suspend (Int) -> Boolean
+    ) {
+        if (depth <= 1) {
+            for (e in this) {
+                if (predicate(e)) sink.send(e)
+            }
+        } else {
+            deeplyNestedFilter(sink, depth - 1, predicate)
+            require(true) // tail-call
+        }
+    }
+
+    private suspend inline fun <E, R> ReceiveChannel<E>.fold(initial: R, operation: (acc: R, E) -> R): R {
+        var accumulator = initial
+        consumeEach {
+            accumulator = operation(accumulator, it)
+        }
+        return accumulator
+    }
+}
+
diff --git a/kotlinx-coroutines-core/common/src/Builders.common.kt b/kotlinx-coroutines-core/common/src/Builders.common.kt
index 6ef1a8da..25ade352 100644
--- a/kotlinx-coroutines-core/common/src/Builders.common.kt
+++ b/kotlinx-coroutines-core/common/src/Builders.common.kt
@@ -207,25 +207,17 @@ private class LazyStandaloneCoroutine(
 }
 
 // Used by withContext when context changes, but dispatcher stays the same
-private class UndispatchedCoroutine<in T>(
+internal expect class UndispatchedCoroutine<in T>(
     context: CoroutineContext,
     uCont: Continuation<T>
-) : ScopeCoroutine<T>(context, uCont) {
-    override fun afterResume(state: Any?) {
-        // resume undispatched -- update context by stay on the same dispatcher
-        val result = recoverResult(state, uCont)
-        withCoroutineContext(uCont.context, null) {
-            uCont.resumeWith(result)
-        }
-    }
-}
+) : ScopeCoroutine<T>
 
 private const val UNDECIDED = 0
 private const val SUSPENDED = 1
 private const val RESUMED = 2
 
 // Used by withContext when context dispatcher changes
-private class DispatchedCoroutine<in T>(
+internal class DispatchedCoroutine<in T>(
     context: CoroutineContext,
     uCont: Continuation<T>
 ) : ScopeCoroutine<T>(context, uCont) {
diff --git a/kotlinx-coroutines-core/common/src/CoroutineContext.common.kt b/kotlinx-coroutines-core/common/src/CoroutineContext.common.kt
index 51374603..17ad66c1 100644
--- a/kotlinx-coroutines-core/common/src/CoroutineContext.common.kt
+++ b/kotlinx-coroutines-core/common/src/CoroutineContext.common.kt
@@ -19,5 +19,6 @@ internal expect val DefaultDelay: Delay
 
 // countOrElement -- pre-cached value for ThreadContext.kt
 internal expect inline fun <T> withCoroutineContext(context: CoroutineContext, countOrElement: Any?, block: () -> T): T
+internal expect inline fun <T> withContinuationContext(continuation: Continuation<*>, countOrElement: Any?, block: () -> T): T
 internal expect fun Continuation<*>.toDebugString(): String
 internal expect val CoroutineContext.coroutineName: String?
diff --git a/kotlinx-coroutines-core/common/src/internal/DispatchedContinuation.kt b/kotlinx-coroutines-core/common/src/internal/DispatchedContinuation.kt
index 1807e2ad..20b77bfe 100644
--- a/kotlinx-coroutines-core/common/src/internal/DispatchedContinuation.kt
+++ b/kotlinx-coroutines-core/common/src/internal/DispatchedContinuation.kt
@@ -235,7 +235,7 @@ internal class DispatchedContinuation<in T>(
 
     @Suppress("NOTHING_TO_INLINE") // we need it inline to save us an entry on the stack
     inline fun resumeUndispatchedWith(result: Result<T>) {
-        withCoroutineContext(context, countOrElement) {
+        withContinuationContext(continuation, countOrElement) {
             continuation.resumeWith(result)
         }
     }
diff --git a/kotlinx-coroutines-core/common/src/internal/DispatchedTask.kt b/kotlinx-coroutines-core/common/src/internal/DispatchedTask.kt
index caf87f14..ce05979d 100644
--- a/kotlinx-coroutines-core/common/src/internal/DispatchedTask.kt
+++ b/kotlinx-coroutines-core/common/src/internal/DispatchedTask.kt
@@ -85,9 +85,9 @@ internal abstract class DispatchedTask<in T>(
         try {
             val delegate = delegate as DispatchedContinuation<T>
             val continuation = delegate.continuation
-            val context = continuation.context
-            val state = takeState() // NOTE: Must take state in any case, even if cancelled
-            withCoroutineContext(context, delegate.countOrElement) {
+            withContinuationContext(continuation, delegate.countOrElement) {
+                val context = continuation.context
+                val state = takeState() // NOTE: Must take state in any case, even if cancelled
                 val exception = getExceptionalResult(state)
                 /*
                  * Check whether continuation was originally resumed with an exception.
diff --git a/kotlinx-coroutines-core/js/src/CoroutineContext.kt b/kotlinx-coroutines-core/js/src/CoroutineContext.kt
index c0b0c511..aed03277 100644
--- a/kotlinx-coroutines-core/js/src/CoroutineContext.kt
+++ b/kotlinx-coroutines-core/js/src/CoroutineContext.kt
@@ -4,6 +4,7 @@
 
 package kotlinx.coroutines
 
+import kotlinx.coroutines.internal.*
 import kotlin.browser.*
 import kotlin.coroutines.*
 
@@ -49,5 +50,13 @@ public actual fun CoroutineScope.newCoroutineContext(context: CoroutineContext):
 
 // No debugging facilities on JS
 internal actual inline fun <T> withCoroutineContext(context: CoroutineContext, countOrElement: Any?, block: () -> T): T = block()
+internal actual inline fun <T> withContinuationContext(continuation: Continuation<*>, countOrElement: Any?, block: () -> T): T = block()
 internal actual fun Continuation<*>.toDebugString(): String = toString()
 internal actual val CoroutineContext.coroutineName: String? get() = null // not supported on JS
+
+internal actual class UndispatchedCoroutine<in T> actual constructor(
+    context: CoroutineContext,
+    uCont: Continuation<T>
+) : ScopeCoroutine<T>(context, uCont) {
+    override fun afterResume(state: Any?) = uCont.resumeWith(recoverResult(state, uCont))
+}
diff --git a/kotlinx-coroutines-core/jvm/src/Builders.kt b/kotlinx-coroutines-core/jvm/src/Builders.kt
index 799267fe..c1b878ce 100644
--- a/kotlinx-coroutines-core/jvm/src/Builders.kt
+++ b/kotlinx-coroutines-core/jvm/src/Builders.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ * Copyright 2016-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */
 
 @file:JvmMultifileClass
diff --git a/kotlinx-coroutines-core/jvm/src/CoroutineContext.kt b/kotlinx-coroutines-core/jvm/src/CoroutineContext.kt
index c6540331..e91bb9fd 100644
--- a/kotlinx-coroutines-core/jvm/src/CoroutineContext.kt
+++ b/kotlinx-coroutines-core/jvm/src/CoroutineContext.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright 2016-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ * Copyright 2016-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */
 
 package kotlinx.coroutines
@@ -7,6 +7,7 @@ package kotlinx.coroutines
 import kotlinx.coroutines.internal.*
 import kotlinx.coroutines.scheduling.*
 import kotlin.coroutines.*
+import kotlin.coroutines.jvm.internal.CoroutineStackFrame
 
 internal const val COROUTINES_SCHEDULER_PROPERTY_NAME = "kotlinx.coroutines.scheduler"
 
@@ -47,6 +48,102 @@ internal actual inline fun <T> withCoroutineContext(context: CoroutineContext, c
     }
 }
 
+/**
+ * Executes a block using a context of a given continuation.
+ */
+internal actual inline fun <T> withContinuationContext(continuation: Continuation<*>, countOrElement: Any?, block: () -> T): T {
+    val context = continuation.context
+    val oldValue = updateThreadContext(context, countOrElement)
+    val undispatchedCompletion = if (oldValue !== NO_THREAD_ELEMENTS) {
+        // Only if some values were replaced we'll go to the slow path of figuring out where/how to restore them
+        continuation.updateUndispatchedCompletion(context, oldValue)
+    } else {
+        null // fast path -- don't even try to find undispatchedCompletion as there's nothing to restore in the context
+    }
+    try {
+        return block()
+    } finally {
+        if (undispatchedCompletion == null || undispatchedCompletion.clearThreadContext()) {
+            restoreThreadContext(context, oldValue)
+        }
+    }
+}
+
+internal fun Continuation<*>.updateUndispatchedCompletion(context: CoroutineContext, oldValue: Any?): UndispatchedCoroutine<*>? {
+    if (this !is CoroutineStackFrame) return null
+    /*
+     * Fast-path to detect whether we have unispatched coroutine at all in our stack.
+     *
+     * Implementation note.
+     * If we ever find that stackwalking for thread-locals is way too slow, here is another idea:
+     * 1) Store undispatched coroutine right in the `UndispatchedMarker` instance
+     * 2) To avoid issues with cross-dispatch boundary, remove `UndispatchedMarker`
+     *    from the context when creating dispatched coroutine in `withContext`.
+     *    Another option is to "unmark it" instead of removing to save an allocation.
+     *    Both options should work, but it requires more careful studying of the performance
+     *    and, mostly, maintainability impact.
+     */
+    val potentiallyHasUndispatchedCorotuine = context[UndispatchedMarker] !== null
+    if (!potentiallyHasUndispatchedCorotuine) return null
+    val completion = undispatchedCompletion()
+    completion?.saveThreadContext(context, oldValue)
+    return completion
+}
+
+internal tailrec fun CoroutineStackFrame.undispatchedCompletion(): UndispatchedCoroutine<*>? {
+    // Find direct completion of this continuation
+    val completion: CoroutineStackFrame = when (this) {
+        is DispatchedCoroutine<*> -> return null
+        else -> callerFrame ?: return null // something else -- not supported
+    }
+    if (completion is UndispatchedCoroutine<*>) return completion // found UndispatchedCoroutine!
+    return completion.undispatchedCompletion() // walk up the call stack with tail call
+}
+
+/**
+ * Marker indicating that [UndispatchedCoroutine] exists somewhere up in the stack.
+ * Used as a performance optimization to avoid stack walking where it is not nesessary.
+ */
+private object UndispatchedMarker: CoroutineContext.Element, CoroutineContext.Key<UndispatchedMarker> {
+    override val key: CoroutineContext.Key<*>
+        get() = this
+}
+
+// Used by withContext when context changes, but dispatcher stays the same
+internal actual class UndispatchedCoroutine<in T>actual constructor (
+    context: CoroutineContext,
+    uCont: Continuation<T>
+) : ScopeCoroutine<T>(if (context[UndispatchedMarker] == null) context + UndispatchedMarker else context, uCont) {
+
+    private var savedContext: CoroutineContext? = null
+    private var savedOldValue: Any? = null
+
+    fun saveThreadContext(context: CoroutineContext, oldValue: Any?) {
+        savedContext = context
+        savedOldValue = oldValue
+    }
+
+    fun clearThreadContext(): Boolean {
+        if (savedContext == null) return false
+        savedContext = null
+        savedOldValue = null
+        return true
+    }
+
+    override fun afterResume(state: Any?) {
+        savedContext?.let { context ->
+            restoreThreadContext(context, savedOldValue)
+            savedContext = null
+            savedOldValue = null
+        }
+        // resume undispatched -- update context but stay on the same dispatcher
+        val result = recoverResult(state, uCont)
+        withContinuationContext(uCont, null) {
+            uCont.resumeWith(result)
+        }
+    }
+}
+
 internal actual val CoroutineContext.coroutineName: String? get() {
     if (!DEBUG) return null
     val coroutineId = this[CoroutineId] ?: return null
diff --git a/kotlinx-coroutines-core/jvm/src/internal/ThreadContext.kt b/kotlinx-coroutines-core/jvm/src/internal/ThreadContext.kt
index 9d9d30e4..18c2ce04 100644
--- a/kotlinx-coroutines-core/jvm/src/internal/ThreadContext.kt
+++ b/kotlinx-coroutines-core/jvm/src/internal/ThreadContext.kt
@@ -7,8 +7,8 @@ package kotlinx.coroutines.internal
 import kotlinx.coroutines.*
 import kotlin.coroutines.*
 
-
-private val ZERO = Symbol("ZERO")
+@JvmField
+internal val NO_THREAD_ELEMENTS = Symbol("NO_THREAD_ELEMENTS")
 
 // Used when there are >= 2 active elements in the context
 private class ThreadState(val context: CoroutineContext, n: Int) {
@@ -60,12 +60,13 @@ private val restoreState =
 internal actual fun threadContextElements(context: CoroutineContext): Any = context.fold(0, countAll)!!
 
 // countOrElement is pre-cached in dispatched continuation
+// returns NO_THREAD_ELEMENTS if the contest does not have any ThreadContextElements
 internal fun updateThreadContext(context: CoroutineContext, countOrElement: Any?): Any? {
     @Suppress("NAME_SHADOWING")
     val countOrElement = countOrElement ?: threadContextElements(context)
     @Suppress("IMPLICIT_BOXING_IN_IDENTITY_EQUALS")
     return when {
-        countOrElement === 0 -> ZERO // very fast path when there are no active ThreadContextElements
+        countOrElement === 0 -> NO_THREAD_ELEMENTS // very fast path when there are no active ThreadContextElements
         //    ^^^ identity comparison for speed, we know zero always has the same identity
         countOrElement is Int -> {
             // slow path for multiple active ThreadContextElements, allocates ThreadState for multiple old values
@@ -82,7 +83,7 @@ internal fun updateThreadContext(context: CoroutineContext, countOrElement: Any?
 
 internal fun restoreThreadContext(context: CoroutineContext, oldState: Any?) {
     when {
-        oldState === ZERO -> return // very fast path when there are no ThreadContextElements
+        oldState === NO_THREAD_ELEMENTS -> return // very fast path when there are no ThreadContextElements
         oldState is ThreadState -> {
             // slow path with multiple stored ThreadContextElements
             oldState.start()
diff --git a/kotlinx-coroutines-core/jvm/test/ThreadContextElementRestoreTest.kt b/kotlinx-coroutines-core/jvm/test/ThreadContextElementRestoreTest.kt
new file mode 100644
index 00000000..e2ab4d72
--- /dev/null
+++ b/kotlinx-coroutines-core/jvm/test/ThreadContextElementRestoreTest.kt
@@ -0,0 +1,198 @@
+/*
+ * Copyright 2016-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package kotlinx.coroutines
+
+import org.junit.Test
+import kotlin.coroutines.*
+import kotlin.test.*
+
+class ThreadContextElementRestoreTest : TestBase() {
+    private val tl = ThreadLocal<String?>()
+
+    // Checks that ThreadLocal context is properly restored after executing the given block inside
+    // withContext(tl.asContextElement("OK")) code running in different outer contexts
+    private inline fun check(crossinline block: suspend () -> Unit) = runTest {
+        val mainDispatcher = coroutineContext[ContinuationInterceptor] as CoroutineDispatcher
+        // Scenario #1: withContext(ThreadLocal) direct from runTest
+        withContext(tl.asContextElement("OK")) {
+            block()
+            assertEquals("OK", tl.get())
+        }
+        assertEquals(null, tl.get())
+        // Scenario #2: withContext(ThreadLocal) from coroutineScope
+        coroutineScope {
+            withContext(tl.asContextElement("OK")) {
+                block()
+                assertEquals("OK", tl.get())
+            }
+            assertEquals(null, tl.get())
+        }
+        // Scenario #3: withContext(ThreadLocal) from undispatched withContext
+        withContext(CoroutineName("NAME")) {
+            withContext(tl.asContextElement("OK")) {
+                block()
+                assertEquals("OK", tl.get())
+            }
+            assertEquals(null, tl.get())
+        }
+        // Scenario #4: withContext(ThreadLocal) from dispatched withContext
+        withContext(wrapperDispatcher()) {
+            withContext(tl.asContextElement("OK")) {
+                block()
+                assertEquals("OK", tl.get())
+            }
+            assertEquals(null, tl.get())
+        }
+        // Scenario #5: withContext(ThreadLocal) from withContext(ThreadLocal)
+        withContext(tl.asContextElement(null)) {
+            withContext(tl.asContextElement("OK")) {
+                block()
+                assertEquals("OK", tl.get())
+            }
+            assertEquals(null, tl.get())
+        }
+        // Scenario #6: withContext(ThreadLocal) from withTimeout
+        withTimeout(1000) {
+            withContext(tl.asContextElement("OK")) {
+                block()
+                assertEquals("OK", tl.get())
+            }
+            assertEquals(null, tl.get())
+        }
+        // Scenario #7: withContext(ThreadLocal) from withContext(Unconfined)
+        withContext(Dispatchers.Unconfined) {
+            withContext(tl.asContextElement("OK")) {
+                block()
+                assertEquals("OK", tl.get())
+            }
+            assertEquals(null, tl.get())
+        }
+        // Scenario #8: withContext(ThreadLocal) from withContext(Default)
+        withContext(Dispatchers.Default) {
+            withContext(tl.asContextElement("OK")) {
+                block()
+                assertEquals("OK", tl.get())
+            }
+            assertEquals(null, tl.get())
+        }
+        // Scenario #9: withContext(ThreadLocal) from withContext(mainDispatcher)
+        withContext(mainDispatcher) {
+            withContext(tl.asContextElement("OK")) {
+                block()
+                assertEquals("OK", tl.get())
+            }
+            assertEquals(null, tl.get())
+        }
+    }
+
+    @Test
+    fun testSimpleNoSuspend() =
+        check {}
+
+    @Test
+    fun testSimpleDelay() = check {
+        delay(1)
+    }
+
+    @Test
+    fun testSimpleYield() = check {
+        yield()
+    }
+
+    private suspend fun deepDelay() {
+        deepDelay2(); deepDelay2()
+    }
+
+    private suspend fun deepDelay2() {
+        delay(1); delay(1)
+    }
+
+    @Test
+    fun testDeepDelay() = check {
+        deepDelay()
+    }
+
+    private suspend fun deepYield() {
+        deepYield2(); deepYield2()
+    }
+
+    private suspend fun deepYield2() {
+        yield(); yield()
+    }
+
+    @Test
+    fun testDeepYield() = check {
+        deepYield()
+    }
+
+    @Test
+    fun testCoroutineScopeDelay() = check {
+        coroutineScope {
+            delay(1)
+        }
+    }
+
+    @Test
+    fun testCoroutineScopeYield() = check {
+        coroutineScope {
+            yield()
+        }
+    }
+
+    @Test
+    fun testWithContextUndispatchedDelay() = check {
+        withContext(CoroutineName("INNER")) {
+            delay(1)
+        }
+    }
+
+    @Test
+    fun testWithContextUndispatchedYield() = check {
+        withContext(CoroutineName("INNER")) {
+            yield()
+        }
+    }
+
+    @Test
+    fun testWithContextDispatchedDelay() = check {
+        withContext(wrapperDispatcher()) {
+            delay(1)
+        }
+    }
+
+    @Test
+    fun testWithContextDispatchedYield() = check {
+        withContext(wrapperDispatcher()) {
+            yield()
+        }
+    }
+
+    @Test
+    fun testWithTimeoutDelay() = check {
+        withTimeout(1000) {
+            delay(1)
+        }
+    }
+
+    @Test
+    fun testWithTimeoutYield() = check {
+        withTimeout(1000) {
+            yield()
+        }
+    }
+
+    @Test
+    fun testWithUnconfinedContextDelay() = check {
+        withContext(Dispatchers.Unconfined) {
+            delay(1)
+        }
+    }
+    @Test
+    fun testWithUnconfinedContextYield() = check {
+        withContext(Dispatchers.Unconfined) {
+            yield()
+        }
+    }
+}
diff --git a/kotlinx-coroutines-core/native/src/CoroutineContext.kt b/kotlinx-coroutines-core/native/src/CoroutineContext.kt
index 4ec1289e..86ffa8de 100644
--- a/kotlinx-coroutines-core/native/src/CoroutineContext.kt
+++ b/kotlinx-coroutines-core/native/src/CoroutineContext.kt
@@ -4,6 +4,7 @@
 
 package kotlinx.coroutines
 
+import kotlinx.coroutines.internal.*
 import kotlin.coroutines.*
 import kotlin.native.concurrent.*
 
@@ -38,5 +39,13 @@ public actual fun CoroutineScope.newCoroutineContext(context: CoroutineContext):
 
 // No debugging facilities on native
 internal actual inline fun <T> withCoroutineContext(context: CoroutineContext, countOrElement: Any?, block: () -> T): T = block()
+internal actual inline fun <T> withContinuationContext(continuation: Continuation<*>, countOrElement: Any?, block: () -> T): T = block()
 internal actual fun Continuation<*>.toDebugString(): String = toString()
 internal actual val CoroutineContext.coroutineName: String? get() = null // not supported on native
+
+internal actual class UndispatchedCoroutine<in T> actual constructor(
+    context: CoroutineContext,
+    uCont: Continuation<T>
+) : ScopeCoroutine<T>(context, uCont) {
+    override fun afterResume(state: Any?) = uCont.resumeWith(recoverResult(state, uCont))
+}
\ No newline at end of file