PR #7 Move ThreadContextElement to common

knisht · knisht · commit 09a2663561d5 · 2025-02-06T12:41:38.000+01:00
diff --git a/IntelliJ-patches.md b/IntelliJ-patches.md
@@ -90,3 +90,8 @@ Why not just let `next` return a maybe wrapped value? That's because it is heavi
 Changes were made to lambda parameter `onElementRetrieved` in `BufferedChannel<E>` methods: now they accept `Any?` instead of `E` because now they may be given a wrapped value.
 
 `SelectImplementation.complete` now uses `debuggerCapture` to properly propagate value that might come from flows. 
+
+## `ThreadContextElement` is moved to common
+
+The motivation of this change is enabling `ThreadContextElements` in Kotlin/Native, Kotlin/JS and Kotlin/wasm.
+The API is left intact.
diff --git a/kotlinx-coroutines-core/common/src/ThreadContextElement.common.kt b/kotlinx-coroutines-core/common/src/ThreadContextElement.common.kt
@@ -0,0 +1,82 @@
+package kotlinx.coroutines
+
+import kotlin.coroutines.*
+
+/**
+ * Defines elements in [CoroutineContext] that are installed into thread context
+ * every time the coroutine with this element in the context is resumed on a thread.
+ *
+ * Implementations of this interface define a type [S] of the thread-local state that they need to store on
+ * resume of a coroutine and restore later on suspend. The infrastructure provides the corresponding storage.
+ *
+ * Example usage looks like this:
+ *
+ * ```
+ * // Appends "name" of a coroutine to a current thread name when coroutine is executed
+ * class CoroutineName(val name: String) : ThreadContextElement<String> {
+ *     // declare companion object for a key of this element in coroutine context
+ *     companion object Key : CoroutineContext.Key<CoroutineName>
+ *
+ *     // provide the key of the corresponding context element
+ *     override val key: CoroutineContext.Key<CoroutineName>
+ *         get() = Key
+ *
+ *     // this is invoked before coroutine is resumed on current thread
+ *     override fun updateThreadContext(context: CoroutineContext): String {
+ *         val previousName = Thread.currentThread().name
+ *         Thread.currentThread().name = "$previousName # $name"
+ *         return previousName
+ *     }
+ *
+ *     // this is invoked after coroutine has suspended on current thread
+ *     override fun restoreThreadContext(context: CoroutineContext, oldState: String) {
+ *         Thread.currentThread().name = oldState
+ *     }
+ * }
+ *
+ * // Usage
+ * launch(Dispatchers.Main + CoroutineName("Progress bar coroutine")) { ... }
+ * ```
+ *
+ * Every time this coroutine is resumed on a thread, UI thread name is updated to
+ * "UI thread original name # Progress bar coroutine" and the thread name is restored to the original one when
+ * this coroutine suspends.
+ *
+ * To use [ThreadLocal] variable within the coroutine use [ThreadLocal.asContextElement][asContextElement] function.
+ *
+ * ### Reentrancy and thread-safety
+ *
+ * Correct implementations of this interface must expect that calls to [restoreThreadContext]
+ * may happen in parallel to the subsequent [updateThreadContext] and [restoreThreadContext] operations.
+ * See [CopyableThreadContextElement] for advanced interleaving details.
+ *
+ * All implementations of [ThreadContextElement] should be thread-safe and guard their internal mutable state
+ * within an element accordingly.
+ */
+public interface ThreadContextElement<S> : CoroutineContext.Element {
+    /**
+     * Updates context of the current thread.
+     * This function is invoked before the coroutine in the specified [context] is resumed in the current thread
+     * when the context of the coroutine this element.
+     * The result of this function is the old value of the thread-local state that will be passed to [restoreThreadContext].
+     * This method should handle its own exceptions and do not rethrow it. Thrown exceptions will leave coroutine which
+     * context is updated in an undefined state and may crash an application.
+     *
+     * @param context the coroutine context.
+     */
+    public fun updateThreadContext(context: CoroutineContext): S
+
+    /**
+     * Restores context of the current thread.
+     * This function is invoked after the coroutine in the specified [context] is suspended in the current thread
+     * if [updateThreadContext] was previously invoked on resume of this coroutine.
+     * The value of [oldState] is the result of the previous invocation of [updateThreadContext] and it should
+     * be restored in the thread-local state by this function.
+     * This method should handle its own exceptions and do not rethrow it. Thrown exceptions will leave coroutine which
+     * context is updated in an undefined state and may crash an application.
+     *
+     * @param context the coroutine context.
+     * @param oldState the value returned by the previous invocation of [updateThreadContext].
+     */
+    public fun restoreThreadContext(context: CoroutineContext, oldState: S)
+}
diff --git a/kotlinx-coroutines-core/common/src/internal/ThreadContext.common.kt b/kotlinx-coroutines-core/common/src/internal/ThreadContext.common.kt
@@ -1,5 +1,56 @@
 package kotlinx.coroutines.internal
 
