How Kotlin Coroutines Work Internally (Kotlin Standard Library)(Part-2.1)

This is the second part of the article series, “How Kotlin Coroutines Work Internally.” If you haven’t already, you can read the first part here. In this article, I will explore the implementation details of Kotlin coroutines within the Kotlin standard library.

If you want to explore the source code referenced in this article from the Kotlin repository, you can access it here.

CoroutineContext

CoroutineContext is essentially a map-like collection with unique keys, implemented using the composite design pattern. You can learn more about this here. The reason it's not implemented as a traditional map is to preserve type information, which would be lost in a standard map structure.

CoroutineContext Interface:-

interface CoroutineContext { 
 
 
    operator fun <E : Element> get(key: Key<E>): E? 
 
    fun <R> fold(initial: R, operation: (R, Element) -> R): R 
 
    operator fun plus(context: CoroutineContext): CoroutineContext { 
    } 
 
    fun minusKey(key: Key<*>): CoroutineContext 
 
    interface Key<E : Element> 
 
    interface Element : CoroutineContext { 
 
        val key: Key<*> 
 
        override fun <E : Element> get(key: Key<E>): E? { 
            return if (this.key == key) this as E else null 
        } 
 
        override fun <R> fold(initial: R, operation: (R, Element) -> R): R { 
            return operation(initial, this) 
        } 
 
        override fun minusKey(key: Key<*>): CoroutineContext { 
           return if(this.key == key) EmptyCoroutineContext else this 
        } 
 
    } 
}

As you can see above, the CoroutineContext interface defines four functions. I will discuss the plus function a bit later. For now, let's go through the other three functions one by one:

• get Function: This function takes a parameter named key of type Key, which is also defined within the CoroutineContext interface. Key is an empty interface, primarily used for referencing. The key parameter is a generic type that must be a subtype of the Element interface, also defined within CoroutineContext. The Element interface represents a single element within the context, which is how the CoroutineContext collection maintains type safety. The return type of the get function matches the generic type provided with the key parameter. It is implemented as an operator function, allowing you to access elements using the index access syntax.
• fold Function: This function allows us to iterate through all the values within the CoroutineContext collection. I will explain this in more detail a little bit later.
• minusKey Function: This function removes the element associated with the specified key, if it exists, and returns the remaining collection.
• plus Function: This function is responsible for combining multiple CoroutineContext instances into a single collection.

EmptyCoroutineContext

object EmptyCoroutineContext : CoroutineContext { 
   
    override fun <E : CoroutineContext.Element> get(key: CoroutineContext.Key<E>): E? = null 
    override fun <R> fold(initial: R, operation: (R, CoroutineContext.Element) -> R): R = initial 
    override fun plus(context: CoroutineContext): CoroutineContext = context 
    override fun minusKey(key: CoroutineContext.Key<*>): CoroutineContext = this 
}

EmptyCoroutineContext represents a collection that contains no elements, effectively an empty CoroutineContext. In its implementation, the get method returns null since there are no elements present. The minusKey function returns the EmptyCoroutineContext itself, as it remains an empty collection. The plus function returns another CoroutineContext since EmptyCoroutineContext adds no elements of its own. Lastly, the fold function returns the initial value, as there are no elements to iterate over.

Element

This interface is defined within the CoroutineContext interface and also inherits from CoroutineContext. It represents a CoroutineContext collection that contains only one element. An important detail to note is that each element also holds a reference to its key. Let's delve into the implementation details.

• get Function: In this collection, which contains only one element, this function checks if the provided key is equal to the key of that element. If they match, it returns the same element; otherwise, it returns null.
• minusKey Function: In this collection, which contains only one element, the function checks if the provided key matches the key of that element. If they are equal, it returns EmptyCoroutineContext, as there are no elements left. If they do not match, it returns the same context, indicating that it still does not contain any element associated with that key.
• fold Function: This function performs a single operation, as specified by the parameter, using the initial value, since the collection contains only one element.

The plus function is implemented in the CoroutineContext interface, which I will discuss in more detail shortly.

CombinedContext

Now the question arises: what happens if there are more than two elements in the CoroutineContext interface? The CombinedContext class is used to represent a CoroutineContext collection that contains multiple elements. Let's examine the implementation details of this class:

