Kotlin 协程源码阅读笔记 —— SharedFlow 和 StateFlow

Kotlin 协程源码阅读笔记 ------ SharedFlow 和 StateFlow

我在前面的文章中讲了 Flow，它是 Kotlin 协程中的冷流：Kotlin 协程源码阅读笔记 ------ Flow。

前面的文章中也讲过 Channel：Kotlin 协程源码阅读笔记 ------ Channel，它可以勉强地算是一种热流吧，首先每次发送对应的只有一个地方能够收到这次的数据；而且它也不支持 Flow 中非常灵活的操作符，用它来写简单的生产者消费者模型非常方便。

今天要讲的内容是 Kotlin 协程中的 SharedFlow 和 StateFlow，他们是 Kotlin 协程中的热流的实现。对应到 RxJava 中就是 PublishSubject 和 BehaviorSubject。但是 SharedFlow 与 StateFlow 的功能比 RxJava 中的 BehaviorSubject 与 PublishSubject 更加强大，它们也可以使用 Flow 中功能非常强大的各种操作符。

源码阅读基于 Kotlin 协程 1.8.0-RC2

SharedFlow 和 StateFlow 的简单使用

SharedFlow

简单发送消息和接收消息

runBlocking(Dispatchers.Default) {
    val sharedFlow = MutableSharedFlow<Int>()
    launch {
        repeat(10) {
            delay(100)
            sharedFlow.emit(it)
            println("Coroutine1: Send $it")
        }
    }
    launch {
        delay(200)
        sharedFlow.collect {
            println("Coroutine2: Receive $it")
        }
    }
    launch {
        delay(500)
        sharedFlow.collect {
            println("Coroutine3: Receive $it")
        }
    }
}

我通过协程1来发送 10 个元素，而且，每次发送元素耗时 100ms。我通过协程2和协程3来接收元素，他们分别延迟 200ms 和 500ms 才订阅 ShreadFlow。

协程2打印结果如下：

Coroutine2: Receive 2
Coroutine2: Receive 3
Coroutine2: Receive 4
Coroutine2: Receive 5
Coroutine2: Receive 6
Coroutine2: Receive 7
Coroutine2: Receive 8
Coroutine2: Receive 9

由于它延迟了 200ms，所以丢弃了元素 0 和 1。

协程3打印结果如下：

Coroutine3: Receive 4
Coroutine3: Receive 5
Coroutine3: Receive 6
Coroutine3: Receive 7
Coroutine3: Receive 8
Coroutine3: Receive 9

由于协程3延迟的时间更长，它丢弃了 0，1，2，3。

所以到这里理解 ShrareFlow 和普通 Flow 的区别了吗？这也是热流和冷流的区别。

理解 replay 参数

默认情况下 replay 参数设置是 0，它是用来描述需要发送给新订阅者缓存的数据数量，说起来有点抽象，来看例子吧。

runBlocking(Dispatchers.Default) {
    val sharedFlow = MutableSharedFlow<Int>()
    repeat(3) {
        sharedFlow.emit(it)
    }
    launch {
        sharedFlow.collect {
            println("Coroutine1: Receive $it")
        }
    }
}

我在订阅前发送了 3 个元素，由于是订阅前发送的，所以协程1一个元素也接收不到。我们将上面的代码改改，设置 replay 为 2：

runBlocking(Dispatchers.Default) {
    val sharedFlow = MutableSharedFlow<Int>(replay = 2)
    repeat(3) {
        sharedFlow.emit(it)
    }
    launch {
        sharedFlow.collect {
            println("Coroutine1: Receive $it")
        }
    }
}

最后的输出是：

Coroutine1: Receive 1
Coroutine1: Receive 2

我们惊奇地发现它能够收到两个元素了，通过 replay 参数可以帮我们缓存对应数量的元素，如果到达 replay 的上线就会丢弃一条最旧的元素。

理解 extraBufferCapacity 参数。

replay + extraBufferCapacity 就是最终的 BufferCapacity，如果发送元素的速度大于订阅者处理这些元素的速度时，新的发送的数据就会存放在 BufferCapacity 中，当 BufferCapacity 满了后，就会触发数据溢出的处理策略 SUSPEND （挂起 emit() 函数，直到订阅者处理完成后再恢复），DROP_OLDEST（丢弃 BufferCapacity 中的最久的一条数据），DROP_LATEST（丢弃 BufferCacaity 中最新的一条数据），我们以溢出策略 SUSPEND (这也是默认的策略) 来看看 extraBufferCapacity 参数。

runBlocking(Dispatchers.Default) {
    val sharedFlow = MutableSharedFlow<Int>()
    launch {
        delay(100)
        repeat(10) {
            val start = System.currentTimeMillis()
            sharedFlow.emit(it)
            val end = System.currentTimeMillis()
            println("Coroutine1: Send $it cost ${end - start}ms.")
        }
    }
    launch {
        sharedFlow.collect {
            delay(100)
            println("Coroutine2: Receive $it")
        }
    }
}

我在协程1中发送元素，记录了它发送元素的耗时，我在协程2中接收元素每次接收都耗时 100ms，最后的输入日志是：