+import kotlinx.coroutines.*
 import kotlin.coroutines.*
+import kotlin.jvm.*
 
-internal expect fun threadContextElements(context: CoroutineContext): Any
+@JvmField
+internal val NO_THREAD_ELEMENTS = Symbol("NO_THREAD_ELEMENTS")
+
+// Used when there are >= 2 active elements in the context
+@Suppress("UNCHECKED_CAST")
+internal class ThreadState(@JvmField val context: CoroutineContext, n: Int) {
+    private val values = arrayOfNulls<Any>(n)
+    private val elements = arrayOfNulls<ThreadContextElement<Any?>>(n)
+    private var i = 0
+
+    fun append(element: ThreadContextElement<*>, value: Any?) {
+        values[i] = value
+        elements[i++] = element as ThreadContextElement<Any?>
+    }
+
+    fun restore(context: CoroutineContext) {
+        for (i in elements.indices.reversed()) {
+            elements[i]!!.restoreThreadContext(context, values[i])
+        }
+    }
+}
+
+// Counts ThreadContextElements in the context
+// Any? here is Int | ThreadContextElement (when count is one)
+private val countAll =
+    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
+        }
+        return countOrElement
+    }
+
+// Find one (first) ThreadContextElement in the context, it is used when we know there is exactly one
+internal val findOne =
+    fun (found: ThreadContextElement<*>?, element: CoroutineContext.Element): ThreadContextElement<*>? {
+        if (found != null) return found
+        return element as? ThreadContextElement<*>
+    }
+
+// Updates state for ThreadContextElements in the context using the given ThreadState
+internal val updateState =
+    fun (state: ThreadState, element: CoroutineContext.Element): ThreadState {
+        if (element is ThreadContextElement<*>) {
+            state.append(element, element.updateThreadContext(state.context))
+        }
+        return state
+    }
+
+internal fun threadContextElements(context: CoroutineContext): Any = context.fold(0, countAll)!!
diff --git a/kotlinx-coroutines-core/jsAndWasmShared/src/internal/ThreadContext.kt b/kotlinx-coroutines-core/jsAndWasmShared/src/internal/ThreadContext.kt
@@ -1,5 +1,41 @@
 package kotlinx.coroutines.internal
 
+import kotlinx.coroutines.*
 import kotlin.coroutines.*
 
-internal actual fun threadContextElements(context: CoroutineContext): Any = 0
+// 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
+        countOrElement is Int -> {
+            // slow path for multiple active ThreadContextElements, allocates ThreadState for multiple old values
+            context.fold(ThreadState(context, countOrElement), updateState)
+        }
+        else -> {
+            // fast path for one ThreadContextElement (no allocations, no additional context scan)
+            @Suppress("UNCHECKED_CAST")
+            val element = countOrElement as ThreadContextElement<Any?>
+            element.updateThreadContext(context)
+        }
+    }
+}
+
+internal fun restoreThreadContext(context: CoroutineContext, oldState: Any?) {
+    when {
+        oldState === NO_THREAD_ELEMENTS -> return // very fast path when there are no ThreadContextElements
+        oldState is ThreadState -> {
+            // slow path with multiple stored ThreadContextElements
+            oldState.restore(context)
+        }
+        else -> {
+            // fast path for one ThreadContextElement, but need to find it
+            @Suppress("UNCHECKED_CAST")
+            val element = context.fold(null, findOne) as ThreadContextElement<Any?>
+            element.restoreThreadContext(context, oldState)
+        }
+    }
+}
diff --git a/kotlinx-coroutines-core/jvm/src/ThreadContextElement.kt b/kotlinx-coroutines-core/jvm/src/ThreadContextElement.kt
@@ -3,85 +3,6 @@ package kotlinx.coroutines
 import kotlinx.coroutines.internal.*
 import kotlin.coroutines.*
 
