Kotlin协程如何高效实现任务执行后切换指定CoroutineContext

需求背景 #

我们需要实现一个suspend函数，满足以下要求：

• 在指定的on上下文（通常是阻塞IO线程池）执行传入的阻塞函数f
• 函数执行完成后，强制在continueOn上下文（通常是非阻塞异步线程池）恢复协程执行
• 支持从包括Unconfined在内的任意协程上下文调用
• 尽可能减少线程切换次数，最优情况仅需2次切换，避免现有嵌套withContext实现最多4次切换的性能损耗

最优实现代码 #

import kotlinx.coroutines.Job
import kotlinx.coroutines.intrinsics.intercepted
import kotlinx.coroutines.intrinsics.suspendCoroutineUninterceptedOrReturn
import kotlin.coroutines.ContinuationInterceptor
import kotlin.coroutines.intrinsics.COROUTINE_SUSPENDED
import kotlin.coroutines.resumeWithException

suspend fun <A> eval(f: () -> A, on: CoroutineContext, continueOn: CoroutineContext): A {
    // 继承当前协程的Job，保证取消链路正常
    val currentJob = coroutineContext[Job]
    // 给目标上下文绑定当前Job，支持取消传播
    val targetOnContext = on + currentJob
    val targetContinueContext = continueOn + currentJob

    return suspendCoroutineUninterceptedOrReturn { cont ->
        // 先把续体拦截为continueOn上下文，后续恢复直接走该上下文的调度器
        val resumeContinuation = cont.intercepted(targetContinueContext)

        // 绑定取消逻辑：父协程取消时直接取消续体
        currentJob?.invokeOnCompletion { cause ->
            cause?.let { resumeContinuation.resumeWithException(it) }
        }

        // 直接调度到on上下文执行阻塞函数
        targetOnContext[ContinuationInterceptor]!!.dispatch(targetOnContext, Runnable runnable@{
            // 执行前检查是否已取消，避免无用执行
            if (currentJob?.isActive == false) return@runnable
            val result = runCatching(f)
            // 直接通过预拦截的续体恢复，自动切到continueOn上下文
            resumeContinuation.resumeWith(result)
        })

        COROUTINE_SUSPENDED
    }
}

实现原理说明 #

1. 上下文切换优化
   跳过了嵌套withContext实现中两次不必要的上下文恢复步骤：无需先切换到continueOn再切到on，也无需执行完on后先切回continueOn再切回调用方上下文，仅保留两次必要切换：

   • 第一次：从调用方上下文切换到on上下文执行阻塞函数f
   • 第二次：f执行完成后切换到continueOn上下文恢复协程

2. Unconfined上下文兼容
   我们提前将续体拦截为continueOn上下文，不管调用方是什么上下文，最终协程恢复一定会走continueOn的调度器，不会出现Unconfined调用后停留在on阻塞线程的问题。

3. 取消兼容性
   实现中继承了当前协程的Job，支持取消传播：父协程取消时会自动终止f的执行逻辑，不会出现任务泄漏。

常见疑问解答 #

• 为什么不用嵌套withContext实现？
  嵌套实现withContext(continueOn) { withContext(on) { f() } }最多会产生4次线程切换：调用方→continueOn→on→continueOn→调用方，额外的两次切换完全无意义，在高并发场景下会产生不必要的性能损耗。
• 为什么不直接用withContext(on) { f() }？
  如果调用方是Unconfined上下文，withContext执行完成后会直接在on的线程恢复执行，不符合我们要求强制回到continueOn上下文的需求。
• 为什么不手动写两次withContext？
  手动写withContext(on) { f() }.also { withContext(continueOn) {} }的写法不仅使用成本高，依然存在Unconfined上下文下第二次withContext执行前停留在on线程的问题，封装到函数中也无法避免额外的切换开销。

内容的提问来源于stack exchange，提问作者Sergey Alaev