Coroutine1: Send 0 cost 53ms.
Coroutine2: Receive 0
Coroutine1: Send 1 cost 111ms.
Coroutine2: Receive 1
Coroutine1: Send 2 cost 106ms.
Coroutine2: Receive 2
Coroutine1: Send 3 cost 107ms.
Coroutine2: Receive 3
Coroutine1: Send 4 cost 107ms.
Coroutine2: Receive 4
Coroutine1: Send 5 cost 111ms.
Coroutine2: Receive 5
Coroutine1: Send 6 cost 106ms.
Coroutine2: Receive 6
Coroutine1: Send 7 cost 107ms.
Coroutine2: Receive 7
Coroutine1: Send 8 cost 103ms.
Coroutine2: Receive 8
Coroutine1: Send 9 cost 108ms.
Coroutine2: Receive 9

我们看到除了第一次，其他的发送数据耗时都接近 100ms，这是因为接收的耗时会让发送时等待。我们试试设置 extraBufferCapacity 为 3 再来看看。

runBlocking(Dispatchers.Default) {
    val sharedFlow = MutableSharedFlow<Int>(extraBufferCapacity = 3)
    launch {
        delay(100)
        repeat(10) {
            val start = System.currentTimeMillis()
            sharedFlow.emit(it)
            val end = System.currentTimeMillis()
            println("Coroutine1: Send $it cost ${end - start}ms.")
        }
    }
    launch {
        sharedFlow.collect {
            delay(100)
            println("Coroutine2: Receive $it")
        }
    }
}

最后的输出结果就是：

Coroutine1: Send 0 cost 4ms.
Coroutine1: Send 1 cost 0ms.
Coroutine1: Send 2 cost 0ms.
Coroutine1: Send 3 cost 55ms.
Coroutine2: Receive 0
Coroutine1: Send 4 cost 113ms.
Coroutine2: Receive 1
Coroutine1: Send 5 cost 106ms.
Coroutine2: Receive 2
Coroutine1: Send 6 cost 106ms.
Coroutine2: Receive 3
Coroutine1: Send 7 cost 106ms.
Coroutine2: Receive 4
Coroutine1: Send 8 cost 109ms.
Coroutine2: Receive 5
Coroutine1: Send 9 cost 101ms.
Coroutine2: Receive 6
Coroutine2: Receive 7
Coroutine2: Receive 8
Coroutine2: Receive 9

我们看到前三条数据几乎不耗时，而后面的数据又和上面测试的一样耗时接近 100ms。所以到这里理解这个 extraBufferCapacity 了吧，我只测试了 SUSPEND 这种 BufferOverflow 的策略，DROP_OLDEST 与 DROP_LATEST 他们就是丢弃数据了，大家可以自己测试一下。

StateFlow

StateFlow 就相当于 SharedFlow 的 replay 为 1，同时它必须有一个默认的元素，而且它的 emit() 方法永远不会被挂起，订阅时总是能够拿到它的最新的一个元素。

举一个例子：

runBlocking(Dispatchers.Default) {
    val stateFlow = MutableStateFlow(-1)

    launch {
        delay(100)
        repeat(10) {
            val start = System.currentTimeMillis()
            stateFlow.emit(it)
            val end = System.currentTimeMillis()
            println("Coroutine1: Send $it cost ${end - start}ms")
        }
    }

    launch {
        stateFlow.collect {
            println("Coroutine2: Receive $it")
        }
    }
}

最后的输出是：

Coroutine2: Receive -1
Coroutine1: Send 0 cost 0ms
Coroutine1: Send 1 cost 0ms
Coroutine1: Send 2 cost 0ms
Coroutine1: Send 3 cost 0ms
Coroutine1: Send 4 cost 0ms
Coroutine1: Send 5 cost 0ms
Coroutine1: Send 6 cost 0ms
Coroutine1: Send 7 cost 0ms
Coroutine1: Send 8 cost 0ms
Coroutine1: Send 9 cost 0ms
Coroutine2: Receive 9

StateFlow 还有一个特性就是会过滤掉连续发送的相同的元素，我这里就不针对这个特性再写例子了，大家可以自己试试。

综上 SharedFlow 更加适合用来处理某种事件，而 StateFlow 更容易来保存和发送某种状态。

SharedFlow 工作原理

在开始分析源码之前这里有必要大致描述一下 SharedFlow 的工作原理，它的实现类是 SharedFlowImpl，它是继承于 AbstractSharedFlow。

当调用订阅时会将创建一个 FlowSlot 它内部记录了获取元素开始的 index（就是 Buffer 中的 index，表示从这个 index 开始获取元素），然后就是一个死循环从 Buffer 中去获取对应的 index 中的元素，如果获取失败就会将当前的函数挂起，对应的 suspend 函数的 Continuation 会被保存在 FlowSlot 中，等待发送数据后将其唤醒继续去获取元素；如果获取数据成功就将 index 的值加 1，表示下次获取 Buffer 中的下一个元素。

