Make MPP ThreadContextElement compile and run

Author: dkhalanskyjb

Diff for: kotlinx-coroutines-core/api/kotlinx-coroutines-core.klib.api

@@ -186,6 +186,16 @@ abstract interface <#A: kotlin/Any?> kotlinx.coroutines/CompletableDeferred : ko
     abstract fun completeExceptionally(kotlin/Throwable): kotlin/Boolean // kotlinx.coroutines/CompletableDeferred.completeExceptionally|completeExceptionally(kotlin.Throwable){}[0]
 }
 
+abstract interface <#A: kotlin/Any?> kotlinx.coroutines/CopyableThreadContextElement : kotlinx.coroutines/ThreadContextElement<#A> { // kotlinx.coroutines/CopyableThreadContextElement|null[0]
+    abstract fun copyForChild(): kotlinx.coroutines/CopyableThreadContextElement<#A> // kotlinx.coroutines/CopyableThreadContextElement.copyForChild|copyForChild(){}[0]
+    abstract fun mergeForChild(kotlin.coroutines/CoroutineContext.Element): kotlin.coroutines/CoroutineContext // kotlinx.coroutines/CopyableThreadContextElement.mergeForChild|mergeForChild(kotlin.coroutines.CoroutineContext.Element){}[0]
+}
+
+abstract interface <#A: kotlin/Any?> kotlinx.coroutines/ThreadContextElement : kotlin.coroutines/CoroutineContext.Element { // kotlinx.coroutines/ThreadContextElement|null[0]
+    abstract fun restoreThreadContext(kotlin.coroutines/CoroutineContext, #A) // kotlinx.coroutines/ThreadContextElement.restoreThreadContext|restoreThreadContext(kotlin.coroutines.CoroutineContext;1:0){}[0]
+    abstract fun updateThreadContext(kotlin.coroutines/CoroutineContext): #A // kotlinx.coroutines/ThreadContextElement.updateThreadContext|updateThreadContext(kotlin.coroutines.CoroutineContext){}[0]
+}
+
 abstract interface <#A: kotlin/Throwable & kotlinx.coroutines/CopyableThrowable<#A>> kotlinx.coroutines/CopyableThrowable { // kotlinx.coroutines/CopyableThrowable|null[0]
     abstract fun createCopy(): #A? // kotlinx.coroutines/CopyableThrowable.createCopy|createCopy(){}[0]
 }

Diff for: kotlinx-coroutines-core/common/src/Builders.common.kt

@@ -9,6 +9,7 @@ import kotlinx.atomicfu.*
 import kotlinx.coroutines.internal.*
 import kotlinx.coroutines.intrinsics.*
 import kotlinx.coroutines.selects.*
+import kotlin.concurrent.Volatile
 import kotlin.contracts.*
 import kotlin.coroutines.*
 import kotlin.coroutines.intrinsics.*
@@ -206,13 +207,7 @@ private class LazyStandaloneCoroutine(
 }
 
 // Used by withContext when context changes, but dispatcher stays the same
