Kotlin协程之flow工作原理

Posted 2021-05-28 斯音

概述

最近想学习一下 Kotlin 中 flow 的用法, Google 上搜了搜发现很多比较 RxJava 和 flow 的文章，其实我在实际业务中从来没有用过 RxJava, 倒不是因为它不好，而是…我一直傻傻不太会用 RxJava 的操作符，看不太懂，又一直没花时间(懒惰)去研究它那些操作符的原理，就一直不怎么敢用。这次看到了 flow, 想着还是先去了解了解它内部几个操作符的原理吧，不然用起来总是不太踏实。

需要注意的是 Flow 需要在协程中使用, 因此配合协程，可以方便地切线程。分析 flow 工作流程离不开协程的工作原理，关于 Kotlin 协程的解析可以参考下列文章：

• Kotlin协程之基础使用
• Kotlin协程之深入理解协程工作原理
• Kotlin协程之协程取消与异常处理

首先看一下 Flow 接口的源码，内部只有一个 collect 方法：

public interface Flow<out T> {
    // 是一个 suspend 方法，意味着会挂起当前协程
    @InternalCoroutinesApi
    public suspend fun collect(collector: FlowCollector<T>)
}

public interface FlowCollector<in T> {
    // 数据的发射方
    public suspend fun emit(value: T)
}

flow {}

以下面代码为例，讲解 flow 工作的基本流程：

flow { emit(1) }.collect { println(it) }

首先看一下 flow {} 的源码：

public fun <T> flow(@BuilderInference block: suspend FlowCollector<T>.() -> Unit): Flow<T> = SafeFlow(block)

上面就是以 block 代码块为参数创建了一个 SafeFlow 对象，SafeFlow 实现了 Flow 接口，于是接着看其 collect 方法。

collect

除了一开始贴的实现 Flow 接口调用 collect 方法的方式, Kotlin 还提供了调用 collect 的两个扩展函数，最后都是调用的 fun collect(collector: FlowCollector<T>) 方法：

public suspend fun Flow<*>.collect(): Unit = collect(NopCollector)

public suspend inline fun <T> Flow<T>.collect(crossinline action: suspend (value: T) -> Unit): Unit =
    collect(object : FlowCollector<T> {
        override suspend fun emit(value: T) = action(value)
    })

于是我们接着上面的示例，看一下 SafeFlow.collect 方法：

private class SafeFlow<T>(private val block: suspend FlowCollector<T>.() -> Unit) : AbstractFlow<T>() {
    override suspend fun collectSafely(collector: FlowCollector<T>) {
        collector.block()
    }
}

// collect 方法在父类 AbstractFlow 中
public abstract class AbstractFlow<T> : Flow<T>, CancellableFlow<T> {
    public final override suspend fun collect(collector: FlowCollector<T>) {
        val safeCollector = SafeCollector(collector, coroutineContext)
        try {
            collectSafely(safeCollector)
        } finally {
            safeCollector.releaseIntercepted()
        }
    }

    public abstract suspend fun collectSafely(collector: FlowCollector<T>)
}

可以看到 collect 方法中通过 collector 封装了一个 SafeCollector 对象，并以其为参数执行了 SafeFlow.collectSafely 方法，而 collectSafely 方法只是执行了 block 代码块(collector.block())，它是一个扩展函数，所以执行的示例代码中的 emit(1) 其实就是调用了 SafeCollector.emit(1), 然后在 SafeCollector 中对 FlowCollector 做了一层安全校验后，最后还是会调用 FlowCollector.emit 方法，即创建 SafeCollector 时传入的 collector 对象的 emit 方法。这里只关注核心流程，故不贴出具体代码了。

根据上面我们看到的 collect {} 扩展函数的源码，可以知道其 emit 方法其实就是执行 collect {} 中传入的 action 代码块，参数为 emit 发射的值 – 1.

小结：flow {} 方式(或flowOf, asFlow)创建的 Flow 实例是 SafeFlow 类型，其父类是 AbstractFlow 抽象类，当调用其 collect(FlowCollector) 方法时，首先会执行该 Flow 对象传入的 block 代码块，代码块中一般会有 emit 方法发射值，这个 emit 调用的就是 AbstractFlow.emit 方法，在其中做了安全判定后，会接着调用到 collect 中传入的 FlowCollector.emit 方法，对于 collect {} 的情况，emit 方法内部就是执行 collect 传入的 action 代码块。因为它在每次调用 collect 时才去触发发送数据的动作，所以说 Flow 是冷流。

主要流程如下图：

flowOn

学习 flow 一个绕不开的操作符就是 flowOn 了，以下面示例代码为例, flow 需要在协程中使用，下面的 emit(1) 会在 Dispatchers.Default 指定的线程中执行，而 println(it) 会在父协程所在线程中执行：

flow { emit(1) }.flowOn(Dispatchers.Default).collect { println(it) }

flow {} 的源码在上面已经看过了，就是以 block 代码块为参数创建了一个 SafeFlow 对象，接下来看一下 Flow.flowOn 的逻辑：