当发送数据时就需要判断 Buffer 是否已经达到最大的 Capacity，如果没有达到最大的 Capacity 那么就直接存入到对应的 Buffer 中，同时去查找满足条件的订阅者的 FlowSlot 然后，调用它的 Contiuation#resumeWith() 方法去唤醒对应的 suspend 方法，然后去读取这个元素；如果 Capacity 已经达到最大值，那么就会触发 BufferOverflow 的策略，DROP_OLDEST 和 DROP_LATEST 就相对比较简单，直接暴力地丢掉最旧和最新的一个元素，如果是 SUSPEND 就要复杂一些，这个时候就会挂起 emit() 方法，然后将其对应要发送的元素和 Continuation 对象存放在 Emitter 中，当所有的订阅者都读取了 Emitter 中的元素后，就会调用其 Contianution#resumeWith() 方法恢复对应的 emit() 方法。

为了方便理解源码，这里介绍一下其中对应的重要的成员变量：

// Stored state
private var buffer: Array<Any?>? = null // allocated when needed, allocated size always power of two
private var replayIndex = 0L // minimal index from which new collector gets values
private var minCollectorIndex = 0L // minimal index of active collectors, equal to replayIndex if there are none
private var bufferSize = 0 // number of buffered values
private var queueSize = 0 // number of queued emitters

// Computed state
private val head: Long get() = minOf(minCollectorIndex, replayIndex)
private val replaySize: Int get() = (head + bufferSize - replayIndex).toInt()
private val totalSize: Int get() = bufferSize + queueSize
private val bufferEndIndex: Long get() = head + bufferSize
private val queueEndIndex: Long get() = head + bufferSize + queueSize

• buffer
  存放发送数据的缓存，也包括挂起 emit() 方法时的 Emitter 对象。
• replayIndex
  依赖于我们在介绍 SharedFlow 使用时的 replay 参数，它是用来描述 replay 开始时在 buffer 中的 index。
• bufferSize
  普通缓存的数量，也就是不算 Emitter 这种挂起的数据。
• queueSize
Emitter 中的数量，Emitter 是存放在普通的缓存后面。
• minCollectorIndex
  所有的订阅者中处理的最小的元素的 index，这之前的 Buffer 中的数据都会被清除掉。
• head
  从 Buffer 中开始读取数据的 index，它是取的 minCollectorIndex 与 replayIndex 中的较小值。
• replaySize 也就是当前需要 replay 的数量。
• totalSize
bufferSize 与 queueSize 的和。
• bufferEndIndex
  普通缓存结束时的 index。
• queueEndIndex
  被挂起 Emitter 结束时的 index。

这里再给一张官方描述 Buffer 的图：

订阅

我们看看 SharedFlowImpl#collect() 方法的实现：

override suspend fun collect(collector: FlowCollector<T>): Nothing {
    // 创建一个 FlowSlot
    val slot = allocateSlot()
    try {
        if (collector is SubscribedFlowCollector) collector.onSubscription()
        val collectorJob = currentCoroutineContext()[Job]
        // 死循环去读取元素
        while (true) {
            var newValue: Any?
            while (true) {
                // 获取元素
                newValue = tryTakeValue(slot) // attempt no-suspend fast path first
                // 获取元素成功直接退出循环
                if (newValue !== NO_VALUE) break
                // 获取元素失败挂起，等待唤醒下次继续获取。
                awaitValue(slot) // await signal that the new value is available
            }
            collectorJob?.ensureActive()
            // 将获取成功的元素发送给 FlowCollector
            collector.emit(newValue as T)
        }
    } finally {
        freeSlot(slot)
    }
}

首先调用 AbstractSharedFlow#allocateSlot() 方法去创建一个 FlowSlot；然后进入死循环尝试读取数据，通过 tryTakeValue() 方法去尝试读取需要的数据，如果读取成功，直接通过 emit() 方法发送给订阅者，如果读取失败就会通过 awaitValue() 方法挂起当前的方法，等待有数据被唤醒后重新去读取。

我们看看 AbstractSharedFlow#allocateSlot() 方法是如何创建一个 FlowSlot 的。

protected fun allocateSlot(): S {
    // Actually create slot under lock
    val subscriptionCount: SubscriptionCountStateFlow?
    val slot = synchronized(this) {
        val slots = when (val curSlots = slots) {
            // 如果 slots 为空，创建一个数组大小为 2 的 slot
            null -> createSlotArray(2).also { slots = it }
            else -> if (nCollectors >= curSlots.size) {
                // 如果订阅者的数量大于等于原来数组的数量，将原来的数组扩展为 2 倍
                curSlots.copyOf(2 * curSlots.size).also { slots = it }
            } else {
                curSlots
            }
        }
        // 新创建的 slot 的 index
        var index = nextIndex
        var slot: S
        while (true) {
            // 调用 createSlot() 方法创建一个 slot
            slot = slots[index] ?: createSlot().also { slots[index] = it }
            index++
            if (index >= slots.size) index = 0
            
            // 通过 allocateLocked() 方法去初始化 Slot
            if ((slot as AbstractSharedFlowSlot<Any>).allocateLocked(this)) break // break when found and allocated free slot
        }
        // 更新下次插入的 slot 的 index
        nextIndex = index
        // 更新订阅者数量
        nCollectors++
        subscriptionCount = _subscriptionCount // retrieve under lock if initialized
        slot
    }
    // increments subscription count
    subscriptionCount?.increment(1)
    return slot
}