class CombinedContext( 
    val left: CoroutineContext, 
    val element: CoroutineContext.Element 
) : CoroutineContext { 
 
    override fun <E : CoroutineContext.Element> get(key: CoroutineContext.Key<E>): E? { 
        var cur = this //1 
        while (true) { 
            cur.element.get(key)?.let { return it } //2 
            val next = cur.left  //3 
            if (next is CombinedContext) { //4 
                cur = next //5 
            } else { 
                return next[key] //6 
            } 
        } 
    } 
 
    override fun <R> fold(initial: R, operation: (R, CoroutineContext.Element) -> R): R { 
        return operation(left.fold(initial, operation), element) //1 
    } 
 
    override fun minusKey(key: CoroutineContext.Key<*>): CoroutineContext { 
        element[key]?.let { return left } //1 
        val newLeft = left.minusKey(key) //2 
        return when { 
            newLeft === this -> this //3 
            newLeft === EmptyCoroutineContext -> element //4 
            else -> CombinedContext(newLeft, element) //5 
        } 
    } 
 
 
}

If you examine the constructor of CombinedContext, you'll notice it takes two parameters: one of type CoroutineContext (named left) and another of type CoroutineContext.Element (named element). The idea here is that the element has been added to the existing context (left). Now, let's dive into how the get, fold, and minusKey functions work in this context:

Finding Elements in context collection:

get Function: The left parameter is of type CoroutineContext, which can have multiple elements, a single element, or none at all, while the element parameter represents a single element. In the get function, our goal is to find the element whose key matches the key passed as a parameter.

If you look at the implementation (at mark 2), it first checks if the key matches the key of the element within the CombinedContext. If they match, it returns that element. The cur variable holds the next processable CombinedContext element.

If the key does not match the element's key, we need to search for it in the context (the left parameter). At mark 3, we store the left context in the next variable. If next is a CombinedContext, which means it has more than one element, the loop iterates over all elements in that context, updating the next variable accordingly.

If next is not a CombinedContext, it is a single element. In this case, the function checks if the key of that element matches the key passed as a parameter. If they match, it returns that element; otherwise, it returns null.

Removing elements from context collection:

minusKey Function: The left parameter is of type CoroutineContext, which can contain multiple elements, one element, or none, while the element is of type CoroutineContext.Element and represents a single element. Our goal is to remove the element from this collection whose key matches the key passed as a parameter.

At mark 1, we check if the key matches the key of the element within the CombinedContext. If they match, the element is removed, leaving only the left context, which we return.

If the key does not match, we call the minusKey function on left. For now, assume that this function attempts to remove the element with the specified key from left. I’ll explain this in more detail later.

After potentially removing that element (stored as newLeft), we use a when expression to check the state:

• At mark 3, we check whether newLeft is still equal to the original left. If so, this means there was no element with the specified key in left, so we return the original CombinedContext with the element intact.
• At mark 4, if newLeft is equal to EmptyCoroutineContext, it means that left initially contained only one element, which matched the key and was removed. Therefore, we return the element of the CombinedContext.
• At mark 5, we return a new CombinedContext with newLeft and the element of the original CombinedContext. This occurs because the element matching the key was removed from left, and newLeft now represents the collection without that element.

Now, let me explain the logic behind mark 2. You might have noticed that we didn’t loop through all the elements to find the one matching the key — so how does it work? Let me break it down.

Every time a CombinedContext is created, the minusKey function is overridden. At mark 2, we call the minusKey function on the left context. This creates a chain of calls.

Imagine we have 10 elements in the collection, and the element we are looking for is at the 5th position. What happens is that the minusKey function of the 10th element calls the minusKey function of the 9th element, and the 9th calls the 8th, and so on, until it reaches the 5th element, where it finds the match and removes it.

This chain of calls allows the function to work efficiently without having to manually loop through all elements, enabling a recursive search through the CombinedContext structure.

Fold Operation on Context Collection:

fold Function: The fold function also works by forming a chain, similar to the minusKey function.

At mark 1, the function first applies the provided operation to the left context. Then, it applies the same operation to the element. If you use this CombinedContext as the left context for creating another CombinedContext, the fold function is overridden again, allowing the operation to be applied to the new element.

