Fix native-mt ensureNeverFrozen with coroutine launch

Fixes #2064

Author: elizarov

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

@@ -137,7 +137,6 @@ public abstract class AbstractCoroutine<in T>(
      * * [LAZY] does nothing.
      */
     public fun <R> start(start: CoroutineStart, receiver: R, block: suspend R.() -> T) {
-        initParentJob()
-        startCoroutine(start, receiver, this, block)
+        startAbstractCoroutine(start, receiver, this, block)
     }
 }

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

@@ -156,8 +156,7 @@ public suspend fun <T> withContext(
     }
     // SLOW PATH -- use new dispatcher
     val coroutine = DispatchedCoroutine(newContext, uCont)
-    coroutine.initParentJob()
-    startCoroutine(CoroutineStart.DEFAULT, coroutine, coroutine, block)
+    coroutine.start(CoroutineStart.DEFAULT, coroutine, block)
     coroutine.getResult()
 }
 
@@ -194,6 +193,14 @@ internal expect fun <T, R> startCoroutine(
     block: suspend R.() -> T
 )
 
+// initParentJob + startCoroutine
+internal expect fun <T, R> startAbstractCoroutine(
+    start: CoroutineStart,
+    receiver: R,
+    coroutine: AbstractCoroutine<T>,
+    block: suspend R.() -> T
+)
+
 /**
  * On JVM & JS lazy coroutines are eagerly started (to record creation trace), the started later.
  * On Native the block is saved so that it can be shared with another worker, the created and started later.

kotlinx-coroutines-core/js/src/Builders.kt

@@ -17,6 +17,17 @@ internal actual inline fun <T, R> startCoroutine(
 ) =
     startCoroutineImpl(start, receiver, completion, block)
 
+@Suppress("NOTHING_TO_INLINE") // Save an entry on call stack
+internal actual inline fun <T, R> startAbstractCoroutine(
+    start: CoroutineStart,
+    receiver: R,
+    coroutine: AbstractCoroutine<T>,
+    noinline block: suspend R.() -> T
+) {
+    coroutine.initParentJob()
+    startCoroutineImpl(start, receiver, coroutine, block)
+}
+
 @Suppress("NOTHING_TO_INLINE") // Save an entry on call stack
 internal actual inline fun <T, R> saveLazyCoroutine(
     coroutine: AbstractCoroutine<T>,

kotlinx-coroutines-core/jvm/src/Builders.kt

@@ -105,6 +105,17 @@ internal actual inline fun <T, R> startCoroutine(
 ) =
     startCoroutineImpl(start, receiver, completion, block)
 
+@Suppress("NOTHING_TO_INLINE") // Save an entry on call stack
+internal actual inline fun <T, R> startAbstractCoroutine(
+    start: CoroutineStart,
+    receiver: R,
+    coroutine: AbstractCoroutine<T>,
+    noinline block: suspend R.() -> T
+) {
+    coroutine.initParentJob()
+    startCoroutineImpl(start, receiver, coroutine, block)
+}
+
 @Suppress("NOTHING_TO_INLINE") // Save an entry on call stack
 internal actual inline fun <T, R> saveLazyCoroutine(
     coroutine: AbstractCoroutine<T>,

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

@@ -88,25 +88,52 @@ internal fun runEventLoop(eventLoop: EventLoop?, isCompleted: () -> Boolean) {
 
 // --------------- Kotlin/Native specialization hooks ---------------
 
+// just start
 internal actual fun <T, R> startCoroutine(
     start: CoroutineStart,
     receiver: R,
     completion: Continuation<T>,
     block: suspend R.() -> T
+) =
+    startCoroutine(start, receiver, completion, block) {}
+
+// initParentJob + startCoroutine
+internal actual fun <T, R> startAbstractCoroutine(
+    start: CoroutineStart,
+    receiver: R,
+    coroutine: AbstractCoroutine<T>,
+    block: suspend R.() -> T
+) {
+    // See https://github.com/Kotlin/kotlinx.coroutines/issues/2064
+    // We shall do initParentJob only after freezing the block
+    startCoroutine(start, receiver, coroutine, block) {
+        coroutine.initParentJob()
+    }
+}
+
+private fun <T, R> startCoroutine(
+    start: CoroutineStart,
+    receiver: R,
+    completion: Continuation<T>,
+    block: suspend R.() -> T,
+    initParentJobIfNeeded: () -> Unit
 ) {
     val curThread = currentThread()
     val newThread = completion.context[ContinuationInterceptor].thread()
     if (newThread != curThread) {
         check(start != CoroutineStart.UNDISPATCHED) {
             "Cannot start an undispatched coroutine in another thread $newThread from current $curThread"
         }
+        block.freeze() // freeze the block, safely get FreezingException if it cannot be frozen
+        initParentJobIfNeeded() // only initParentJob here if needed
         if (start != CoroutineStart.LAZY) {
             newThread.execute {
                 startCoroutineImpl(start, receiver, completion, block)
             }
         }
         return
     }
+    initParentJobIfNeeded()
     startCoroutineImpl(start, receiver, completion, block)
 }
 

kotlinx-coroutines-core/native/test/WorkerDispatcherTest.kt

@@ -326,5 +326,17 @@ class WorkerDispatcherTest : TestBase() {
         finish(2)
     }
 
+    @Test
+    fun testEnsureNeverFrozenWithContext() = runTest {
+        expect(1)
+        val x = Data("OK")
+        x.ensureNeverFrozen()
+        assertFailsWith<FreezingException> {
+            val s = withContext(dispatcher) { x.s }
+            println(s) // does not actually execute
+        }
+        finish(2)
+    }
+
     private data class Data(val s: String)
 }