public fun <T> Flow<T>.flowOn(context: CoroutineContext): Flow<T> {
    checkFlowContext(context)
    return when {
        // 返回自身 Flow 实例
        // 这里我们传入了 Dispatchers.Default, 所以不符合这个条件
        context == EmptyCoroutineContext -> this
        // SafeFlow 不是该类型，因此也不走这个流程，实际上 FusibleFlow 是当连续多次调用 flowOn 后会创建的 Flow 对象
        this is FusibleFlow -> fuse(context = context)
        // 逻辑走到这里
        else -> ChannelFlowOperatorImpl(this, context = context)
    }
}

在上面已经对流程注释了一下，因此上述实例代码转换一下即为: SafeFlow.flowOn.collect {} --> ChannelFlowOperatorImpl.collect {}, 这里注意一下创建 ChannelFlowOperatorImpl 对象时传入的两个参数，第一个 this 指的是之前的 SafeFlow 对象，第二个 context 参数即是我们传入的调度器，它是一个协程上下文。

ChannelFlowOperatorImpl.collect 实现在父类 ChannelFlowOperator.collect 中，该方法如果发现传入的 coroutineContext 上下文中没有携带调度器，即我们调用 flowOn 时没有传入 Dispatchers 等调度器，则会直接调用上一层 SafeFlow 的 collect 方法(代码不贴了)，否则接着调用父类 ChannelFlow 中的 collect 方法，我们直接看 flowOn 中传入了调度器后的逻辑：

internal abstract class ChannelFlowOperator<S, T>(
    @JvmField protected val flow: Flow<S>,
    context: CoroutineContext,
    capacity: Int,
    onBufferOverflow: BufferOverflow
) : ChannelFlow<T>(context, capacity, onBufferOverflow) {
    override suspend fun collect(collector: FlowCollector<T>) {
        // 判断 coroutineContext 逻辑
        // ...
        super.collect(collector) // 调用父类 ChannelFlow 中方法
    }
}

public abstract class ChannelFlow<T>(
    // upstream context
    @JvmField public val context: CoroutineContext,
    // buffer capacity between upstream and downstream context
    @JvmField public val capacity: Int,
    // buffer overflow strategy
    @JvmField public val onBufferOverflow: BufferOverflow
) : FusibleFlow<T> {
    override suspend fun collect(collector: FlowCollector<T>): Unit =
        coroutineScope {
            collector.emitAll(produceImpl(this))
        }

    public open fun produceImpl(scope: CoroutineScope): ReceiveChannel<T> =
        scope.produce(context, produceCapacity, onBufferOverflow, start = CoroutineStart.ATOMIC, block = collectToFun)
}

这里可以看到 ChannelFlowOperatorImpl.collect 最后会走到 collector.emitAll(produceImpl(this)) 生产消费的逻辑，我们分步骤看一下生产和接收的流程。

生产数据

首先看上面 produceImpl 方法：

internal fun <E> CoroutineScope.produce(
    context: CoroutineContext = EmptyCoroutineContext,
    capacity: Int = 0,
    onBufferOverflow: BufferOverflow = BufferOverflow.SUSPEND,
    start: CoroutineStart = CoroutineStart.DEFAULT,
    onCompletion: CompletionHandler? = null,
    @BuilderInference block: suspend ProducerScope<E>.() -> Unit
): ReceiveChannel<E> {
    val channel = Channel<E>(capacity, onBufferOverflow)
    val newContext = newCoroutineContext(context)
    val coroutine = ProducerCoroutine(newContext, channel)
    if (onCompletion != null) coroutine.invokeOnCompletion(handler = onCompletion)
    coroutine.start(start, coroutine, block)
    return coroutine
}

看到这个方法，是不是很熟悉呢？参考之前的 Kotlin之深入理解协程工作原理 的文章可以知道，这里的 produce 方法其实就是启动了一个新的协程，该协程执行的代码块 block 是传入的 collectToFun 参数，接着找 collectToFun 可以发现它会取 ChannelFlowOperator.collectTo 方法：

// ChannelFlowOperator
protected override suspend fun collectTo(scope: ProducerScope<T>) =
    // flowCollect 方法实现在子类 ChannelFlowOperatorImpl 中
    flowCollect(SendingCollector(scope))

// ChannelFlowOperatorImpl
internal class ChannelFlowOperatorImpl<T>(
    flow: Flow<T>,
    context: CoroutineContext = EmptyCoroutineContext,
    capacity: Int = Channel.OPTIONAL_CHANNEL,
    onBufferOverflow: BufferOverflow = BufferOverflow.SUSPEND
) : ChannelFlowOperator<T, T>(flow, context, capacity, onBufferOverflow) {
    override suspend fun flowCollect(collector: FlowCollector<T>) =
        // 这个 flow 就是上层传入的 SafeFlow 对象
        flow.collect(collector)
}

根据之前的解析, flow.collect(collector) 中的 flow 是 SafeFlow 对象，其 collect 方法会执行 SafeFlow 中传入的代码块(即flow {})，这个代码块中调用了 collector.emit(1) 方法(上面代码可以看出此时的 collector 是 SendingCollector 实例)，因此我们看看 SendingCollector.emit 方法做了什么：

