Review branch for #4208

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]
 }

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,10 +207,124 @@ private class LazyStandaloneCoroutine(
 }
 
 // Used by withContext when context changes, but dispatcher stays the same
-internal expect class UndispatchedCoroutine<in T>(
+internal class UndispatchedCoroutine<in T>(
     context: CoroutineContext,
     uCont: Continuation<T>
-) : ScopeCoroutine<T>
+) : ScopeCoroutine<T>(if (context[UndispatchedMarker] == null) context + UndispatchedMarker else context, uCont) {
+
+    /**
+     * The state of [ThreadContextElement]s associated with the current undispatched coroutine.
+     * It is stored in a thread local because this coroutine can be used concurrently in suspend-resume race scenario.
+     * See the followin, boiled down example with inlined `withContinuationContext` body:
+     * ```
+     * val state = saveThreadContext(ctx)
+     * try {
+     *     invokeSmthWithThisCoroutineAsCompletion() // Completion implies that 'afterResume' will be called
+     *     // COROUTINE_SUSPENDED is returned
+     * } finally {
+     *     thisCoroutine().clearThreadContext() // Concurrently the "smth" could've been already resumed on a different thread
+     *     // and it also calls saveThreadContext and clearThreadContext
+     * }
+     * ```
+     *
+     * Usage note:
+     *
+     * This part of the code is performance-sensitive.
+     * It is a well-established pattern to wrap various activities into system-specific undispatched
+     * `withContext` for the sake of logging, MDC, tracing etc., meaning that there exists thousands of
+     * undispatched coroutines.
+     * Each access to Java's [ThreadLocal] leaves a footprint in the corresponding Thread's `ThreadLocalMap`
+     * that is cleared automatically as soon as the associated thread-local (-> UndispatchedCoroutine) is garbage collected
+     * when either the corresponding thread is GC'ed or it cleans up its stale entries on other TL accesses.
+     * When such coroutines are promoted to old generation, `ThreadLocalMap`s become bloated and an arbitrary accesses to thread locals
+     * start to consume significant amount of CPU because these maps are open-addressed and cleaned up incrementally on each access.
+     * (You can read more about this effect as "GC nepotism").
+     *
+     * To avoid that, we attempt to narrow down the lifetime of this thread local as much as possible:
+     * - 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 = commonThreadLocal<Pair<CoroutineContext, Any?>?>(Symbol("UndispatchedCoroutine"))
+
+    /*
+     * Indicates that a coroutine has at least one thread context element associated with it
+     * and that 'threadStateToRecover' is going to be set in case of dispatchhing in order to preserve them.
+     * Better than nullable thread-local for easier debugging.
+     *
+     * It is used as a performance optimization to avoid 'threadStateToRecover' initialization
+     * (note: tl.get() initializes thread local),
+     * and is prone to false-positives as it is never reset: otherwise
+     * it may lead to logical data races between suspensions point where
+     * coroutine is yet being suspended in one thread while already being resumed
+     * in another.
+     */
+    @Volatile
+    private var threadLocalIsSet = false
+
+    init {
+        /*
+         * This is a hack for a very specific case in #2930 unless #3253 is implemented.
+         * 'ThreadLocalStressTest' covers this change properly.
+         *
+         * The scenario this change covers is the following:
+         * 1) The coroutine is being started as plain non kotlinx.coroutines related suspend function,
+         *    e.g. `suspend fun main` or, more importantly, Ktor `SuspendFunGun`, that is invoking
+         *    `withContext(tlElement)` which creates `UndispatchedCoroutine`.
+         * 2) It (original continuation) is then not wrapped into `DispatchedContinuation` via `intercept()`
+         *    and goes neither through `DC.run` nor through `resumeUndispatchedWith` that both
+         *    do thread context element tracking.
+         * 3) So thread locals never got chance to get properly set up via `saveThreadContext`,
+         *    but when `withContext` finishes, it attempts to recover thread locals in its `afterResume`.
+         *
+         * Here we detect precisely this situation and properly setup context to recover later.
+         *
+         */
+        if (uCont.context[ContinuationInterceptor] !is CoroutineDispatcher) {
+            /*
+             * We cannot just "read" the elements as there is no such API,
+             * so we update-restore it immediately and use the intermediate value
+             * as the initial state, leveraging the fact that thread context element
+             * is idempotent and such situations are increasingly rare.
+             */
+            val values = updateThreadContext(context, null)
+            restoreThreadContext(context, values)
+            saveThreadContext(context, values)
+        }
+    }
+
+    fun saveThreadContext(context: CoroutineContext, oldValue: Any?) {
+        threadLocalIsSet = true // Specify that thread-local is touched at all
+        threadStateToRecover.set(context to oldValue)
+    }
+
+    fun clearThreadContext(): Boolean {
+        return !(threadLocalIsSet && threadStateToRecover.get() == null).also {
+            threadStateToRecover.remove()
+        }
+    }
+
+    override fun afterCompletionUndispatched() {
+        clearThreadLocal()
+    }
+
+    override fun afterResume(state: Any?) {
+        clearThreadLocal()
+        // resume undispatched -- update context but stay on the same dispatcher
+        val result = recoverResult(state, uCont)
+        withContinuationContext(uCont, null) {
+            uCont.resumeWith(result)
+        }
+    }
+
+    private fun clearThreadLocal() {
+        if (threadLocalIsSet) {
+            threadStateToRecover.get()?.let { (ctx, value) ->
+                restoreThreadContext(ctx, value)
+            }
+            threadStateToRecover.remove()
+        }
+    }
+}
 
 private const val UNDECIDED = 0
 private const val SUSPENDED = 1

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