上面代码比较简单不再细说，这里有两个关键点，会调用 createSlot() 方法来创建 FlowSlot，然后通过 FlowSlot#allocateLocked() 方法来初始化 FlowSlot，他们都是在 SharedFlowImpl 中实现的。

我们看看 SharedFlowImpl#createSlot() 方法的实现：

override fun createSlot() = SharedFlowSlot()

直接返回一个 SharedFlowSlot 对象，我们看看它的实现：

internal class SharedFlowSlot : AbstractSharedFlowSlot<SharedFlowImpl<*>>() {
    @JvmField
    var index = -1L // current "to-be-emitted" index, -1 means the slot is free now

    @JvmField
    var cont: Continuation<Unit>? = null // collector waiting for new value

    override fun allocateLocked(flow: SharedFlowImpl<*>): Boolean {
        if (index >= 0) return false // not free
        index = flow.updateNewCollectorIndexLocked()
        return true
    }

    override fun freeLocked(flow: SharedFlowImpl<*>): Array<Continuation<Unit>?> {
        assert { index >= 0 }
        val oldIndex = index
        index = -1L
        cont = null // cleanup continuation reference
        return flow.updateCollectorIndexLocked(oldIndex)
    }
}

在 allocateLocked() 初始化的方法中会通过 SharedFlowImpl#updateNewCollectorIndexLocked() 方法来计算下次获取元素的 index，其实就是 replayIndex。通过 freeLocked() 方法来释放 ShraredFlowSlot，主要是把 index 设置为 -1，Continuation 设置为空，然后重新计算 SharedFlowImpl 中的各种状态。

collect() 方法中创建 FlowSlot 看完了，我们再看看 tryTakeValue() 如何获取元素：

private fun tryTakeValue(slot: SharedFlowSlot): Any? {
    var resumes: Array<Continuation<Unit>?> = EMPTY_RESUMES
    val value = synchronized(this) {
        // 获取 slot 中的 index
        val index = tryPeekLocked(slot)
        if (index < 0) {
    
            // 不可用
            NO_VALUE
        } else {
            val oldIndex = slot.index
            // 从 buffer 中获取对应的元素
            val newValue = getPeekedValueLockedAt(index)
            // 将 Slot 中的 index 加 1.
            slot.index = index + 1 // points to the next index after peeked one
            // 更新各种状态和计算需要恢复的 emit() 方法的 Continuation
            resumes = updateCollectorIndexLocked(oldIndex)
            newValue
        }
    }
    // 恢复需要恢复的 emit() 方法
    for (resume in resumes) resume?.resume(Unit)
    return value
}

private fun getPeekedValueLockedAt(index: Long): Any? =
    when (val item = buffer!!.getBufferAt(index)) {
        is Emitter -> item.value
        else -> item
    }

解释一下上面的代码:

1. 通过 tryPeekLocked() 方法计算当前的 FlowSlot 中的 index 是否可以读取元素，如果可以读取返回值大于 0，否则返回 -1，如果不能读取直接返回函数。
2. 通过 getPeekedValueLockedAt() 方法去读取 buffer 中对应的元素，然后将原来的 FlowSlot 中的 index 加 1。
3. 通过 updateCollectorIndexLocked() 方法更新 Flow 的各种状态，并获取需要恢复的 emit() 函数的 Continuation 对象。
4. 恢复满足要求的 emit() 方法。

再来看看 collect() 方法中挂起时调用的方法 awaitValue()：

private suspend fun awaitValue(slot: SharedFlowSlot): Unit = suspendCancellableCoroutine { cont ->
    synchronized(this) lock@{
        // 再次判断 Slot 中的 index 是否可用
        val index = tryPeekLocked(slot) // recheck under this lock
        if (index < 0) {
            // 不可用将当前方法挂起，并把 Continuation 对象存储在 Slot 中
            slot.cont = cont // Ok -- suspending
        } else {
            // 当前变为可用了，直接恢复当前方法
            cont.resume(Unit) // has value, no need to suspend
            return@lock
        }
        slot.cont = cont // suspend, waiting
    }
}

上面的方法中还会再次通过 tryPeekLocked() 方法判断，FlowSlot 中的 index 是否可用，如果不可用，直接将 Continuation 保存到 FlowSlot 中，等待下次发送元素时恢复；如果可用直接恢复当前方法，尝试再去获取数据。

发送数据

直接看 SharedFlowImpl#emit() 方法：

override suspend fun emit(value: T) {
    if (tryEmit(value)) return // fast-path
    emitSuspend(value)
}

首先通过 tryEmit() 方法发送数据，如果返回 true 就结束了，如果返回 false 通过 emitSuspend() 方法来将 emit() 函数挂起。

我们先看看 tryEmit() 方法的实现：

