CopyableThreadContextElement implementation

kotlinx-coroutines-core/api/kotlinx-coroutines-core.api

@@ -141,7 +141,8 @@ public final class kotlinx/coroutines/CompletionHandlerException : java/lang/Run
 }
 
 public abstract interface class kotlinx/coroutines/CopyableThreadContextElement : kotlinx/coroutines/ThreadContextElement {
-	public abstract fun copyForChildCoroutine ()Lkotlinx/coroutines/CopyableThreadContextElement;
+	public abstract fun copyForChild ()Lkotlinx/coroutines/CopyableThreadContextElement;
+	public abstract fun mergeForChild (Lkotlin/coroutines/CoroutineContext$Element;)Lkotlin/coroutines/CoroutineContext;
 }
 
 public final class kotlinx/coroutines/CopyableThreadContextElement$DefaultImpls {
@@ -156,6 +157,7 @@ public abstract interface class kotlinx/coroutines/CopyableThrowable {
 }
 
 public final class kotlinx/coroutines/CoroutineContextKt {
+	public static final fun newCoroutineContext (Lkotlin/coroutines/CoroutineContext;Lkotlin/coroutines/CoroutineContext;)Lkotlin/coroutines/CoroutineContext;
 	public static final fun newCoroutineContext (Lkotlinx/coroutines/CoroutineScope;Lkotlin/coroutines/CoroutineContext;)Lkotlin/coroutines/CoroutineContext;
 }
 

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

@@ -148,7 +148,8 @@ public suspend fun <T> withContext(
     return suspendCoroutineUninterceptedOrReturn sc@ { uCont ->
         // compute new context
         val oldContext = uCont.context
-        val newContext = oldContext + context
+        // Copy CopyableThreadContextElement if necessary
+        val newContext = oldContext.newCoroutineContext(context)
         // always check for cancellation of new context
         newContext.ensureActive()
         // FAST PATH #1 -- new context is the same as the old one

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

@@ -8,10 +8,18 @@ import kotlin.coroutines.*
 
 /**
  * Creates a context for the new coroutine. It installs [Dispatchers.Default] when no other dispatcher or
- * [ContinuationInterceptor] is specified, and adds optional support for debugging facilities (when turned on).
+ * [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 context for coroutine builder functions that do not launch a new coroutine, namely [withContext].
+ * @suppress
+ */
+@InternalCoroutinesApi
+public expect fun CoroutineContext.newCoroutineContext(addedContext: CoroutineContext): CoroutineContext
+
 @PublishedApi
 @Suppress("PropertyName")
 internal expect val DefaultDelay: Delay

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

@@ -42,6 +42,10 @@ public actual fun CoroutineScope.newCoroutineContext(context: CoroutineContext):
         combined + Dispatchers.Default else combined
 }
 
+public actual fun CoroutineContext.newCoroutineContext(addedContext: CoroutineContext): CoroutineContext {
+    return this + addedContext
+}
+
 // 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()

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

@@ -10,35 +10,84 @@ import kotlin.coroutines.jvm.internal.CoroutineStackFrame
 
 /**
  * Creates context for the new coroutine. It installs [Dispatchers.Default] when no other dispatcher nor
- * [ContinuationInterceptor] is specified, and adds optional support for debugging facilities (when turned on).
+ * [ContinuationInterceptor] is specified, and adds optional support for
+ * copyable thread context [elements][CopyableThreadContextElement] and debugging facilities (when turned on).
  *
  * See [DEBUG_PROPERTY_NAME] for description of debugging facilities on JVM.
  */
 @ExperimentalCoroutinesApi
 public actual fun CoroutineScope.newCoroutineContext(context: CoroutineContext): CoroutineContext {
-    val combined = coroutineContext.foldCopiesForChildCoroutine() + context
+    val combined = foldCopies(coroutineContext, context, true)
     val debug = if (DEBUG) combined + CoroutineId(COROUTINE_ID.incrementAndGet()) else combined
     return if (combined !== Dispatchers.Default && combined[ContinuationInterceptor] == null)
         debug + Dispatchers.Default else debug
 }
 
 /**
- * Returns the [CoroutineContext] for a child coroutine to inherit.
- *
- * If any [CopyableThreadContextElement] is in the [this], calls
- * [CopyableThreadContextElement.copyForChildCoroutine] on each, returning a new [CoroutineContext]
- * by folding the returned copied elements into [this].
- *
- * Returns [this] if `this` has zero [CopyableThreadContextElement] in it.
+ * Creates context for coroutine builder functions that do not launch a new coroutine, namely [withContext].
+ * @suppress
  */
-private fun CoroutineContext.foldCopiesForChildCoroutine(): CoroutineContext {
-    val hasToCopy = fold(false) { result, it ->
-        result || it is CopyableThreadContextElement<*>
+@InternalCoroutinesApi
+public actual 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 element.
+     */
+    if (!addedContext.fold(false, hasCopyableElements)) return this + addedContext
+    return foldCopies(this, addedContext, false)
+}
+
+private val hasCopyableElements: (Boolean, CoroutineContext.Element) -> Boolean = { result, it ->
+    result || it is CopyableThreadContextElement<*>
+}
+
+/**
+ * Folds two contexts properly applying [CopyableThreadContextElement] rules when necessary.
+ * The rules are the following:
+ * * If both context do not have CTCE, the sum of two contexts is returned
+ * * Every CTCE from left-hand side context that does not have matching (by key) element from right-hand side context
+ *   is [copied][CopyableThreadContextElement.copyForChild]
+ * * Every CTCE from left-hand side context that has matching element in right-hand side context is [merged][CopyableThreadContextElement.mergeForChild]
+ * * Every CTCE from right-hand side context that hasn't been merged is copied
+ * * Everything else is added to the resulting context as is.
+ */
+private fun foldCopies(originalContext: CoroutineContext, appendContext: CoroutineContext, isNewCoroutine: Boolean): CoroutineContext {
+    // Do we have something to copy left-hand side?
+    val hasElementsLeft = originalContext.fold(false, hasCopyableElements)
+    val hasElementsRight = appendContext.fold(false, hasCopyableElements)
+
+    // Nothing to fold, so just return the sum of contexts
+    if (!hasElementsLeft && !hasElementsRight) {
+        return originalContext + appendContext
     }
-    if (!hasToCopy) return this
-    return fold<CoroutineContext>(EmptyCoroutineContext) { combined, it ->
-        combined + if (it is CopyableThreadContextElement<*>) it.copyForChildCoroutine() else it
+
+    var leftoverContext = appendContext
+    val folded = originalContext.fold<CoroutineContext>(EmptyCoroutineContext) { result, element ->
+        if (element !is CopyableThreadContextElement<*>) return@fold result + element
+        // Will this element be overwritten?
+        val newElement = leftoverContext[element.key]
+        // No, just copy it
+        if (newElement == null) {
+            // For 'withContext'-like builders we do not copy as the element is not shared
+            return@fold result + if (isNewCoroutine) element.copyForChild() else element
+        }
+        // Yes, then first remove the element from append context
+        leftoverContext = leftoverContext.minusKey(element.key)
+        // Return the sum
+        @Suppress("UNCHECKED_CAST")
+        return@fold result + (element as CopyableThreadContextElement<Any?>).mergeForChild(newElement)
+    }
+
+    if (hasElementsRight) {
+        leftoverContext = leftoverContext.fold<CoroutineContext>(EmptyCoroutineContext) { result, element ->
+            // We're appending new context element -- we have to copy it, otherwise it may be shared with others
+            if (element is CopyableThreadContextElement<*>) {
+                return@fold result + element.copyForChild()
+            }
+            return@fold result + element
+        }
     }
+    return folded + leftoverContext
 }
 
 /**

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

@@ -97,7 +97,7 @@ public interface ThreadContextElement<S> : CoroutineContext.Element {
  * is in a coroutine:
  *
  * ```
- * class TraceContextElement(private val traceData: TraceData?) : CopyableThreadContextElement<TraceData?> {
+ * class TraceContextElement(private val traceData: TraceData?) : CopyableThreadContextElement<TraceData?, TraceContextElement> {
  *     companion object Key : CoroutineContext.Key<TraceContextElement>
  *     override val key: CoroutineContext.Key<TraceContextElement> = Key
  *
@@ -111,32 +111,52 @@ public interface ThreadContextElement<S> : CoroutineContext.Element {
  *         traceThreadLocal.set(oldState)
  *     }
  *
- *     override fun copyForChildCoroutine(): CopyableThreadContextElement<TraceData?> {
+ *     override fun copyForChild(): TraceContextElement {
  *         // Copy from the ThreadLocal source of truth at child coroutine launch time. This makes
  *         // ThreadLocal writes between resumption of the parent coroutine and the launch of the
  *         // child coroutine visible to the child.
  *         return TraceContextElement(traceThreadLocal.get()?.copy())
  *     }
+ *
+ *     override fun mergeForChild(overwritingElement: CoroutineContext.Element): CoroutineContext
+ *         // Merge operation defines how to handle situation when both
+ *         // parent coroutine has the element in the context and it was also
+ *         // explicitly passed to a child coroutine.
+ *         // If merge is not defined, the copy of the element can be returned.
+ *         return TraceContextElement(traceThreadLocal.get()?.copy())
+ *     }
  * }
  * ```
  *
  * A coroutine using this mechanism can safely call Java code that assumes it's called using a
  * `Thread`.
  */
+@DelicateCoroutinesApi
 @ExperimentalCoroutinesApi
 public interface CopyableThreadContextElement<S> : ThreadContextElement<S> {
 
     /**
      * Returns a [CopyableThreadContextElement] to replace `this` `CopyableThreadContextElement` in the child
-     * coroutine's context that is under construction.
+     * coroutine's context that is under construction if the child coroutine's context does not contain
+     * an element with the same [key].
      *
      * This function is called on the element each time a new coroutine inherits a context containing it,
      * and the returned value is folded into the context given to the child.
      *
      * Since this method is called whenever a new coroutine is launched in a context containing this
      * [CopyableThreadContextElement], implementations are performance-sensitive.
      */
-    public fun copyForChildCoroutine(): CopyableThreadContextElement<S>
+    public fun copyForChild(): CopyableThreadContextElement<S>
+
+    /**
+     * Returns a [CopyableThreadContextElement] to replace `this` `CopyableThreadContextElement` in the child
+     * coroutine's context when the child coroutine's context contains an element with the same [key] as the
+     * current one.
+     *
+     * This function is called on current element, supplied with an element retrieved from child's
+     * coroutine context by the current [key].
+     */
+    public fun mergeForChild(overwritingElement: CoroutineContext.Element): CoroutineContext
 }
 
 /**

kotlinx-coroutines-core/jvm/test/ThreadContextElementTest.kt

@@ -126,8 +126,7 @@ class ThreadContextElementTest : TestBase() {
     @Test
     fun testCopyableThreadContextElementImplementsWriteVisibility() = runTest {
         newFixedThreadPoolContext(nThreads = 4, name = "withContext").use {
-            val startData = MyData()
-            withContext(it + CopyForChildCoroutineElement(startData)) {
+            withContext(it + CopyForChildCoroutineElement(MyData())) {
                 val forBlockData = MyData()
                 myThreadLocal.setForBlock(forBlockData) {
                     assertSame(myThreadLocal.get(), forBlockData)
@@ -153,7 +152,7 @@ class ThreadContextElementTest : TestBase() {
                         assertSame(myThreadLocal.get(), forBlockData)
                     }
                 }
-                assertSame(myThreadLocal.get(), startData) // Asserts value was restored.
+                assertNull(myThreadLocal.get()) // Asserts value was restored to its origin
             }
         }
     }
@@ -187,7 +186,7 @@ class MyElement(val data: MyData) : ThreadContextElement<MyData?> {
 }
 
 /**
- * A [ThreadContextElement] that implements copy semantics in [copyForChildCoroutine].
+ * A [ThreadContextElement] that implements copy semantics in [copyForChild].
  */
 class CopyForChildCoroutineElement(val data: MyData?) : CopyableThreadContextElement<MyData?> {
     companion object Key : CoroutineContext.Key<CopyForChildCoroutineElement>
@@ -201,6 +200,10 @@ class CopyForChildCoroutineElement(val data: MyData?) : CopyableThreadContextEle
         return oldState
     }
 
+    override fun mergeForChild(overwritingElement: CoroutineContext.Element): CopyForChildCoroutineElement {
+        TODO("Not used in tests")
+    }
+
     override fun restoreThreadContext(context: CoroutineContext, oldState: MyData?) {
         myThreadLocal.set(oldState)
     }
@@ -216,7 +219,7 @@ class CopyForChildCoroutineElement(val data: MyData?) : CopyableThreadContextEle
      * will be reflected in the parent coroutine's [CopyForChildCoroutineElement] when it yields the
      * thread and calls [restoreThreadContext].
      */
-    override fun copyForChildCoroutine(): CopyableThreadContextElement<MyData?> {
+    override fun copyForChild(): CopyForChildCoroutineElement {
         return CopyForChildCoroutineElement(myThreadLocal.get())
     }
 }