-internal expect class UndispatchedCoroutine<in T>(
-    context: CoroutineContext,
-    uCont: Continuation<T>
-) : ScopeCoroutine<T>
-
-// Used by withContext when context changes, but dispatcher stays the same
-internal actual class UndispatchedCoroutine<in T>actual constructor (
+internal class UndispatchedCoroutine<in T>(
     context: CoroutineContext,
     uCont: Continuation<T>
 ) : ScopeCoroutine<T>(if (context[UndispatchedMarker] == null) context + UndispatchedMarker else context, uCont) {
@@ -249,7 +244,7 @@ internal actual class UndispatchedCoroutine<in T>actual constructor (
      * - It's never accessed when we are sure there are no thread context elements
      * - It's cleaned up via [ThreadLocal.remove] as soon as the coroutine is suspended or finished.
      */
-    private val threadStateToRecover = ThreadLocal<Pair<CoroutineContext, Any?>>()
+    private val threadStateToRecover = commonThreadLocal<Pair<CoroutineContext, Any?>?>(Symbol("UndispatchedCoroutine"))
 
     /*
      * Indicates that a coroutine has at least one thread context element associated with it

Diff for: kotlinx-coroutines-core/common/src/CoroutineContext.common.kt

@@ -3,34 +3,18 @@ package kotlinx.coroutines
 import kotlinx.coroutines.internal.*
 import kotlin.coroutines.*
 
-/**
- * Creates a context for a new coroutine. It installs [Dispatchers.Default] when no other dispatcher or
- * [ContinuationInterceptor] is specified and adds optional support for debugging facilities (when turned on)
- * and copyable-thread-local facilities on JVM.
- */
-public expect fun CoroutineScope.newCoroutineContext(context: CoroutineContext): CoroutineContext
-
-/**
- * Creates a context for coroutine builder functions that do not launch a new coroutine, e.g. [withContext].
- * @suppress
- */
-@InternalCoroutinesApi
-public expect fun CoroutineContext.newCoroutineContext(addedContext: CoroutineContext): CoroutineContext
-
 @PublishedApi // to have unmangled name when using from other modules via suppress
 @Suppress("PropertyName")
 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?
+internal expect fun wrapContextWithDebug(context: CoroutineContext): CoroutineContext
 
 /**
  * Executes a block using a given coroutine context.
  */
-internal actual inline fun <T> withCoroutineContext(context: CoroutineContext, countOrElement: Any?, block: () -> T): T {
+internal inline fun <T> withCoroutineContext(context: CoroutineContext, countOrElement: Any?, block: () -> T): T {
     val oldValue = updateThreadContext(context, countOrElement)
     try {
         return block()
@@ -42,7 +26,7 @@ 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 {
+internal 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) {
@@ -60,7 +44,7 @@ internal actual inline fun <T> withContinuationContext(continuation: Continuatio
     }
 }
 
-internal fun Continuation<*>.updateUndispatchedCompletion(context: CoroutineContext, oldValue: Any?): UndispatchedCoroutine<*>? {
+private fun Continuation<*>.updateUndispatchedCompletion(context: CoroutineContext, oldValue: Any?): UndispatchedCoroutine<*>? {
     if (this !is CoroutineStackFrame) return null
     /*
      * Fast-path to detect whether we have undispatched coroutine at all in our stack.
@@ -81,7 +65,7 @@ internal fun Continuation<*>.updateUndispatchedCompletion(context: CoroutineCont
     return completion
 }
 
-internal tailrec fun CoroutineStackFrame.undispatchedCompletion(): UndispatchedCoroutine<*>? {
+private tailrec fun CoroutineStackFrame.undispatchedCompletion(): UndispatchedCoroutine<*>? {
     // Find direct completion of this continuation
     val completion: CoroutineStackFrame = when (this) {
         is DispatchedCoroutine<*> -> return null
@@ -95,7 +79,7 @@ internal tailrec fun CoroutineStackFrame.undispatchedCompletion(): UndispatchedC
  * Marker indicating that [UndispatchedCoroutine] exists somewhere up in the stack.
  * Used as a performance optimization to avoid stack walking where it is not necessary.
  */
-private object UndispatchedMarker: CoroutineContext.Element, CoroutineContext.Key<UndispatchedMarker> {
+internal object UndispatchedMarker: CoroutineContext.Element, CoroutineContext.Key<UndispatchedMarker> {
     override val key: CoroutineContext.Key<*>
         get() = this
 }

Diff for: kotlinx-coroutines-core/common/src/CoroutineContext.sharedWithJvm.kt

@@ -1,19 +1,22 @@
+// This file should be a part of `CoroutineContext.common.kt`, but adding `JvmName` to that fails: KT-75248
+@file:JvmName("CoroutineContextKt")
+@file:JvmMultifileClass
 package kotlinx.coroutines
 
+import kotlin.coroutines.ContinuationInterceptor
 import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.EmptyCoroutineContext
-
+import kotlin.jvm.JvmMultifileClass
+import kotlin.jvm.JvmName
 
 /**
  * Creates a context for a new coroutine. It installs [Dispatchers.Default] when no other dispatcher or
- * [ContinuationInterceptor] is specified and adds optional support for debugging facilities (when turned on)
- * and copyable-thread-local facilities on JVM.
- * See [DEBUG_PROPERTY_NAME] for description of debugging facilities on JVM.
+ * [ContinuationInterceptor] is specified and
  */
 @ExperimentalCoroutinesApi
-public actual fun CoroutineScope.newCoroutineContext(context: CoroutineContext): CoroutineContext {
+public fun CoroutineScope.newCoroutineContext(context: CoroutineContext): CoroutineContext {
     val combined = foldCopies(coroutineContext, context, true)
-    val debug = if (DEBUG) combined + CoroutineId(COROUTINE_ID.incrementAndGet()) else combined
+    val debug = wrapContextWithDebug(combined)
     return if (combined !== Dispatchers.Default && combined[ContinuationInterceptor] == null)
         debug + Dispatchers.Default else debug
 }
@@ -23,7 +26,7 @@ public actual fun CoroutineScope.newCoroutineContext(context: CoroutineContext):
  * @suppress
  */
 @InternalCoroutinesApi
-public actual fun CoroutineContext.newCoroutineContext(addedContext: CoroutineContext): CoroutineContext {
+public fun CoroutineContext.newCoroutineContext(addedContext: CoroutineContext): CoroutineContext {
     /*
      * Fast-path: we only have to copy/merge if 'addedContext' (which typically has one or two elements)
      * contains copyable elements.

Diff for: kotlinx-coroutines-core/common/src/internal/ThreadContext.common.kt

@@ -4,8 +4,6 @@ import kotlinx.coroutines.ThreadContextElement
 import kotlin.coroutines.*
 import kotlin.jvm.JvmField
 
-internal expect fun threadContextElements(context: CoroutineContext): Any
-
 @JvmField
 internal val NO_THREAD_ELEMENTS = Symbol("NO_THREAD_ELEMENTS")
 
@@ -29,9 +27,9 @@ private class ThreadState(@JvmField val context: CoroutineContext, n: Int) {
 }
 
 // Counts ThreadContextElements in the context
-// Any? here is Int | ThreadContextElement (when count is one)
+// Any here is Int | ThreadContextElement (when count is one)
 private val countAll =
-    fun (countOrElement: Any?, element: CoroutineContext.Element): Any? {
+    fun (countOrElement: Any, element: CoroutineContext.Element): Any {
         if (element is ThreadContextElement<*>) {
             val inCount = countOrElement as? Int ?: 1
             return if (inCount == 0) element else inCount + 1
@@ -55,17 +53,15 @@ private val updateState =
         return state
     }
 
-internal actual fun threadContextElements(context: CoroutineContext): Any = context.fold(0, countAll)!!
+internal expect inline fun isZeroCount(countOrElement: Any?): Boolean
 
 // 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 -> NO_THREAD_ELEMENTS // very fast path when there are no active ThreadContextElements
-        //    ^^^ identity comparison for speed, we know zero always has the same identity
+        isZeroCount(countOrElement) -> NO_THREAD_ELEMENTS // very fast path when there are no active ThreadContextElements
         countOrElement is Int -> {
             // slow path for multiple active ThreadContextElements, allocates ThreadState for multiple old values
             context.fold(ThreadState(context, countOrElement), updateState)
@@ -94,3 +90,5 @@ internal fun restoreThreadContext(context: CoroutineContext, oldState: Any?) {
         }
     }
 }
+
+internal fun threadContextElements(context: CoroutineContext): Any = context.fold(0, countAll)

Diff for: kotlinx-coroutines-core/common/src/internal/ThreadLocal.common.kt

@@ -3,6 +3,7 @@ package kotlinx.coroutines.internal
 internal expect class CommonThreadLocal<T> {
     fun get(): T
     fun set(value: T)
+    fun remove()
 }
 
 /**

Diff for: kotlinx-coroutines-core/common/test/ThreadContextElementTest.kt

@@ -1,8 +1,13 @@
 package kotlinx.coroutines
 
+import kotlinx.atomicfu.atomic
+import kotlinx.atomicfu.loop
 import kotlinx.coroutines.testing.*
 import kotlin.test.*
 import kotlinx.coroutines.flow.*
+import kotlinx.coroutines.internal.CommonThreadLocal
+import kotlinx.coroutines.internal.Symbol
+import kotlinx.coroutines.internal.commonThreadLocal
 import kotlin.coroutines.*
 
 class ThreadContextElementTest: TestBase() {
@@ -38,20 +43,18 @@ class ThreadContextElementTest: TestBase() {
      */
     @Test
     fun testWithContextJobAccess() = runTest {
-        val executor = Executors.newSingleThreadExecutor()
         // Emulate non-equal dispatchers
-        val executor1 = object : ExecutorService by executor {}
-        val executor2 = object : ExecutorService by executor {}
-        val dispatcher1 = executor1.asCoroutineDispatcher()
-        val dispatcher2 = executor2.asCoroutineDispatcher()
+        val dispatcher = Dispatchers.Default.limitedParallelism(1)
+        val dispatcher1 = dispatcher.limitedParallelism(1, "dispatcher1")
+        val dispatcher2 = dispatcher.limitedParallelism(1, "dispatcher2")
         val captor = JobCaptor()
         val manuallyCaptured = mutableListOf<String>()
 
         fun registerUpdate(job: Job?) = manuallyCaptured.add("Update: $job")
         fun registerRestore(job: Job?) = manuallyCaptured.add("Restore: $job")
 
         var rootJob: Job? = null
-        runBlocking(captor + dispatcher1) {
+        withContext(captor + dispatcher1) {
             rootJob = coroutineContext.job
             registerUpdate(rootJob)
             var undispatchedJob: Job? = null
@@ -84,7 +87,6 @@ class ThreadContextElementTest: TestBase() {
         val expected = manuallyCaptured.filter { it.startsWith("Update: ") }.joinToString(separator = "\n")
         val actual = captor.capturees.filter { it.startsWith("Update: ") }.joinToString(separator = "\n")
         assertEquals(expected, actual)
-        executor.shutdownNow()
     }
 
     @Test
@@ -96,7 +98,7 @@ class ThreadContextElementTest: TestBase() {
             assertEquals(myData, myThreadLocal.get())
             emit(1)
         }
-            .flowOn(myThreadLocal.asContextElement() + Dispatchers.Default)
+            .flowOn(myThreadLocal.asCtxElement() + Dispatchers.Default)
             .single()
         myThreadLocal.set(null)
         finish(2)
@@ -105,7 +107,7 @@ class ThreadContextElementTest: TestBase() {
 
 class MyData
 
-class JobCaptor(val capturees: MutableList<String> = CopyOnWriteArrayList()) : ThreadContextElement<Unit> {
+private class JobCaptor(val capturees: CopyOnWriteList<String> = CopyOnWriteList()) : ThreadContextElement<Unit> {
 
     companion object Key : CoroutineContext.Key<MyElement>
 
@@ -121,7 +123,7 @@ class JobCaptor(val capturees: MutableList<String> = CopyOnWriteArrayList()) : T
 }
 
 // declare thread local variable holding MyData
-private val myThreadLocal = ThreadLocal<MyData?>()
+internal val myThreadLocal = commonThreadLocal<MyData?>(Symbol("myElement"))
 
 // declare context element holding MyData
 class MyElement(val data: MyData) : ThreadContextElement<MyData?> {
@@ -144,3 +146,44 @@ class MyElement(val data: MyData) : ThreadContextElement<MyData?> {
         myThreadLocal.set(oldState)
     }
 }
+
+
+private class CommonThreadLocalContextElement<T>(
+    private val threadLocal: CommonThreadLocal<T>,
+    private val value: T = threadLocal.get()
+): ThreadContextElement<T>, CoroutineContext.Key<CommonThreadLocalContextElement<T>> {
+    // provide the key of the corresponding context element
+    override val key: CoroutineContext.Key<CommonThreadLocalContextElement<T>>
+        get() = this
+
+    // this is invoked before coroutine is resumed on current thread
+    override fun updateThreadContext(context: CoroutineContext): T {
+        val oldState = threadLocal.get()
+        threadLocal.set(value)
+        return oldState
+    }
+
+    // this is invoked after coroutine has suspended on current thread
+    override fun restoreThreadContext(context: CoroutineContext, oldState: T) {
+        threadLocal.set(oldState)
+    }
+}
+
+// overload resolution issues if this is called `asContextElement`
+internal fun <T> CommonThreadLocal<T>.asCtxElement(value: T = get()): ThreadContextElement<T> =
+    CommonThreadLocalContextElement(this, value)
+
+private class CopyOnWriteList<T> private constructor(list: List<T>) {
+    private val field = atomic(list)
+
+    constructor() : this(emptyList())
+
+    fun add(value: T) {
+        field.loop { current ->
+            val new = current + value
+            if (field.compareAndSet(current, new)) return
+        }
+    }
+
+    fun filter(predicate: (T) -> Boolean): List<T> = field.value.filter(predicate)
+}

Diff for: kotlinx-coroutines-core/common/test/ThreadContextMutableCopiesTest.kt

@@ -2,12 +2,16 @@ package kotlinx.coroutines
 
 import kotlinx.coroutines.testing.*
 import kotlinx.coroutines.flow.*
+import kotlinx.coroutines.internal.Symbol
+import kotlinx.coroutines.internal.commonThreadLocal
 import kotlin.coroutines.*
 import kotlin.test.*
 
 class ThreadContextMutableCopiesTest : TestBase() {
     companion object {
-        val threadLocalData: ThreadLocal<MutableList<String>> = ThreadLocal.withInitial { ArrayList() }
+        internal val threadLocalData = commonThreadLocal<MutableList<String>>(Symbol("ThreadLocalData")).also {
+            it.set(mutableListOf())
+        }
     }
 
     class MyMutableElement(
@@ -42,7 +46,7 @@ class ThreadContextMutableCopiesTest : TestBase() {
     @Test
     fun testDataIsCopied() = runTest {
         val root = MyMutableElement(ArrayList())
-        runBlocking(root) {
+        launch(root) {
             val data = threadLocalData.get()
             expect(1)
             launch(root) {
@@ -56,7 +60,7 @@ class ThreadContextMutableCopiesTest : TestBase() {
     @Test
     fun testDataIsNotOverwritten() = runTest {
         val root = MyMutableElement(ArrayList())
-        runBlocking(root) {
+        withContext(root) {
             expect(1)
             val originalData = threadLocalData.get()
             threadLocalData.get().add("X")
@@ -75,7 +79,7 @@ class ThreadContextMutableCopiesTest : TestBase() {
     @Test
     fun testDataIsMerged() = runTest {
         val root = MyMutableElement(ArrayList())
-        runBlocking(root) {
+        withContext(root) {
             expect(1)
             val originalData = threadLocalData.get()
             threadLocalData.get().add("X")
@@ -94,7 +98,7 @@ class ThreadContextMutableCopiesTest : TestBase() {
     @Test
     fun testDataIsNotOverwrittenWithContext() = runTest {
         val root = MyMutableElement(ArrayList())
-        runBlocking(root) {
+        withContext(root) {
             val originalData = threadLocalData.get()
             threadLocalData.get().add("X")
             expect(1)
@@ -114,7 +118,7 @@ class ThreadContextMutableCopiesTest : TestBase() {
     fun testDataIsCopiedForRunBlocking() = runTest {
         val root = MyMutableElement(ArrayList())
         val originalData = root.mutableData
-        runBlocking(root) {
+        withContext(root) {
             assertNotSame(originalData, threadLocalData.get())
         }
     }