override fun tryEmit(value: T): Boolean {
    var resumes: Array<Continuation<Unit>?> = EMPTY_RESUMES
    val emitted = synchronized(this) {
        // 发送元素
        if (tryEmitLocked(value)) {
            // 发送成功，从 Slot 中查找可以恢复的 collect() 方法。
            resumes = findSlotsToResumeLocked(resumes)
            true
        } else {
            false
        }
    }
    // 恢复 collect() 方法
    for (cont in resumes) cont?.resume(Unit)
    return emitted
}

先调用 tryEmitLocked() 去发送数据，发送成功后通过 findSlotsToResumeLocked() 去查找满足条件可以恢复的 FlowSlot，也就是恢复被挂起的 collect() 函数。

我们先看看 tryEmitLocked() 方法：

private fun tryEmitLocked(value: T): Boolean {
    // Fast path without collectors -> no buffering
    // 如果没有订阅者，直接调用 tryEmitNoCollectorsLocked() 方法处理，它始终返回 true
    if (nCollectors == 0) return tryEmitNoCollectorsLocked(value) // always returns true
    // With collectors we'll have to buffer
    // cannot emit now if buffer is full & blocked by slow collectors
    // 判断当前是否已经达到容量的上限，如果达到上限触发 Overflow 策略
    if (bufferSize >= bufferCapacity && minCollectorIndex <= replayIndex) {
        when (onBufferOverflow) {
            // 直接返回 false
            BufferOverflow.SUSPEND -> return false // will suspend
            // 直接返回 true，并丢弃当前发送的消息
            BufferOverflow.DROP_LATEST -> return true // just drop incoming
            // 丢弃最旧的一条数据，和没有达到上限时的逻辑一样，只是后面的逻辑会多一个移除最旧的一条数据
            BufferOverflow.DROP_OLDEST -> {} // force enqueue & drop oldest instead
        }
    }
    // 数据插入 Buffer
    enqueueLocked(value)
    // 更新 bufferSize
    bufferSize++ // value was added to buffer
    // drop oldest from the buffer if it became more than bufferCapacity
    
    // 如果达到最大的 Capacity 丢弃最旧的一条数据，也就是 DROP_OLDEST 逻辑
    if (bufferSize > bufferCapacity) dropOldestLocked()
    // keep replaySize not larger that needed
    // 如果 replaySize 大于了 replay 的数量，需要将当前的 replayIndex 向后移动一位
    if (replaySize > replay) { // increment replayIndex by one
        updateBufferLocked(replayIndex + 1, minCollectorIndex, bufferEndIndex, queueEndIndex)
    }
    return true
}

上面代码稍微复杂了一点，这里我整理一下：

1. 如果订阅者的数量为 0，直接调用 tryEmitNoCollectorsLocked() 方法处理，而且它永远返回 true，也就是说当 ShreadFlowImpl 没有订阅者时，emit() 方法永远不会被挂起。如果有订阅者，继续执行后续流程。
2. 判断当前是否已经达到了 Buffer 的容量上限，如果达到了上限触发不同的 BufferOverflow 策略：

• SUSPEND: 直接返回 false，不执行后续逻辑了，也就是后续会进入挂起逻辑。
• DROP_LATEST: 直接返回 true，不执行后续逻辑了，后续也不会挂起，而是直接丢弃当前发送的数据。
• DROP_OLDEST: 继续执行后续逻辑，也就是和没有达到缓存上限时执行的流程一样。

3. 通过 enqueueLocked() 方法向 Buffer 中插入数据，同时将 bufferSize 的数量加 1.
4. 判断 bufferSize 是否已经达到最大的数量，如果达到了调用 dropOldestLocked() 方法丢弃最旧的一条数据，也就是 DROP_OLDTEST 的逻辑。
5. 如果当前的 replaySize 数量大于了设置的 replay，将 replayIndex 向后移动一位。
6. 最后返回 true，总的来说只有 SUSPEND 达到最大的 Buffer 上限后才会返回 false，其他情况全部返回 true。

我们再看看没有订阅者时调用的 tryEmitNoCollectorsLocked() 方法：

private fun tryEmitNoCollectorsLocked(value: T): Boolean {
    assert { nCollectors == 0 }
    // 没有设置 replay 直接返回，丢弃当前数据
    if (replay == 0) return true // no need to replay, just forget it now
    // 插入数据
    enqueueLocked(value) // enqueue to replayCache
    bufferSize++ // value was added to buffer
    // drop oldest from the buffer if it became more than replay
    // 当插入的数据大于 replay 后，丢弃最旧的一条数据.
    if (bufferSize > replay) dropOldestLocked()
    minCollectorIndex = head + bufferSize // a default value (max allowed)
    return true
}

上面的逻辑非常简单，也就是如果没有设置 replay 直接丢弃当前元素返回。如果有设置，直接插入到 Buffer 中，如果当前 Buffer 中的元素数量大于 replay，丢弃最旧的一条数据。

我们再来看看插入数据的逻辑 enqueueLocked() 方法的实现：

