Kotlin协程：嵌套runBlocking场景下保留协程上下文的方案

问题背景 #

作为Kotlin协程新手，在处理runBlocking与协程上下文交互时遇到了上下文丢失的场景。首先定义了一个自定义上下文元素：

class ExampleContext(val s: String) : AbstractCoroutineContextElement(Key) {
    companion object Key : CoroutineContext.Key<ExampleContext>
}

符合预期的行为 #

以下场景中上下文能正常传递：

runBlocking(ExampleContext("foo")) {
    println(coroutineContext[ExampleContext.Key]?.s) // 输出"foo"
}

runBlocking(ExampleContext("foo")) {
    launch {
        println(coroutineContext[ExampleContext.Key]?.s) // 输出"foo"
    }
}

runBlocking(ExampleContext("foo")) {
    launch(ExampleContext("bar")) {
        println(coroutineContext[ExampleContext.Key]?.s) // 输出"bar"
    }
}

上下文丢失的异常场景 #

嵌套无参runBlocking时，自定义上下文会丢失：

runBlocking(ExampleContext("foo")) {
    runBlocking {
        println(coroutineContext[ExampleContext.Key]?.s) // 输出null
    }
}

库开发场景 #

开发的库需要为第三方代码提供自定义上下文（类似拦截器），初始伪代码：

class MyLibrary(otherPeoplesLogic: OtherPeoplesBusinessLogic) {

    fun <IN, OUT> execute(input: IN): OUT {
        ... 执行库逻辑，添加自定义上下文元素 ...
        try {
            return otherPeoplesLogic.execute(input)
        } finally {
            ... 执行库的清理操作 ...
        }
    }
}

为支持协程修改后的代码：

class MyLibrary(otherPeoplesLogic: OtherPeoplesBusinessLogic) {

    fun <IN, OUT> execute(input: IN): OUT {
        ... 执行库逻辑 ...
        runBlocking(myCustomContext) {
            try {
                return otherPeoplesLogic.execute(input)
            } finally {
                ... 执行库的清理操作 ...
            }
        }
    }
}

核心问题：除约束用户不嵌套runBlocking外，是否有办法防止自定义上下文丢失？

解决方案 #

1. 提供自定义runBlocking扩展函数 #

在库中封装一个替代标准无参runBlocking的扩展，强制传递当前协程上下文：

fun <T> CoroutineScope.runBlockingWithContext(block: suspend CoroutineScope.() -> T): T {
    return kotlinx.coroutines.runBlocking(coroutineContext, block)
}

在库文档中明确告知用户使用该扩展替代标准无参runBlocking，这样嵌套调用时会自动继承当前上下文。

2. 用ThreadLocal做兜底传递 #

如果无法约束用户使用自定义扩展，可以将上下文存储在ThreadLocal中，作为协程上下文的兜底来源：

private val exampleContextThreadLocal = ThreadLocal<ExampleContext?>()

// 库启动runBlocking时设置ThreadLocal
runBlocking(myCustomContext) {
    exampleContextThreadLocal.set(myCustomContext[ExampleContext.Key])
    try {
        return otherPeoplesLogic.execute(input)
    } finally {
        exampleContextThreadLocal.remove()
    }
}

// 修改上下文元素，优先从ThreadLocal读取
class ExampleContext(val s: String) : AbstractCoroutineContextElement(Key) {
    companion object Key : CoroutineContext.Key<ExampleContext> {
        operator fun invoke(): ExampleContext? {
            return exampleContextThreadLocal.get()
        }
    }
}

注意：跨线程调度（如Dispatchers.IO）时需要额外处理ThreadLocal的传递，可结合ThreadContextElement实现自动同步。

3. 改用CoroutineScope作为入口（推荐长期方案） #

如果库的架构允许，让第三方代码基于CoroutineScope编写逻辑，而非直接调用runBlocking：

class MyLibrary(otherPeoplesLogic: OtherPeoplesBusinessLogic) {

    fun <IN, OUT> execute(input: IN, scope: CoroutineScope): OUT {
        ... 执行库逻辑 ...
        val scopedContext = scope.coroutineContext + myCustomContext
        return runBlocking {
            coroutineScope {
                async(scopedContext) {
                    otherPeoplesLogic.execute(input)
                }.await()
            }
        }
    }
}

这种方式通过CoroutineScope层级传递上下文，从根源上避免了嵌套runBlocking的上下文丢失问题。

结论 #

完全禁止用户嵌套runBlocking不现实，最可靠的方案是提供自定义runBlocking扩展并明确规范，同时结合ThreadLocal作为兜底。如果库的设计允许，改用CoroutineScope作为入口是更优雅的长期解决方案。

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