Kotlin 协程 select（选择器)的源码分析

前言

多路复用（Multiplexing） 是一种技术，允许在同一通信介质或信道上传输多个信号或数据流。其目的是通过有效地共享有限的资源（如带宽、时间、物理设备等），提高通信效率，减少资源浪费。kotlin中的select和我们常见的还不太一样，这里的是一次性的操作。

Kotlin 中的 Select

协程的 select 是一种用于异步操作的选择器，它允许同时等待多个挂起函数的结果，并在其中一个完成时执行相应的操作。

select

从注释中，可以看到，支持Job, Deferred, Channel。

kotlin 复制代码

public suspend inline fun <R> select(crossinline builder: SelectBuilder<R>.() -> Unit): R {
    contract {
        callsInPlace(builder, InvocationKind.EXACTLY_ONCE)
    }
    return suspendCoroutineUninterceptedOrReturn { uCont ->
        val scope = SelectBuilderImpl(uCont)
        try {
            builder(scope)
        } catch (e: Throwable) {
            scope.handleBuilderException(e)
        }
        scope.getResult()
    }
}

invoke

传入的参数，是一个SelectBuidler接口的扩展函数。如下：

kotlin 复制代码

// SelectClause0, SelectClause1， SelectClause2的区别是传入的参数不同
public interface SelectClause0 {
    /**
     * Registers this clause with the specified [select] instance and [block] of code.
     * @suppress **This is unstable API and it is subject to change.**
     */
    @InternalCoroutinesApi
    public fun <R> registerSelectClause0(select: SelectInstance<R>, block: suspend () -> R)
}

public interface SelectBuilder<in R> {
    /**
     * Registers a clause in this [select] expression without additional parameters that does not select any value.
     */
    public operator fun SelectClause0.invoke(block: suspend () -> R)

    /**
     * Registers clause in this [select] expression without additional parameters that selects value of type [Q].
     */
    public operator fun <Q> SelectClause1<Q>.invoke(block: suspend (Q) -> R)

    public operator fun <P, Q> SelectClause2<P, Q>.invoke(param: P, block: suspend (Q) -> R)

    public operator fun <P, Q> SelectClause2<P?, Q>.invoke(block: suspend (Q) -> R): Unit = invoke(null, block)

    public fun onTimeout(timeMillis: Long, block: suspend () -> R)
}

这个地方使用了 operator 重载运算符,重载了 invoke 方法，所以我们可以像调用函数一下，直接调用这个对象，即 SelectClause0({}) 或 SelectClause0{}。

所以如果我们也想让我们自定义的类有这样的功能，也可以重载这个 invoke 方法。（猜测：kotlin中默认能直接用invoke方法的那些对象，应该都是重载了的。）

由上面的代码我们就可以知道，实际上我们调用的是实现了SelectBuilder接口的SelectBuilderImpl的拓展函数。下面看看 invoke 的具体实现：

kotlin 复制代码

override fun SelectClause0.invoke(block: suspend () -> R) {
    registerSelectClause0(this@SelectBuilderImpl, block)
}

override fun <Q> SelectClause1<Q>.invoke(block: suspend (Q) -> R) {
    registerSelectClause1(this@SelectBuilderImpl, block)
}

override fun <P, Q> SelectClause2<P, Q>.invoke(param: P, block: suspend (Q) -> R) {
    registerSelectClause2(this@SelectBuilderImpl, param, block)
}

实际上调用的是对应 Job, Deferred, Channel 中具体实现的代码。

registerSelectClauseX

kotlin 复制代码

//==================================================================================

    public final override val onJoin: SelectClause0
            get() = this
    // registerSelectJoin
    public final override fun <R> registerSelectClause0(select: SelectInstance<R>, block: suspend () -> R) {
        // fast-path -- check state and select/return if needed
        loopOnState { state ->
            if (select.isSelected) return
            if (state !is Incomplete) {
                // already complete -- select result
                if (select.trySelect()) {
                    block.startCoroutineUnintercepted(select.completion)
                }
                return
            }
            if (startInternal(state) == 0) {
                // slow-path -- register waiter for completion
                select.disposeOnSelect(invokeOnCompletion(handler = SelectJoinOnCompletion(select, block).asHandler))
                return
            }
        }
    }
//==================================================================================
 
 // registerSelectAwaitInternal
    @Suppress("UNCHECKED_CAST")
    internal fun <T, R> registerSelectClause1Internal(select: SelectInstance<R>, block: suspend (T) -> R) {
        // fast-path -- check state and select/return if needed
        loopOnState { state ->
            if (select.isSelected) return
            if (state !is Incomplete) {
                // already complete -- select result
                if (select.trySelect()) {
                    if (state is CompletedExceptionally) {
                        select.resumeSelectWithException(state.cause)
                    }
                    else {
                        block.startCoroutineUnintercepted(state.unboxState() as T, select.completion)
                    }
                }
                return
            }
            if (startInternal(state) == 0) {
                // slow-path -- register waiter for completion
                select.disposeOnSelect(invokeOnCompletion(handler = SelectAwaitOnCompletion(select, block).asHandler))
                return
            }
        }
    }
    
//==================================================================================
    
    private fun <R> registerSelectSend(select: SelectInstance<R>, element: E, block: suspend (SendChannel<E>) -> R) {
        while (true) {
            if (select.isSelected) return
            if (isFullImpl) {
                val node = SendSelect(element, this, select, block)
                val enqueueResult = enqueueSend(node)
                when {
                    enqueueResult == null -> { // enqueued successfully
                        select.disposeOnSelect(node)
                        return
                    }
                    enqueueResult is Closed<*> -> throw recoverStackTrace(helpCloseAndGetSendException(element, enqueueResult))
                    enqueueResult === ENQUEUE_FAILED -> {} // try to offer
                    enqueueResult is Receive<*> -> {} // try to offer
                    else -> error("enqueueSend returned $enqueueResult ")
                }
            }
            // hm... receiver is waiting or buffer is not full. try to offer
            val offerResult = offerSelectInternal(element, select)
            when {
                offerResult === ALREADY_SELECTED -> return
                offerResult === OFFER_FAILED -> {} // retry
                offerResult === RETRY_ATOMIC -> {} // retry
                offerResult === OFFER_SUCCESS -> {
                    block.startCoroutineUnintercepted(receiver = this, completion = select.completion)
                    return
                }
                offerResult is Closed<*> -> throw recoverStackTrace(helpCloseAndGetSendException(element, offerResult))
                else -> error("offerSelectInternal returned $offerResult")
            }
        }
    }

    private inline fun loopOnState(block: (Any?) -> Unit): Nothing {
        while (true) {
            block(state)
        }
    }

从上述代码，可以看到，大致的逻辑还是差不多的。在循环中去判断状态，直到协程执行结束后调用select中的resume方法，返回结果给select

总结

kotlin 中的 select 中主要是使用到了续体风格传递，将结果通过Result传递给外层的续体。当多个协程要唤醒同一续体时，会通过 CAS 去判断续体的状态。所以就能达到使用少资源去监听多任务的目的，当有返回值时，就直接返回做早的结果。

主要知识：

invoke的重载
续体传递、状态机（本文解析未涉及到）
协程调度（本文解析未涉及到）