private fun enqueueLocked(item: Any?) {
    val curSize = totalSize
    val buffer = when (val curBuffer = buffer) {
        // Buffer 为空创建一个大小为 2 的 Buffer
        null -> growBuffer(null, 0, 2)
        // 如果 Buffer 的数量达到了上限，将原来的 Buffer 容量扩展至 2 倍
        else -> if (curSize >= curBuffer.size) growBuffer(curBuffer, curSize,curBuffer.size * 2) else curBuffer
    }
    // 将值添加到 Buffer 中。
    buffer.setBufferAt(head + curSize, item)
}

再看看 dropOldestLocked() 丢弃最旧一条数据的实现：

private fun dropOldestLocked() {
    // 将第一个元素设置为空
    buffer!!.setBufferAt(head, null)
    // buffer 数量减 1
    bufferSize--
    // head 向后移动一个位置
    val newHead = head + 1
    // 计算新的 replayIndex
    if (replayIndex < newHead) replayIndex = newHead
    // 如果有 collector 需要这个被移除的 index，将他们的 index 修改成新的 head.
    if (minCollectorIndex < newHead) correctCollectorIndexesOnDropOldest(newHead)
    assert { head == newHead } // since head = minOf(minCollectorIndex, replayIndex) it should have updated
}

我们再来看看 tryEmit() 方法中调用 findSlotsToResumeLocked() 去查找可以恢复的 collect() 方法的实现：

private fun findSlotsToResumeLocked(resumesIn: Array<Continuation<Unit>?>): Array<Continuation<Unit>?> {
    var resumes: Array<Continuation<Unit>?> = resumesIn
    var resumeCount = resumesIn.size
    // 遍历 Slot
    forEachSlotLocked loop@{ slot ->
        // Continuation 为空，跳过
        val cont = slot.cont ?: return@loop // only waiting slots
        // Slot 的 index 还不可用，跳过。
        if (tryPeekLocked(slot) < 0) return@loop // only slots that can peek a value
        if (resumeCount >= resumes.size) resumes = resumes.copyOf(maxOf(2, 2 * resumes.size))
        // 将 Continuation 添加到结果的数组中
        resumes[resumeCount++] = cont
        // 将 Slot 中的 Continuation 设置为空
        slot.cont = null // not waiting anymore
    }
    return resumes
}

tryEmit() 方法中的主要逻辑就分析完毕了，这里再总结一下主要逻辑：如果缓存没有达到上限，元素直接插入到 Buffer 中，然后去查找可以恢复的 collect() 方法，执行恢复操作，然后就结束了。如果达到了上限，DROP_LATEST 策略直接跳过当前元素的处理，DROP_OLDEST 处理逻辑和没有达到缓存上限类似，只是多了一个简单移除最旧消息的逻辑，如果是 SUSPEND，就会直接进入后续的挂起 emit() 方法的逻辑，也是后续我们要继续分析的代码。

直接看看挂起 emit() 方法的实现：

private suspend fun emitSuspend(value: T) = suspendCancellableCoroutine<Unit> sc@{ cont ->
    var resumes: Array<Continuation<Unit>?> = EMPTY_RESUMES
    val emitter = synchronized(this) lock@{
        // recheck buffer under lock again (make sure it is really full)
        // 再次检查是否可以发送数据
        if (tryEmitLocked(value)) {
            // 如果可以发送直接恢复当前 emit() 方法和查找可以恢复的 collect() 函数
            cont.resume(Unit)
            resumes = findSlotsToResumeLocked(resumes)
            return@lock null
        }
        // add suspended emitter to the buffer
        // 将当前的 Continuation 和 元素都封装在 Emitter 对象中
        Emitter(this, head + totalSize, value, cont).also {
            // 将 Emitter 添加到 Buffer 中
            enqueueLocked(it)
            // 更新 queueSize
            queueSize++ // added to queue of waiting emitters
            // synchronous shared flow might rendezvous with waiting emitter
            // 如果 Capacity 为空，需要手动去查询一次需要恢复的 collect() 函数
            if (bufferCapacity == 0) resumes = findSlotsToResumeLocked(resumes)
        }
    }
    // outside of the lock: register dispose on cancellation
    // 当前的协程取消后，通知 Emitter.
    emitter?.let { cont.disposeOnCancellation(it) }
    // outside of the lock: resume slots if needed
    // 恢复挂起的 collect() 方法
    for (r in resumes) r?.resume(Unit)
}

上面的代码都比较简单，又有注释，我就不再多说了。

StateFlow 工作原理

StateFlow 也是继承于 AbstractSharedFlow，它本身的代码逻辑也比 SharedFlow 简单挺多，如果看懂了 SharedFlow 的代码再来看 StateFlow 的代码会简单很多。

在开始之前需要简单介绍一下它的成员变量，它的成员变量就非常简单了：

private class StateFlowImpl<T>(
    initialState: Any // T | NULL
) : AbstractSharedFlow<StateFlowSlot>(), MutableStateFlow<T>, CancellableFlow<T>, FusibleFlow<T> {
    private val _state = atomic(initialState) // T | NULL
    private var sequence = 0 // serializes updates, value update is in process when sequence is odd
    // ...
}

