SafeCollector performance improvements:

  * Cache result of the context check thus speeding up safe collector on happy path by a factor of three
  * Separate fast and slow paths in JobSupport to drastically change inlining decisions of the JVM that are crucial for leaf coroutines with flows

Author: qwwdfsad

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

@@ -783,6 +783,11 @@ public open class JobSupport constructor(active: Boolean) : Job, ChildJob, Paren
             if (!tryFinalizeSimpleState(state, proposedUpdate, mode)) return COMPLETING_RETRY
             return COMPLETING_COMPLETED
         }
+        // The separate slow-path function to simplify profiling
+        return tryMakeCompletingSlowPath(state, proposedUpdate, mode)
+    }
+
+    private fun tryMakeCompletingSlowPath(state: Incomplete, proposedUpdate: Any?, mode: Int): Int {
         // get state's list or else promote to list to correctly operate on child lists
         val list = getOrPromoteCancellingList(state) ?: return COMPLETING_RETRY
         // promote to Finishing state if we are not in it yet
@@ -1202,7 +1207,7 @@ public open class JobSupport constructor(active: Boolean) : Job, ChildJob, Paren
  * Class to represent object as the final state of the Job
  */
 private class IncompleteStateBox(@JvmField val state: Incomplete)
-private fun Any?.boxIncomplete(): Any? = if (this is Incomplete) IncompleteStateBox(this) else this
+internal fun Any?.boxIncomplete(): Any? = if (this is Incomplete) IncompleteStateBox(this) else this
 internal fun Any?.unboxState(): Any? = (this as? IncompleteStateBox)?.state ?: this
 
 // --------------- helper classes & constants for job implementation

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

@@ -16,14 +16,26 @@ internal class SafeCollector<T>(
 ) : FlowCollector<T>, SynchronizedObject() {
 
     private val collectContext = collectContext.minusKey(Job).minusId()
+    private var lastObservedContext: CoroutineContext? = null
 
     override suspend fun emit(value: T)  {
-        val emitContext = coroutineContext.minusKey(Job).minusId()
-        if (emitContext != collectContext) {
-            error(
-                "Flow invariant is violated: flow was collected in $collectContext, but emission happened in $emitContext. " +
-                        "Please refer to 'flow' documentation or use 'flowOn' instead")
+        /*
+         * Benign data-race here:
+         * We read potentially racy published coroutineContext, but we only use it for
+         * referential comparison (=> thus safe) and are not using it for actual comparisons.
+         */
+        val currentContext = coroutineContext
+        if (lastObservedContext !== currentContext) {
+            val emitContext = currentContext.minusKey(Job).minusId()
+            if (emitContext != collectContext) {
+                error(
+                    "Flow invariant is violated: flow was collected in $collectContext, but emission happened in $emitContext. " +
+                            "Please refer to 'flow' documentation or use 'flowOn' instead"
+                )
+            }
+            // Racy publication
+            lastObservedContext = currentContext
         }
-        collector.emit(value)
+        collector.emit(value) // TCE
     }
 }

kotlinx-coroutines-core/common/test/flow/FlowInvariantsTest.kt

@@ -34,6 +34,27 @@ class FlowInvariantsTest : TestBase() {
         }
     }
 
+    @Test
+    fun testCachedInvariantCheckResult() = runTest {
+        flow {
+            emit(1)
+
+            try {
+                kotlinx.coroutines.withContext(NamedDispatchers("foo")) {
+                    emit(1)
+                }
+                fail()
+            } catch (e: IllegalStateException) {
+                expect(2)
+            }
+
+            emit(3)
+        }.collect {
+            expect(it)
+        }
+        finish(4)
+    }
+
     @Test
     fun testWithNameContractViolated() = runTest({ it is IllegalStateException }) {
         flow {