@@ -1,27 +1,85 @@
 package kotlinx.coroutines
 
+import kotlinx.coroutines.internal.*
 import kotlin.coroutines.*
 
+@PublishedApi // to have unmangled name when using from other modules via suppress
+@Suppress("PropertyName")
+internal expect val DefaultDelay: Delay
+
+internal expect fun Continuation<*>.toDebugString(): String
+internal expect val CoroutineContext.coroutineName: String?
+internal expect fun wrapContextWithDebug(context: CoroutineContext): CoroutineContext
+
 /**
- * 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.
+ * Executes a block using a given coroutine context.
  */
-public expect fun CoroutineScope.newCoroutineContext(context: CoroutineContext): CoroutineContext
+internal inline fun <T> withCoroutineContext(context: CoroutineContext, countOrElement: Any?, block: () -> T): T {
+    val oldValue = updateThreadContext(context, countOrElement)
+    try {
+        return block()
+    } finally {
+        restoreThreadContext(context, oldValue)
+    }
+}
 
 /**
- * Creates a context for coroutine builder functions that do not launch a new coroutine, e.g. [withContext].
- * @suppress
+ * Executes a block using a context of a given continuation.
  */
-@InternalCoroutinesApi
-public expect fun CoroutineContext.newCoroutineContext(addedContext: CoroutineContext): CoroutineContext
+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) {
+        // 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)
+        }
+    }
+}
 
-@PublishedApi // to have unmangled name when using from other modules via suppress
-@Suppress("PropertyName")
-internal expect val DefaultDelay: Delay
+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.
+     *
+     * 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 potentiallyHasUndispatchedCoroutine = context[UndispatchedMarker] !== null
+    if (!potentiallyHasUndispatchedCoroutine) return null
+    val completion = undispatchedCompletion()
+    completion?.saveThreadContext(context, oldValue)
+    return completion
+}
 
-// 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?
+private 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 necessary.
+ */
+internal object UndispatchedMarker: CoroutineContext.Element, CoroutineContext.Key<UndispatchedMarker> {
+    override val key: CoroutineContext.Key<*>
+        get() = this
+}

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

@@ -0,0 +1,88 @@
+// 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
+ */
+@ExperimentalCoroutinesApi
+public fun CoroutineScope.newCoroutineContext(context: CoroutineContext): CoroutineContext {
+    val combined = foldCopies(coroutineContext, context, true)
+    val debug = wrapContextWithDebug(combined)
+    return if (combined !== Dispatchers.Default && combined[ContinuationInterceptor] == null)
+        debug + Dispatchers.Default else debug
+}
+
+/**
+ * Creates a context for coroutine builder functions that do not launch a new coroutine, e.g. [withContext].
+ * @suppress
+ */
+@InternalCoroutinesApi
+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.
+     */
+    if (!addedContext.hasCopyableElements()) return this + addedContext
+    return foldCopies(this, addedContext, false)
+}
+
+private fun CoroutineContext.hasCopyableElements(): Boolean =
+    fold(false) { result, it -> result || it is CopyableThreadContextElement<*> }
+
+/**
+ * Folds two contexts properly applying [CopyableThreadContextElement] rules when necessary.
+ * The rules are the following:
+ * - If neither context has CTCE, the sum of two contexts is returned
+ * - Every CTCE from the left-hand side context that does not have a matching (by key) element from right-hand side context
+ *   is [copied][CopyableThreadContextElement.copyForChild] if [isNewCoroutine] is `true`.
+ * - Every CTCE from the left-hand side context that has a matching element in the right-hand side context is [merged][CopyableThreadContextElement.mergeForChild]
+ * - Every CTCE from the 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.hasCopyableElements()
+    val hasElementsRight = appendContext.hasCopyableElements()
+
+    // Nothing to fold, so just return the sum of contexts
+    if (!hasElementsLeft && !hasElementsRight) {
+        return originalContext + appendContext
+    }
+
+    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
+}