Re-implement flatMapMerge via the channel to have context preservation property

Fixes #1440

Author: qwwdfsad

benchmarks/src/jmh/kotlin/benchmarks/YieldRelativeCostBenchmark.kt

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2016-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
+ */
+
+package benchmarks.flow
+
+import kotlinx.coroutines.*
+import kotlinx.coroutines.flow.*
+import org.openjdk.jmh.annotations.*
+import java.util.concurrent.*
+
+@Warmup(iterations = 7, time = 1, timeUnit = TimeUnit.SECONDS)
+@Measurement(iterations = 7, time = 1, timeUnit = TimeUnit.SECONDS)
+@Fork(value = 1)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MICROSECONDS)
+@State(Scope.Benchmark)
+open class FlatMapMergeBenchmark {
+
+    // Note: tests only absence of contention on downstream
+
+    @Param("10", "100", "1000")
+    private var iterations = 100
+
+    @Benchmark
+    fun flatMapUnsafe() = runBlocking {
+        benchmarks.flow.scrabble.flow {
+            repeat(iterations) { emit(it) }
+        }.flatMapMerge { value ->
+            flowOf(value)
+        }.collect {
+            if (it == -1) error("")
+        }
+    }
+
+    @Benchmark
+    fun flatMapSafe() = runBlocking {
+        kotlinx.coroutines.flow.flow {
+            repeat(iterations) { emit(it) }
+        }.flatMapMerge { value ->
+            flowOf(value)
+        }.collect {
+            if (it == -1) error("")
+        }
+    }
+
+}

benchmarks/src/jmh/kotlin/benchmarks/flow/FlatMapMergeBenchmark.kt

@@ -3,7 +3,7 @@
  */
 
 
-package benchmarks.flow.misc
+package benchmarks.flow
 
 import benchmarks.flow.scrabble.flow
 import io.reactivex.*
@@ -35,7 +35,7 @@ import java.util.concurrent.*
 @BenchmarkMode(Mode.AverageTime)
 @OutputTimeUnit(TimeUnit.MICROSECONDS)
 @State(Scope.Benchmark)
