Kotlin协程详解——withContext函数

一、原因

CoroutineDispatcher虽然能够提供线程的切换，但这只是单方向的，因为它没有提供线程的恢复。

试想一下，我们有个网络请求，我们通过CoroutineDispatcher将线程切换到Dispatchers.IO，当拿到请求成功的数据之后，所在的线程还是IO线程，这样并不能有利于我们UI操作。所以为了解决这个问题kotlin提供了withContext，它不仅能够接受CoroutineDispatcher来帮助我们切换线程，同时在执行完毕之后还会帮助我们将之前切换掉的线程进恢复，保证协程运行的连贯性。这也是为什么官方推荐使用withContext进行协程线程的切换的原因。

二、源码分析

/**

使用给定的协程上下文调用指定的挂起代码块，挂起直到它执行完成，并返回结果。
该代码块的最终上下文是通过将当前协程上下文与指定的上下文合并得到的
（使用 coroutineContext + context，参见 CoroutineContext.plus）。
这个挂起函数是可取消的。它会立即检查合并后的上下文是否处于活动状态，
如果不处于活动状态，则抛出 CancellationException。

当 withContext 的上下文参数提供的 CoroutineDispatcher 与当前的不同时，
必然需要进行额外的调度：代码块不能立即执行，需要调度到传递的 
CoroutineDispatcher 上执行，然后当代码块执行完成后，执行需要切换回原始的分发器。

请注意，withContext 调用的结果被以可取消的方式调度回原始上下文，
并具有即时取消保证，这意味着如果调用 withContext 时的原始 coroutineContext 
在其分发器开始执行代码之前被取消，它会丢弃 withContext 的结果并抛出 CancellationException。

上述取消行为仅在分发器发生变化时启用。例如，当使用 withContext(NonCancellable) { ... } 时，
分发器没有变化，因此在这个 withContext 代码块内部进入或退出时都不会被取消。

*/

public suspend fun <T> withContext(
    context: CoroutineContext,
    block: suspend CoroutineScope.() -> T
): T {
    contract {
        callsInPlace(block, InvocationKind.EXACTLY_ONCE)
    }
    return suspendCoroutineUninterceptedOrReturn sc@ { uCont ->
        //创建新的协程上下文
        val oldContext = uCont.context
        // Copy CopyableThreadContextElement if necessary
        val newContext = oldContext.newCoroutineContext(context)
        // always check for cancellation of new context
        newContext.ensureActive()
        // FAST PATH #1 -- new context is the same as the old one
        if (newContext === oldContext) {
            val coroutine = ScopeCoroutine(newContext, uCont)
            return@sc coroutine.startUndispatchedOrReturn(coroutine, block)
        }
        // FAST PATH #2 -- the new dispatcher is the same as the old one (something else changed)
        // `equals` is used by design (see equals implementation is wrapper context like ExecutorCoroutineDispatcher)
        if (newContext[ContinuationInterceptor] == oldContext[ContinuationInterceptor]) {
            val coroutine = UndispatchedCoroutine(newContext, uCont)
            // There are changes in the context, so this thread needs to be updated
            withCoroutineContext(newContext, null) {
                return@sc coroutine.startUndispatchedOrReturn(coroutine, block)
            }
        }
        // 使用新的调度器，覆盖外层
        val coroutine = DispatchedCoroutine(newContext, uCont)
        block.startCoroutineCancellable(coroutine, coroutine)
        coroutine.getResult()
    }
}

internal class DispatchedCoroutine<in T>(
    context: CoroutineContext,
    uCont: Continuation<T>
) : ScopeCoroutine<T>(context, uCont) {

    //在complete时会回调
    override fun afterCompletion(state: Any?) {
        // Call afterResume from afterCompletion and not vice-versa, because stack-size is more
        // important for afterResume implementation
        afterResume(state)
    }

    override fun afterResume(state: Any?) {
        //uCont为父协程，context仍是老版context，因此可以切换回原来的线程
        if (tryResume()) return // completed before getResult invocation -- bail out
        // Resume in a cancellable way because we have to switch back to the original dispatcher
        uCont.intercepted().resumeCancellableWith(recoverResult(state, uCont))
    }

}

1. 对于withContext，传入的context会覆盖外层的拦截器并生成一个newContext，因此可以实现线程切换。
2. DispatchedCoroutine作为complete传入协程体的创建函数中，因此协程体执行完成后会回调到afterCompletion中。
3. DispatchedCoroutine中传入的uCont是父协程，它的拦截器仍是外层的拦截器，因此会切换回原来的线程中。

而withContext的线程恢复原理是它内部生成了一个DispatchedCoroutine，保存切换线程时的CoroutineContext与切换之前的Continuation，最后在onCompletionInternal进行恢复。

internal override fun onCompletionInternal(state: Any?, mode: Int, suppressed: Boolean) {
    if (state is CompletedExceptionally) {
        val exception = if (mode == MODE_IGNORE) state.cause else recoverStackTrace(state.cause, uCont)
        uCont.resumeUninterceptedWithExceptionMode(exception, mode)
    } else {
        uCont.resumeUninterceptedMode(state as T, mode)
    }
}

这个uCont就是切换线程之前的Continuation。

推荐文章

https://zhuanlan.zhihu.com/p/301494587