-/**
- * Defines elements in [CoroutineContext] that are installed into thread context
- * every time the coroutine with this element in the context is resumed on a thread.
- *
- * Implementations of this interface define a type [S] of the thread-local state that they need to store on
- * resume of a coroutine and restore later on suspend. The infrastructure provides the corresponding storage.
- *
- * Example usage looks like this:
- *
- * ```
- * // Appends "name" of a coroutine to a current thread name when coroutine is executed
- * class CoroutineName(val name: String) : ThreadContextElement<String> {
- *     // declare companion object for a key of this element in coroutine context
- *     companion object Key : CoroutineContext.Key<CoroutineName>
- *
- *     // provide the key of the corresponding context element
- *     override val key: CoroutineContext.Key<CoroutineName>
- *         get() = Key
- *
- *     // this is invoked before coroutine is resumed on current thread
- *     override fun updateThreadContext(context: CoroutineContext): String {
- *         val previousName = Thread.currentThread().name
- *         Thread.currentThread().name = "$previousName # $name"
- *         return previousName
- *     }
- *
- *     // this is invoked after coroutine has suspended on current thread
- *     override fun restoreThreadContext(context: CoroutineContext, oldState: String) {
- *         Thread.currentThread().name = oldState
- *     }
- * }
- *
- * // Usage
- * launch(Dispatchers.Main + CoroutineName("Progress bar coroutine")) { ... }
- * ```
- *
- * Every time this coroutine is resumed on a thread, UI thread name is updated to
- * "UI thread original name # Progress bar coroutine" and the thread name is restored to the original one when
- * this coroutine suspends.
- *
- * To use [ThreadLocal] variable within the coroutine use [ThreadLocal.asContextElement][asContextElement] function.
- *
- * ### Reentrancy and thread-safety
- *
- * Correct implementations of this interface must expect that calls to [restoreThreadContext]
- * may happen in parallel to the subsequent [updateThreadContext] and [restoreThreadContext] operations.
- * See [CopyableThreadContextElement] for advanced interleaving details.
- *
- * All implementations of [ThreadContextElement] should be thread-safe and guard their internal mutable state
- * within an element accordingly.
- */
-public interface ThreadContextElement<S> : CoroutineContext.Element {
-    /**
-     * Updates context of the current thread.
-     * This function is invoked before the coroutine in the specified [context] is resumed in the current thread
-     * when the context of the coroutine this element.
-     * The result of this function is the old value of the thread-local state that will be passed to [restoreThreadContext].
-     * This method should handle its own exceptions and do not rethrow it. Thrown exceptions will leave coroutine which
-     * context is updated in an undefined state and may crash an application.
-     *
-     * @param context the coroutine context.
-     */
-    public fun updateThreadContext(context: CoroutineContext): S
-
-    /**
-     * Restores context of the current thread.
-     * This function is invoked after the coroutine in the specified [context] is suspended in the current thread
-     * if [updateThreadContext] was previously invoked on resume of this coroutine.
-     * The value of [oldState] is the result of the previous invocation of [updateThreadContext] and it should
-     * be restored in the thread-local state by this function.
-     * This method should handle its own exceptions and do not rethrow it. Thrown exceptions will leave coroutine which
-     * context is updated in an undefined state and may crash an application.
-     *
-     * @param context the coroutine context.
-     * @param oldState the value returned by the previous invocation of [updateThreadContext].
-     */
-    public fun restoreThreadContext(context: CoroutineContext, oldState: S)
-}
-
 /**
  * A [ThreadContextElement] copied whenever a child coroutine inherits a context containing it.
  *
diff --git a/kotlinx-coroutines-core/jvm/src/internal/ThreadContext.kt b/kotlinx-coroutines-core/jvm/src/internal/ThreadContext.kt
@@ -3,57 +3,6 @@ package kotlinx.coroutines.internal
 import kotlinx.coroutines.*
 import kotlin.coroutines.*
 
-@JvmField
-internal val NO_THREAD_ELEMENTS = Symbol("NO_THREAD_ELEMENTS")
-
-// Used when there are >= 2 active elements in the context
-@Suppress("UNCHECKED_CAST")
-private class ThreadState(@JvmField val context: CoroutineContext, n: Int) {
-    private val values = arrayOfNulls<Any>(n)
-    private val elements = arrayOfNulls<ThreadContextElement<Any?>>(n)
-    private var i = 0
-
-    fun append(element: ThreadContextElement<*>, value: Any?) {
-        values[i] = value
-        elements[i++] = element as ThreadContextElement<Any?>
-    }
-
-    fun restore(context: CoroutineContext) {
-        for (i in elements.indices.reversed()) {
-            elements[i]!!.restoreThreadContext(context, values[i])
-        }
-    }
-}
-
-// Counts ThreadContextElements in the context
-// Any? here is Int | ThreadContextElement (when count is one)
-private val countAll =
-    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
-        }
-        return countOrElement
-    }
-
-// Find one (first) ThreadContextElement in the context, it is used when we know there is exactly one
-private val findOne =
-    fun (found: ThreadContextElement<*>?, element: CoroutineContext.Element): ThreadContextElement<*>? {
-        if (found != null) return found
-        return element as? ThreadContextElement<*>
-    }
-
-// Updates state for ThreadContextElements in the context using the given ThreadState
-private val updateState =
-    fun (state: ThreadState, element: CoroutineContext.Element): ThreadState {
-        if (element is ThreadContextElement<*>) {
-            state.append(element, element.updateThreadContext(state.context))
-        }
-        return state
-    }
-
-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? {
diff --git a/kotlinx-coroutines-core/native/src/CoroutineContext.kt b/kotlinx-coroutines-core/native/src/CoroutineContext.kt
diff --git a/kotlinx-coroutines-core/native/src/internal/ThreadContext.kt b/kotlinx-coroutines-core/native/src/internal/ThreadContext.kt
diff --git a/kotlinx-coroutines-core/wasmJs/src/CoroutineContext.kt b/kotlinx-coroutines-core/wasmJs/src/CoroutineContext.kt