每次通过 emit() 发送的元素它都会保存在 _state 变量中，_state 的默认值就是外部传入的默认状态，StateFlow 中发送的元素也就是不停变化的 _state。而 sequence 每次更新新的状态时就会加 2，它的处理逻辑非常有意思，更新前加 1，更新完成后再加 1，也就是如果当前没有正在更新状态那么它一定是偶数，如果有线程正在更新那么它一定是基数，后面我们会分析这部分代码。

还是老规矩开始前我先简单介绍一下 StateFlowImpl 的工作原理：

订阅也是通过 collect() 方法，订阅开始时会创建一个 FlowSlot 对象（和 SharedFlowImpl 中一样，它里面会存储当前 collect() 方法的 Continuation 对象，还有其他的一些状态。），然后开启一个死循环无限去读取 _state 中的状态，如果前一次读取的状态和当前这一次读取的状态不一样就表示有新的元素，这个时候会通过 FlowCollector#emit() 方法通知订阅者；如果前一次的状态和当前这一次读取的状态一样就表示没有新的元素，就需要把当前的 collect() 方法挂起，就会把对应的 Continuation 存放在 FlowSlot 中，等待下次状态更新后再唤醒，然后进入下次读取状态的逻辑。

发送数据是通过 emit() 方法，发送数据就是简单更新 _state 的状态，然后获取所有的能够恢复的 collect() 方法中对应的 FlowSlot，然后调用它们的 Continaution#resumeWith() 方法，去唤醒他们，然后促使 collect() 方法们进入下一次循环去读取新的状态。

订阅

我们直接看看 StateFlowImpl#collect() 方法的实现：

override suspend fun collect(collector: FlowCollector<T>): Nothing {
    // 创建一个 FlowSlot，它的实现类是 StateFlowSlot
    val slot = allocateSlot()
    try {
        if (collector is SubscribedFlowCollector) collector.onSubscription()
        val collectorJob = currentCoroutineContext()[Job]
        var oldState: Any? = null // previously emitted T!! | NULL (null -- nothing emitted yet)
        // The loop is arranged so that it starts delivering current value without waiting first
        // 死循环读取状态
        while (true) {
            // Here the coroutine could have waited for a while to be dispatched,
            // so we use the most recent state here to ensure the best possible conflation of stale values
            // 获取当前这一次循环的状态
            val newState = _state.value
            // always check for cancellation
            collectorJob?.ensureActive()
            // Conflate value emissions using equality
            // 如果上一次状态为空或者当前这一次的循环的状态和上一次的状态发生了改变，那么就把当前的状态通知给订阅者，然后进入下一次循环
            if (oldState == null || oldState != newState) {
                collector.emit(NULL.unbox(newState))
                oldState = newState
            }
            // Note: if awaitPending is cancelled, then it bails out of this loop and calls freeSlot
            // 如果当前没有正在发送数据
            if (!slot.takePending()) { // try fast-path without suspending first
                // 挂起当前的 collect() 函数
                slot.awaitPending() // only suspend for new values when needed
            }
        }
    } finally {
        freeSlot(slot)
    }
}

上面代码比较简单，我就不再多说，创建 FlowSlot 的方法 allocateSlot() 我们在讲 SharedFlow 的时候已经讲过了，其中 FlowSlot 的实现类是 StateFlowSlot，我们先看看它的初始化的方法 StateFlowSlot#allocateLocked() 方法：

override fun allocateLocked(flow: StateFlowImpl<*>): Boolean {
    // No need for atomic check & update here, since allocated happens under StateFlow lock
    if (_state.value != null) return false // not free
    _state.value = NONE // allocated
    return true
}

注意这个 _state 是 StateFlowSlot 中的，不要和 StateFlowImpl 中的 _state 搞混了，它没有初始化之前的值是空，初始化后是 NONE。

我这里需要解释一下 StateFlowSlot 中的 _state 中的各种状态：

• null: 未初始化。
• NONE: 已经初始化，没有在挂起状态，也没有正在发送数据。
• PENDING: emit() 方法正在更新新的状态
• 其他（Continuation）: 表示 collect() 已经被挂起，需要等待发送数据时唤醒。

看看 StateFlowSlot#takePending() 的实现：

fun takePending(): Boolean = _state.getAndSet(NONE)!!.let { state ->
    assert { state !is CancellableContinuationImpl<*> }
    return state === PENDING
}

将状态设置为 NONE，并判断上次的状态是不是 PENDING，如果是 PENDING 就表示有新的状态正在写入，然后 collect() 方法中会进入下一次循环，然后再次尝试读取状态。

然后再看看 StateFlowSlot#awaitPending() 是如何挂起 collect() 方法的：

@Suppress("UNCHECKED_CAST")
suspend fun awaitPending(): Unit = suspendCancellableCoroutine sc@ { cont ->
    assert { _state.value !is CancellableContinuationImpl<*> } // can be NONE or PENDING
    // 将状态设置为 Continuation
    if (_state.compareAndSet(NONE, cont)) return@sc // installed continuation, waiting for pending
    // CAS failed -- the only possible reason is that it is already in pending state now
    // 如果设置失败就表示当前有新的状态正在写入，所以直接唤醒当前的 Continuation，然后进入下一次循环。
    assert { _state.value === PENDING }
    cont.resume(Unit)
}

