Maybe my answer is a little bit late, but you can do it by either
ContinuationInterceptor
CopyableThreadContextElement
More elaborated answer with examples is provided here.
Below I will show an example for ContinuationInterceptor
copied from those answers.
class WrappedDispatcher(
private val dispatcher: ContinuationInterceptor,
private var savedCounter: Int = counterThreadLocal.get() ?: 0
) : AbstractCoroutineContextElement(ContinuationInterceptor), ContinuationInterceptor {
override fun <T> interceptContinuation(continuation: Continuation<T>): Continuation<T> =
dispatcher.interceptContinuation(ContinuationWrapper(continuation))
private inner class ContinuationWrapper<T>(val base: Continuation<T>) : Continuation<T> by base {
override fun resumeWith(result: Result<T>) {
counterThreadLocal.set(savedCounter)
try {
base.resumeWith(result)
} finally {
savedCounter = counterThreadLocal.get()
}
}
}
}
and usage
val counterThreadLocal: ThreadLocal<Int> = ThreadLocal()
fun showCounter() {
println("-------------------------------------------------")
println("Thread: ${Thread.currentThread().name}\n Counter value: ${counterThreadLocal.get()}")
}
fun main() {
runBlocking(WrappedDispatcher(Dispatchers.IO)) {
showCounter()
counterThreadLocal.set(2)
delay(100)
showCounter()
counterThreadLocal.set(3)
withContext(WrappedDispatcher(Dispatchers.Default)) {
println("__________NESTED START___________")
counterThreadLocal.set(4)
showCounter()
println("__________NESTED END_____________")
}
delay(100)
showCounter()
}
}
output will be
-------------------------------------------------
Thread: DefaultDispatcher-worker-1
Counter value: 0
-------------------------------------------------
Thread: DefaultDispatcher-worker-1
Counter value: 2
__________NESTED START___________
-------------------------------------------------
Thread: DefaultDispatcher-worker-3
Counter value: 4
__________NESTED END_____________
-------------------------------------------------
Thread: DefaultDispatcher-worker-3
Counter value: 3