-open class Numbers {
+open class NumbersBenchmark {
 
     companion object {
         private const val primes = 100

benchmarks/src/jmh/kotlin/benchmarks/flow/misc/Numbers.kt renamed to benchmarks/src/jmh/kotlin/benchmarks/flow/NumbersBenchmark.kt

@@ -2,7 +2,7 @@
  * Copyright 2016-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license.
  */
 
-package benchmarks.flow.misc
+package benchmarks.flow
 
 import kotlinx.coroutines.*
 import kotlinx.coroutines.flow.*

benchmarks/src/jmh/kotlin/benchmarks/flow/misc/SafeFlowBenchmark.kt renamed to benchmarks/src/jmh/kotlin/benchmarks/flow/SafeFlowBenchmark.kt

@@ -992,7 +992,7 @@ public final class kotlinx/coroutines/flow/internal/SafeCollectorKt {
 	public static final fun unsafeFlow (Lkotlin/jvm/functions/Function2;)Lkotlinx/coroutines/flow/Flow;
 }
 
-public final class kotlinx/coroutines/flow/internal/SendingCollector : kotlinx/coroutines/flow/internal/ConcurrentFlowCollector {
+public final class kotlinx/coroutines/flow/internal/SendingCollector : kotlinx/coroutines/flow/FlowCollector {
 	public fun <init> (Lkotlinx/coroutines/channels/SendChannel;)V
 	public fun emit (Ljava/lang/Object;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 }

binary-compatibility-validator/reference-public-api/kotlinx-coroutines-core.txt

@@ -61,7 +61,7 @@ public abstract class ChannelFlow<T>(
     internal val collectToFun: suspend (ProducerScope<T>) -> Unit
         get() = { collectTo(it) }
 
-    internal val produceCapacity: Int
+    private val produceCapacity: Int
         get() = if (capacity == Channel.OPTIONAL_CHANNEL) Channel.BUFFERED else capacity
 
     open fun broadcastImpl(scope: CoroutineScope, start: CoroutineStart): BroadcastChannel<T> =
@@ -140,13 +140,11 @@ internal class ChannelFlowOperatorImpl<T>(
 private fun <T> FlowCollector<T>.withUndispatchedContextCollector(emitContext: CoroutineContext): FlowCollector<T> = when (this) {
     // SendingCollector & NopCollector do not care about the context at all and can be used as is
     is SendingCollector, is NopCollector -> this
-    // Original collector is concurrent, so wrap into ConcurrentUndispatchedContextCollector (also concurrent)
-    is ConcurrentFlowCollector -> ConcurrentUndispatchedContextCollector(this, emitContext)
     // Otherwise just wrap into UndispatchedContextCollector interface implementation
     else -> UndispatchedContextCollector(this, emitContext)
 }
 
-private open class UndispatchedContextCollector<T>(
+private class UndispatchedContextCollector<T>(
     downstream: FlowCollector<T>,
     private val emitContext: CoroutineContext
 ) : FlowCollector<T> {
@@ -157,12 +155,6 @@ private open class UndispatchedContextCollector<T>(
         withContextUndispatched(emitContext, countOrElement, emitRef, value)
 }
 
-// named class for a combination of UndispatchedContextCollector & ConcurrentFlowCollector interface
-private class ConcurrentUndispatchedContextCollector<T>(
-    downstream: ConcurrentFlowCollector<T>,
-    emitContext: CoroutineContext
-) : UndispatchedContextCollector<T>(downstream, emitContext), ConcurrentFlowCollector<T>
-
 // Efficiently computes block(value) in the newContext
 private suspend fun <T, V> withContextUndispatched(
     newContext: CoroutineContext,

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

@@ -4,78 +4,17 @@
 
 package kotlinx.coroutines.flow.internal
 
-import kotlinx.atomicfu.*
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.*
-import kotlinx.coroutines.channels.ArrayChannel
 import kotlinx.coroutines.flow.*
 
-internal fun <T> FlowCollector<T>.asConcurrentFlowCollector(): ConcurrentFlowCollector<T> =
-    this as? ConcurrentFlowCollector<T> ?: SerializingCollector(this)
-
-// Flow collector that supports concurrent emit calls.
-// It is internal for now but may be public in the future.
-// Two basic implementations are here: SendingCollector and ConcurrentFlowCollector
-internal interface ConcurrentFlowCollector<T> : FlowCollector<T>
-
 /**
- * Collection that sends to channel. It is marked as [ConcurrentFlowCollector] because it can be used concurrently.
- *
+ * Collection that sends to channel
  * @suppress **This an internal API and should not be used from general code.**
  */
 @InternalCoroutinesApi
 public class SendingCollector<T>(
     private val channel: SendChannel<T>
-) : ConcurrentFlowCollector<T> {
+) : FlowCollector<T> {
     override suspend fun emit(value: T) = channel.send(value)
 }
-
-// Effectively serializes access to downstream collector for merging
-// This is basically a converted from FlowCollector interface to ConcurrentFlowCollector
-private class SerializingCollector<T>(
-    private val downstream: FlowCollector<T>
-) : ConcurrentFlowCollector<T> {
-    // Let's try to leverage the fact that merge is never contended
-    // Should be Any, but KT-30796
-    private val _channel = atomic<ArrayChannel<Any?>?>(null)
-    private val inProgressLock = atomic(false)
-
-    private val channel: ArrayChannel<Any?>
-        get() = _channel.updateAndGet { value ->
-            if (value != null) return value
-            ArrayChannel(Channel.CHANNEL_DEFAULT_CAPACITY)
-        }!!
-
-    public override suspend fun emit(value: T) {
-        if (!inProgressLock.tryAcquire()) {
-            channel.send(value ?: NULL)
-            if (inProgressLock.tryAcquire()) {
-                helpEmit()
-            }
-            return
-        }
-        downstream.emit(value)
-        helpEmit()
-    }
-
-    @Suppress("UNCHECKED_CAST")
-    private suspend fun helpEmit() {
-        while (true) {
-            while (true) {
-                val element = _channel.value?.poll() ?: break // todo: pollOrClosed
-                downstream.emit(NULL.unbox(element))
-            }
-            inProgressLock.release()
-            // Enforce liveness
-            if (_channel.value?.isEmpty != false || !inProgressLock.tryAcquire()) break
-        }
-    }
-}
-
-@Suppress("NOTHING_TO_INLINE")
-private inline fun AtomicBoolean.tryAcquire(): Boolean = compareAndSet(false, true)
-
-@Suppress("NOTHING_TO_INLINE")
-private inline fun AtomicBoolean.release() {
-    value = false
-}

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

@@ -37,21 +37,21 @@ internal class ChannelFlowTransformLatest<T, R>(
     }
 }
 
-internal class ChannelFlowMerge<T>(
+internal class ChannelFlowMerge<T>  (
     flow: Flow<Flow<T>>,
     private val concurrency: Int,
     context: CoroutineContext = EmptyCoroutineContext,
-    capacity: Int = Channel.OPTIONAL_CHANNEL
+    capacity: Int = Channel.BUFFERED
 ) : ChannelFlowOperator<Flow<T>, T>(flow, context, capacity) {
     override fun create(context: CoroutineContext, capacity: Int): ChannelFlow<T> =
         ChannelFlowMerge(flow, concurrency, context, capacity)
 
     override fun produceImpl(scope: CoroutineScope): ReceiveChannel<T> {
-        return scope.flowProduce(context, produceCapacity, block = collectToFun)
+        return scope.flowProduce(context, capacity, block = collectToFun)
     }
 
     // The actual merge implementation with concurrency limit
-    private suspend fun mergeImpl(scope: CoroutineScope, collector: ConcurrentFlowCollector<T>) {
+    private suspend fun mergeImpl(scope: CoroutineScope, collector: SendingCollector<T>) {
         val semaphore = Semaphore(concurrency)
         val job: Job? = coroutineContext[Job]
         flow.collect { inner ->
@@ -72,12 +72,15 @@ internal class ChannelFlowMerge<T>(
         }
     }
 
-    // Fast path in ChannelFlowOperator calls this function (channel was not created yet)
     override suspend fun flowCollect(collector: FlowCollector<T>) {
-        // this function should not have been invoked when channel was explicitly requested
-        assert { capacity == Channel.OPTIONAL_CHANNEL }
-        flowScope {
-            mergeImpl(this, collector.asConcurrentFlowCollector())
+        assert { collector !is SendingCollector<*> }
+        coroutineScope {
+            val output = produce<T>(capacity = capacity) {
+                mergeImpl(this, SendingCollector(this))
+            }
+            output.consumeEach {
+                collector.emit(it)
+            }
         }
     }
 

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

@@ -4,7 +4,9 @@
 
 package kotlinx.coroutines.flow.internal
 
-internal object NopCollector : ConcurrentFlowCollector<Any?> {
+import kotlinx.coroutines.flow.*
+
+internal object NopCollector : FlowCollector<Any?> {
     override suspend fun emit(value: Any?) {
         // does nothing
     }

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

@@ -14,7 +14,7 @@ class FlattenMergeTest : FlatMapMergeBaseTest() {
     @Test
     override fun testFlatMapConcurrency() = runTest {
         var concurrentRequests = 0
-        val flow = (1..100).asFlow().map() { value ->
+        val flow = (1..100).asFlow().map { value ->
             flow {
                 ++concurrentRequests
                 emit(value)
@@ -36,4 +36,19 @@ class FlattenMergeTest : FlatMapMergeBaseTest() {
         consumer.cancelAndJoin()
         finish(3)
     }
+
+    @Test
+    fun testContextPreservationAcrossFlows() = runTest {
+        val result = flow {
+            flowOf(1, 2).flatMapMerge {
+                flow {
+                    yield()
+                    emit(it)
+                }
+            }.collect {
+                emit(it)
+            }
+        }.toList()
+        assertEquals(listOf(1, 2), result)
+    }
 }