发送数据

private val _state = atomic(initialState) // T | NULL
private var sequence = 0 // serializes updates, value update is in process when sequence is odd

@Suppress("UNCHECKED_CAST")
public override var value: T
    get() = NULL.unbox(_state.value)
    set(value) { updateState(null, value ?: NULL) }

// ...
override suspend fun emit(value: T) {
    this.value = value
}
// ...

emit() 方法最终会触发 updateState() 方法，我们来看看它的实现(这个方法中就用了我上面说的通过 sequence 的奇偶来判断是否有新的状态在写入)：

private fun updateState(expectedState: Any?, newState: Any): Boolean {
    var curSequence: Int
    var curSlots: Array<StateFlowSlot?>? // benign race, we will not use it
    synchronized(this) {
        val oldState = _state.value
        if (expectedState != null && oldState != expectedState) return false // CAS support
        // 新旧状态必须不一致
        if (oldState == newState) return true // Don't do anything if value is not changing, but CAS -> true
        _state.value = newState
        curSequence = sequence
        // 通过与 1 做与操作来判断奇偶，结果为 0 表示偶数
        if (curSequence and 1 == 0) { // even sequence means quiescent state flow (no ongoing update)
            // 偶数就表示当前没有别的线程在写入新的状态，将状态加 1，这时就变成了基数。
            curSequence++ // make it odd
            sequence = curSequence
        } else {
            // 基数表示当前有别的线程正在写入新的状态，将状态加 2，这时还是基数，然后直接返回不在做后续的处理。
            // update is already in process, notify it, and return
            sequence = curSequence + 2 // change sequence to notify, keep it odd
            return true // updated
        }
        curSlots = slots // read current reference to collectors under lock
    }
    /*
       Fire value updates outside of the lock to avoid deadlocks with unconfined coroutines.
       Loop until we're done firing all the changes. This is a sort of simple flat combining that
       ensures sequential firing of concurrent updates and avoids the storm of collector resumes
       when updates happen concurrently from many threads.
     */
    while (true) {
        // Benign race on element read from array
        // 遍历 collect 的 FlowSlot，然后通过 `StateFlowSlot#makePending()` 方法唤醒 collect() 方法
        curSlots?.forEach {
            it?.makePending()
        }
        // check if the value was updated again while we were updating the old one
        synchronized(this) {
            // 判断是否有新的状态更新
            if (sequence == curSequence) { // nothing changed, we are done
                // 将 sequence 再加 1，当前又变成了偶数，表示当前的更新已经完成
                sequence = curSequence + 1 // make sequence even again
                return true // done, updated
            }
            // reread everything for the next loop under the lock
            
            // 前后状态不一致，表示别的线程也有在更新状态，进入下一次循环，再次尝试唤醒 collect() 方法。
            curSequence = sequence
            curSlots = slots
        }
    }
}

上面的代码可能有点绕，但是代码本身不复杂，我这里解释一下：

1. 更新状态时会判断当前的状态和当前的状态是否一致，如果不一致继续后续的更新状态逻辑。
2. 判断 sequence 的奇偶状态，如果是偶数表示当前没有别的线程在更新状态，然后将 sequence 加1，继续后续的更新状态的逻辑；如果是奇数表示当前有别的线程正在更新状态，直接将 sequence 加2，然后直接返回，不在继续后续的逻辑。
3. 遍历所有的 FlowSlot，通过 makePending() 方法唤醒对应的 collect() 方法，让其去读取新的状态。
4. 判断前后的 sequence 状态前后是否一致，如果一致就表示更新状态期间没有别的线程来更新状态，直接将 sequece 加 1，这时 sequece 就变成偶数了，然后直接返回方法；如果前后的 sequence 状态不一致，表示在这期间有别的线程来更新了新的状态，然后再次进入第三步，继续遍历 FlowSlot，直到前后的 sequence 一致。

我们再简单看看 StateFlowSlot#makePending() 方法是如何唤醒 collect() 方法的：

fun makePending() {
    _state.loop { state ->
        when {
            state == null -> return // this slot is free - skip it
            state === PENDING -> return // already pending, nothing to do
            state === NONE -> { // mark as pending
                if (_state.compareAndSet(state, PENDING)) return
            }
            else -> { // must be a suspend continuation state
                // we must still use CAS here since continuation may get cancelled and free the slot at any time
                if (_state.compareAndSet(state, NONE)) {
                    (state as CancellableContinuationImpl<Unit>).resume(Unit)
                    return
                }
            }
        }
    }
}

上面方法很简单，如果是 NONE 状态就直接修改成 PEDDING 状态；如果是 Continuation 状态直接调用其 resumeWith() 方法将其唤醒，然后将状态修改成 NONE。

最后

通过本篇文章希望大家能够深入理解 Kotlin 协程中的 SharedFlow 和 StateFlow 实现，如果没有看懂，推荐多看几次。