ThreadLocal/CoroutineContext bridge gaps

Question

I want to maintain an object at the thread level or the coroutine level depending on the type of work the application is doing across different thread/coroutines. Is there a way to achieve this?

let's say for simplicity's sake, I can write a Spring Boot Application where many things are happening based on threads and only certain parts of code use coroutines to leverage their benefits. How do I maintain state based on the current execution? is there a way to do this?

Answer 1

Maybe my answer is a little bit late, but you can do it by either

1. ContinuationInterceptor
2. 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