public class SendingCollector<T>(
    private val channel: SendChannel<T>
) : FlowCollector<T> {
    override suspend fun emit(value: T): Unit = channel.send(value)
}

于是可以知道 produceImpl 方法就是启动了一个新的协程，然后在协程中执行上层 flow 对象(所以 flowOn 会对它上游的部分起作用)中的代码块(里面调用了 SendingCollector.emit 方法)，然后通过 Channel.send 方法把这个 value 发送出去。

接收数据

上面看了启动协程并在其内通过 Channel 发送数据的流程，这里看一下数据是怎么接收的，回到最开始的代码，从 collector.emitAll(channel) 开始，这个 channel 参数就是上一节上调用 send 发送数据的那个 channel 对象：

public suspend fun <T> FlowCollector<T>.emitAll(channel: ReceiveChannel<T>): Unit =
    emitAllImpl(channel, consume = true)

private suspend fun <T> FlowCollector<T>.emitAllImpl(channel: ReceiveChannel<T>, consume: Boolean) {
    ensureActive()
    var cause: Throwable? = null
    try {
        while (true) {
            val result = run { channel.receiveCatching() }
            if (result.isClosed) {
                result.exceptionOrNull()?.let { throw it }
                break // returns normally when result.closeCause == null
            }
            emit(result.getOrThrow())
        }
    } catch (e: Throwable) {
        cause = e
        throw e
    } finally {
        if (consume) channel.cancelConsumed(cause)
    }
}

这里可以看到开了一个无限循环，然后通过 Channel 去接收数据，并通过 emit 方法把接收的值发射出去，至于调用这个 emit 方法的 FlowCollector 对象是谁呢？再回到一开始 flow { emit(1) }.flowOn(Dispatchers.Default).collect { println(it) } 示例中最后面调用的 collect 方法，结合上一章 collect 的解析，可以知道这个 FlowCollector 就是通过 collect 方法传入的代码块创建的对象：

public suspend inline fun <T> Flow<T>.collect(crossinline action: suspend (value: T) -> Unit): Unit =
    collect(object : FlowCollector<T> {
        override suspend fun emit(value: T) = action(value)
    })

于是最终在收到数据并 emit 后，会把 value 传递给 collect {} 中的代码块去执行。

多个flowOn

以下面代码为例：

flow { emit(1) }.flowOn(Dispatchers.IO).flowOn(Dispatchers.Main).collect { println(it) }

由之前的解析可以知道首先会调用到这里：

public fun <T> Flow<T>.flowOn(context: CoroutineContext): Flow<T> {
    checkFlowContext(context)
    return when {
        context == EmptyCoroutineContext -> this
        this is FusibleFlow -> fuse(context = context)
        else -> ChannelFlowOperatorImpl(this, context = context)
    }
}

第一次调用 flowOn 时返回的是 ChannelFlowOperatorImpl 对象，查看其继承关系可以知道它实现了 FusibleFlow 接口，因此第二次调用 flowOn 时会走 fuse(context = context) 逻辑：

public abstract class ChannelFlow<T>(
    // upstream context
    @JvmField public val context: CoroutineContext,
    // buffer capacity between upstream and downstream context
    @JvmField public val capacity: Int,
    // buffer overflow strategy
    @JvmField public val onBufferOverflow: BufferOverflow
) : FusibleFlow<T> {
    public override fun fuse(context: CoroutineContext, capacity: Int, onBufferOverflow: BufferOverflow): Flow<T> {
        val newContext = context + this.context
        val newCapacity: Int
        // 处理 newCapacity 等计算逻辑
        // ...
        if (newContext == this.context && newCapacity == this.capacity && newOverflow == this.onBufferOverflow)
            return this
        return create(newContext, newCapacity, newOverflow)
    }
}

上述 fuse 方法要返回一个 Flow 对象，注意这一行代码: val newContext = context + this.context, 新的上下文 newContext = 传进来的 context + 之前的 this.context, 对协程上下文 CoroutineContext 的结构比较了解的话可以知道，它有点类似于一个 key-value 形式的数据结构，因此对于调度器 Dispatcher 这个 key 而言，最终存在的调度器将会是 this.context, 因为它会覆盖掉传入的 context(因为它在加号前)，而 this.context 是之前的上下文，在这里即是 Dispatchers.IO, 所以第二个 flowOn 就无效了，感兴趣的小伙伴可以自己验证下。

小结

• produceImpl 方法就是启动了一个新的协程，然后在协程中执行上层 flow 对象(所以 flowOn 会对它上游的部分起作用)中的代码块(里面调用了 SendingCollector.emit 方法)，然后通过 Channel.send 方法把这个 value 发送出去。此时运行中 flowOn 指定的线程中。
• collector.emitAll 方法会通过上面的 channel 对象接收数据，并调用 FlowCollector.emit 方法来发射数据，最后执行到 SafeFlow.flowOn.collect {} collect 代码块中的逻辑。此时运行在父协程指定的线程中。

flowOn 流程如下图：