This chaining mechanism ensures that the operation is eventually applied to all the elements within the CoroutineContext collection. The operation starts from the outermost context and proceeds inwards, applying the operation to each element in turn. This recursive application effectively traverses the entire context, ensuring that every element is processed.

Now, let’s move on to how CombinedContext is created by the CoroutineContext. To do this, we need to understand the working mechanism of the plus function defined in the CoroutineContext interface. The implementation is as follows:

operator fun plus(context: CoroutineContext): CoroutineContext { 
    return if (context == EmptyCoroutineContext) this else context.fold(this) { acc, element -> //1 
        val removed = acc.minusKey(key = element.key) //2 
        if (removed == EmptyCoroutineContext) element else { //3 
            val interceptor = removed[ContinuationInterceptor] //4 
            if (interceptor == null) CombinedContext(removed, element) else { 
                val left = removed.minusKey(ContinuationInterceptor) 
                CombinedContext(CombinedContext(left, element), interceptor) 
            } 
        } 
    } 
}

This function is responsible for combining multiple CoroutineContext elements into a single collection. Let's examine its implementation:

In the first line, we check if the context being added is EmptyCoroutineContext. If it is, we return the original context itself since the added context is empty and doesn't alter the collection.

If the context isn’t empty, we use the fold function to perform the addition operation for each element in the context. We start with the context itself as the initial value because we are adding all elements from the incoming context's collection to the existing context.

Step-by-Step Breakdown:

1. Check for Existing Keys:
   At line 2, the fold operation checks if the current accumulator (which represents the evolving context) already contains an element with the same key as the one being added. If it does, the old element needs to be removed before adding the new one. This is because CoroutineContext requires unique keys for each element.
2. Handle Empty Context:
   After removing the element, we check if we’re left with EmptyCoroutineContext. If this is the case, it means the original context only contained the element that was removed, so we return the new element (mark 3).
3. Handling Coroutine Interceptor:
   We then check if the accumulator contains a CoroutineInterceptor element. The CoroutineInterceptor is a special type of CoroutineContext.Element and plays a crucial role in managing the execution of coroutines. For performance reasons, we want this element to be quickly accessible, so we aim to keep it at the top of the context hierarchy.

• If the interceptor is not null (meaning it exists), we first remove it from the accumulator.
• We then create a new CombinedContext with the accumulator (minus the interceptor) and the new element.
• Finally, we create another CombinedContext using the previous combined context as the left context and the interceptor as the element. This ensures the interceptor remains on top, optimizing access performance.

If there is no CoroutineInterceptor, we simply return a CombinedContext with the updated accumulator and the new element.

How to use in Practice

Let’s take a look at an example to see how kotlinx.coroutines utilizes CoroutineContext in practice.

public abstract class AbstractCoroutineContextElement(public override val key: Key<*>) : Element

AbstractCoroutineContextElement is a class implemented in the Kotlin standard library that serves as a foundation for creating custom CoroutineContext elements. Library creators are encouraged to inherit from this class when defining their own coroutine context elements.

This class is straightforward; it inherits from CoroutineContext.Element and overrides the key property. By doing so, it provides a convenient way to establish a unique key for the custom element, ensuring that it integrates seamlessly into the CoroutineContext framework.

CoroutineName (CoroutineContext Example):-

public data class CoroutineName( 
    val name: String 
) : AbstractCoroutineContextElement(CoroutineName) { 
     
    public companion object Key : CoroutineContext.Key<CoroutineName> 
 
    override fun toString(): String = "CoroutineName($name)" 
}

In the implementation above, we have a straightforward CoroutineContext element designed to assign a name to each coroutine. Notice how the companion object is used as the key. This approach simplifies the process of accessing the element from the CoroutineContext collection.

By using a companion object as the key, we ensure that each instance of the coroutine name element is easily retrievable, facilitating efficient management and identification of coroutines within the context.

That concludes this article. CoroutineContext serves as the foundation behind coroutines, which is why I provided a detailed explanation of its components and functionality. In the next part of this series, we will continue exploring the Kotlin standard library for coroutines, as there is still much more to discuss. (I am enhancing this article further by incorporating some sophisticated visuals.)

If you have any questions, feel free to ask by replying! If you enjoyed the article, please give it a clap! Don’t forget to follow me on LinkedIn and Twitter for more updates!

You can read the third part of this series here.

You can read the fourth part of this series here.