Under the Hood: CoroutineContext

In the world of Kotlin Coroutines, the CoroutineContext is the silent engine that holds everything together. Whether it's the Dispatcher, a Job, or a CoroutineName, they all live within this context. But how does it manage to store different types of data while remaining type-safe and easy to manipulate?

public interface CoroutineContext {
    /**
     * Returns the element with the given [key] from this context or `null`.
     */
    public operator fun <E : Element> get(key: Key<E>): E?

    /**
     * Accumulates entries of this context starting with [initial] value and applying [operation]
     * from left to right to current accumulator value and each element of this context.
     */
    public fun <R> fold(initial: R, operation: (R, Element) -> R): R

    /**
     * Returns a context containing elements from this context and elements from  other [context].
     * The elements from this context with the same key as in the other one are dropped.
     */
    public operator fun plus(context: CoroutineContext): CoroutineContext = ...

    /**
     * Returns a context containing elements from this context, but without an element with
     * the specified [key].
     */
    public fun minusKey(key: Key<*>): CoroutineContext

    /**
     * Key for the elements of [CoroutineContext]. [E] is a type of element with this key.
     */
    public interface Key<E : Element>

    /**
     * An element of the [CoroutineContext]. An element of the coroutine context is a singleton context by itself.
     */
    public interface Element : CoroutineContext {
        /**
         * A key of this coroutine context element.
         */
        public val key: Key<*>
        public override operator fun <E : Element> get(key: Key<E>): E? = ...
        public override fun <R> fold(initial: R, operation: (R, Element) -> R): R = ...
        public override fun minusKey(key: Key<*>): CoroutineContext = ...
    }
}

The CoroutineContext interface represents an element or a collection of elements; it is designed to be a Heterogeneous Container.

get: with this operator, we can get a particular CoroutineContext from the CoroutineContext container.

plus: with this operator we can add any CoroutineContext with another CoroutineContext to get a CoroutineContext as a container.

fold and minus are self-explanatory.

Composite Design pattern

To achieve the design of a Heterogeneous Container, CoroutineContext uses composite design pattern. A composite design creates a tree-like structure where an object contains a collection of objects of same type so as to treat individual objects and the composition of objects uniformly. Since it is a tree structure, it will have a leaf object.

• An Element is a leaf (a single context like Dispatchers.IO).

• A CombinedContext is the composite that holds multiple elements together. (Will discuss more on this later)

Key vs. Element

The Element

An Element is actually a CoroutineContext itself, but one that contains exactly one item.

The Key

The Key is the identity of the element.
How does Kotlin guarantee that context[Job] returns a Job? without manual casting?

The secret lies in the generic definition of the Key:

public interface Key<E : Element>

When an element (like Job) is defined, it implements Element and associates itself with a Key that is tied to its own type: (just for illustration, actual implementation is different)

class Job : CoroutineContext.Element {
    companion object Key : CoroutineContext.Key<Job>
    override val key: CoroutineContext.Key<*> = Key
}

This helps in

• Type Association: The compiler knows that Key<Job> can only ever be associated with an instance of Job.

• Polymorphic Dispatch: When you call get(Key), the signature is get(key: Key<E>): E?. Because the key itself carries the type information E, the return type is automatically inferred.

Notice that the declation of Key inside Element has a wildcard Key<*>. The reason is simple: we cannot declare variable with Key<Element>, even before creating the Element interface.

Combined Context: The final piece that completes the composite pattern

CombinedContext holds a CoroutineContext and an Element. It forms a left-skewed binary tree as shown below.

Using this binary tree structure, CombinedContext stores a collection of CoroutineContext.
CoroutineContext uses CombinedContext to add new CoroutineContext in the collection.

Notice the iterative nature of get, size, containsAll methods and the recursive nature of fold, minusKey as we would expect in the implementation of a binary tree.

internal class CombinedContext(
    private val left: CoroutineContext,
    private val element: Element
) : CoroutineContext, Serializable {

    override fun <E : Element> get(key: Key<E>): E? {
        var cur = this
        while (true) {
            cur.element[key]?.let { return it }
            val next = cur.left
            if (next is CombinedContext) {
                cur = next
            } else {
                return next[key]
            }
        }
    }

    public override fun <R> fold(initial: R, operation: (R, Element) -> R): R =
        operation(left.fold(initial, operation), element)

    public override fun minusKey(key: Key<*>): CoroutineContext {
        element[key]?.let { return left }
        val newLeft = left.minusKey(key)
        return when {
            newLeft === left -> this
            newLeft === EmptyCoroutineContext -> element
            else -> CombinedContext(newLeft, element)
        }
    }

    private fun size(): Int {
        var cur = this
        var size = 2
        while (true) {
            cur = cur.left as? CombinedContext ?: return size
            size++
        }
    }

    private fun contains(element: Element): Boolean =
        get(element.key) == element

    private fun containsAll(context: CombinedContext): Boolean {
        var cur = context
        while (true) {
            if (!contains(cur.element)) return false
            val next = cur.left
            if (next is CombinedContext) {
                cur = next
            } else {
                return contains(next as Element)
            }
        }
    }

    override fun equals(other: Any?): Boolean =
        this === other || other is CombinedContext && other.size() == size() && other.containsAll(this)

    override fun hashCode(): Int = left.hashCode() + element.hashCode()

    override fun toString(): String =
        "[" + fold("") { acc, element ->
            if (acc.isEmpty()) element.toString() else "$acc, $element"
        } + "]"

    ...
}

CoroutineContext plus method

plus method add all the context from the current CoroutineContext and the provided CoroutineContext, in case there is a conflict in key, it retains the other CoroutineContext key but drop the current. That means, new context in the right will replace the context in the left, in case of conflict.

The logic uses context.fold(this). This means it iterates over every element in the new context and tries to add them one by one to the accumulator (the current context).

    /**
     * Returns a context containing elements from this context and elements from  other [context].
     * The elements from this context with the same key as in the other one are dropped.
     */
    public operator fun plus(context: CoroutineContext): CoroutineContext =
        if (context === EmptyCoroutineContext) this else // fast path -- avoid lambda creation
            context.fold(this) { acc, element ->
                val removed = acc.minusKey(element.key)
                if (removed === EmptyCoroutineContext) element else {
                    // make sure interceptor is always last in the context (and thus is fast to get when present)
                    val interceptor = removed[ContinuationInterceptor]
                    if (interceptor == null) CombinedContext(removed, element) else {
                        val left = removed.minusKey(ContinuationInterceptor)
                        if (left === EmptyCoroutineContext) CombinedContext(element, interceptor) else
                            CombinedContext(CombinedContext(left, element), interceptor)
                    }
                }
            }

Will cover ContinuationInterceptor in a different blog, here it is ensured that ContinuationInterceptor is always kept on the left as an optimization since it is required frequently in suspend function flow.

Conclusion

CoroutineContext is designed to hold different kinds of context, and it achieves this by using composite design pattern. To ensure Type Safety, it utilizes "Types as Keys," allowing the compiler to guarantee that you always receive the correct object back without risky casting. It used left-skewed binary tree to hold contexts using CombinedContext.