Allow introspection into coroutine context from non-suspending environment

This is required for https://youtrack.jetbrains.com/issue/IJPL-445

Author: knisht

IntelliJ-patches.md

@@ -1,2 +1,21 @@
 # Included IntelliJ-related patches
-- TODO
+
+## Introspection of `coroutineContext`
+
+### Description: 
+To support the threading framework of IntelliJ, 
+we need to cooperate with kotlin coroutines better than they allow it by default.
+
+One example of this cooperation is an ability to get coroutine context from non-suspending places.
+Essentially, we put the coroutine context into a thread local variable on every coroutine resumption, 
+which allows us to read necessary information without a significant change in semantics. 
+This change has a mild performance penalty, namely, modification of a thread local variable. 
+However, coroutines themselves use thread local states via `ThreadLocalContextElement`, which hints that 
+one more thread local variable would not harm.
+
+### API:
+
+We provide a single method `kotlinx.coroutines.internal.intellij.IntellijCoroutines.currentThreadCoroutineContext`.
+The invariant is that the result of this method is always equal to `coroutineContext` in suspending environment, 
+and it does not change during the non-suspending execution within the same thread.
+

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

@@ -25,3 +25,6 @@ internal expect inline fun <T> withCoroutineContext(context: CoroutineContext, c
 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?
+
+// added by IntelliJ
+internal expect inline fun <T> withThreadLocalContext(context: CoroutineContext, block: () -> T) : T

kotlinx-coroutines-core/common/src/internal/Scopes.kt

@@ -25,7 +25,9 @@ internal open class ScopeCoroutine<in T>(
 
     override fun afterResume(state: Any?) {
         // Resume direct because scope is already in the correct context
-        uCont.resumeWith(recoverResult(state, uCont))
+        withThreadLocalContext(uCont.context) {
+            uCont.resumeWith(recoverResult(state, uCont))
+        }
     }
 }
 

kotlinx-coroutines-core/common/src/intrinsics/Undispatched.kt

@@ -58,7 +58,9 @@ private inline fun <T> startDirect(completion: Continuation<T>, block: (Continua
  */
 internal fun <T, R> ScopeCoroutine<T>.startUndispatchedOrReturn(receiver: R, block: suspend R.() -> T): Any? {
     return undispatchedResult({ true }) {
-        block.startCoroutineUninterceptedOrReturn(receiver, this)
+        withThreadLocalContext(this.context) {
+            block.startCoroutineUninterceptedOrReturn(receiver, this)
+        }
     }
 }
 

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

@@ -55,3 +55,5 @@ internal actual class UndispatchedCoroutine<in T> actual constructor(
 ) : ScopeCoroutine<T>(context, uCont) {
     override fun afterResume(state: Any?) = uCont.resumeWith(recoverResult(state, uCont))
 }
+
+internal actual inline fun <T> withThreadLocalContext(context: CoroutineContext, block: () -> T) : T = block()

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

@@ -1,6 +1,7 @@
 package kotlinx.coroutines
 
 import kotlinx.coroutines.internal.*
+import kotlinx.coroutines.internal.intellij.currentContextThreadLocal
 import kotlin.coroutines.*
 import kotlin.coroutines.jvm.internal.CoroutineStackFrame
 
@@ -90,12 +91,22 @@ private fun foldCopies(originalContext: CoroutineContext, appendContext: Corouti
 internal actual inline fun <T> withCoroutineContext(context: CoroutineContext, countOrElement: Any?, block: () -> T): T {
     val oldValue = updateThreadContext(context, countOrElement)
     try {
-        return block()
+        return withThreadLocalContext(context, block)
     } finally {
         restoreThreadContext(context, oldValue)
     }
 }
 
+internal actual inline fun <T> withThreadLocalContext(context: CoroutineContext, block: () -> T) : T {
+    val old = currentContextThreadLocal.get()
+    currentContextThreadLocal.set(context)
+    try {
+        return block()
+    } finally {
+        currentContextThreadLocal.set(old)
+    }
+}
+
 /**
  * Executes a block using a context of a given continuation.
  */
@@ -109,7 +120,7 @@ internal actual inline fun <T> withContinuationContext(continuation: Continuatio
         null // fast path -- don't even try to find undispatchedCompletion as there's nothing to restore in the context
     }
     try {
-        return block()
+        return withThreadLocalContext(context, block)
     } finally {
         if (undispatchedCompletion == null || undispatchedCompletion.clearThreadContext()) {
             restoreThreadContext(context, oldValue)

kotlinx-coroutines-core/jvm/src/flow/internal/SafeCollector.kt

@@ -57,7 +57,9 @@ internal actual class SafeCollector<T> actual constructor(
 
     override fun invokeSuspend(result: Result<Any?>): Any {
         result.onFailure { lastEmissionContext = DownstreamExceptionContext(it, context) }
-        completion_?.resumeWith(result as Result<Unit>)
+        withThreadLocalContext(context) {
+            completion_?.resumeWith(result as Result<Unit>)
+        }
         return COROUTINE_SUSPENDED
     }
 

kotlinx-coroutines-core/jvm/src/internal/intellij/intellij.kt

@@ -1,6 +1,12 @@
+/**
+ * A special file that contains IntelliJ-related functions
+ */
 package kotlinx.coroutines.internal.intellij
 
 import kotlinx.coroutines.InternalCoroutinesApi
+import kotlin.coroutines.*
+
+internal val currentContextThreadLocal : ThreadLocal<CoroutineContext?> = ThreadLocal.withInitial { null }
 
 /**
  * [IntellijCoroutines] exposes the API added as part of IntelliJ patches.
@@ -9,4 +15,14 @@ import kotlinx.coroutines.InternalCoroutinesApi
 @InternalCoroutinesApi
 public object IntellijCoroutines {
 
+    /**
+     * IntelliJ Platform would like to introspect coroutine contexts outside the coroutine framework.
+     * This function is a non-suspend version of [coroutineContext].
+     *
+     * @return null if current thread is not used by coroutine dispatchers,
+     * or [coroutineContext] otherwise.
+     */
+    public fun currentThreadCoroutineContext(): CoroutineContext? {
+        return currentContextThreadLocal.get()
+    }
 }

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

@@ -0,0 +1,91 @@
+package kotlinx.coroutines
+
+import kotlinx.coroutines.flow.*
+import kotlinx.coroutines.internal.intellij.*
+import kotlinx.coroutines.testing.*
+import org.junit.Assert.*
+import org.junit.Test
+import kotlin.coroutines.*
+
+class ExposedThreadContextTest : TestBase() {
+    @Test
+    fun runBlocking() = runBlocking {
+        assertContextEqualUnderResumption()
+    }
+
+    class C : AbstractCoroutineContextElement(C) {
+        companion object Key : CoroutineContext.Key<C>
+    }
+
+    @Test
+    fun withContext() = runBlocking {
+        val element = C()
+        withContext(element) {
+            // the context changed
+            assertContextEqualUnderResumption()
+            withContext(element) {
+                // checking fast path -- the context effectively does not change
+                assertContextEqualUnderResumption()
+            }
+        }
+    }
+
+    @Test
+    fun launch() = runBlocking {
+        for (i in 0..10) {
+            launch {
+                assertContextEqualUnderResumption()
+            }
+        }
+    }
+
+    @Test
+    fun coroutineScope() = runBlocking {
+        coroutineScope {
+            assertContextEqualUnderResumption()
+            coroutineScope {
+                assertContextEqualUnderResumption()
+            }
+        }
+    }
+
+    @Test
+    fun supervisorScope() = runBlocking {
+        supervisorScope {
+            assertContextEqualUnderResumption()
+            supervisorScope {
+                assertContextEqualUnderResumption()
+            }
+        }
+    }
+
+    @Test
+    fun testFlow() = runBlocking {
+        coroutineScope {
+            val flowVar = flow {
+                    repeat(10) {
+                        // Flow has encapsulated context
+                        assertContextEqualUnderResumption()
+                        emit(it)
+                    }
+                }.flowOn(C())
+            flowVar.collect {
+                assertContextEqualUnderResumption()
+            }
+        }
+    }
+
+    private suspend fun assertContextEqualUnderResumption() {
+        // thread context should survive dispatches
+        assertContextsEqual()
+        yield()
+        assertContextsEqual()
+    }
+
+
+    private suspend fun assertContextsEqual() {
+        val coroutineContext = currentCoroutineContext()
+        val threadContext = IntellijCoroutines.currentThreadCoroutineContext()
+        assertEquals(coroutineContext, threadContext)
+    }
+}

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

@@ -51,3 +51,5 @@ internal actual class UndispatchedCoroutine<in T> actual constructor(
 ) : ScopeCoroutine<T>(context, uCont) {
     override fun afterResume(state: Any?) = uCont.resumeWith(recoverResult(state, uCont))
 }
+
+internal actual inline fun <T> withThreadLocalContext(context: CoroutineContext, block: () -> T) : T = block()

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

@@ -45,3 +45,5 @@ internal actual class UndispatchedCoroutine<in T> actual constructor(
 ) : ScopeCoroutine<T>(context, uCont) {
     override fun afterResume(state: Any?) = uCont.resumeWith(recoverResult(state, uCont))
 }
+
+internal actual inline fun <T> withThreadLocalContext(context: CoroutineContext, block: () -> T) : T = block()