Building Kotlin Coroutine Framework from Scratch: Part 2 -Reinventing Dispatchers

Writing a basic Dispatcher and implementing withContext

In the previous article, we introduced a basic launch function to start a new coroutine, but it had some drawbacks that we’ll discuss now

Building KotlinCoroutine Framework from Scratch: Part 1 - Coroutine Hello World

By default, our launch function executes on the caller (main) thread, which may not always be what we want. To address this, we introduce a class called Dispatcher that allows us to specify the desired thread for starting the coroutine

Now we have an ability to specify a dispatcher, but we still encounter a crash in our code after the download is complete

android.view.ViewRoot$CalledFromWrongThreadException: Only the original thread 
that created a view hierarchy can touch its views

The coroutine resumes on the background thread after the download is complete and attempts to update the TextView, resulting in a crash because only the main thread is allowed to interact with UI elements

In order to ensure consistency and predictability, it is important to execute all statements within the launch block on the same Dispatcher, regardless of the thread on which the continuation is resumed

One way to accomplish this is by passing the dispatcher as a parameter to all suspend functions and resume the continuation correctly

Although the code functions as intended, it’s a pretty bad solution when better alternatives exist

Leveraging CoroutineContext

Previously, we observed that suspend functions maintain a consistent CoroutineContext throughout their execution. This is achieved through the compiler-generated Continuation classes which overrides the context with the value from the previous continuation. As a result, the suspend call chain shares the same CoroutineContext throughout its execution

We can utilise this context property to pass our Dispatcher

The Dispatcher now implements the CoroutineContext.Element interface, which, in turn, implements the CoroutineContext interface. Additionally, we override the necessary key property with a CoroutineContext.Key object, which we will discuss further in detail later on

The Dispatcher can now be passed as aCoroutineContext

With these changes, we’ve successfully addressed all issues in our launch function

Implementing withContext

Let’s say we have a situation where we want to run a suspending lambda on a different dispatcher and wait for it to complete before proceeding with our code

Unfortunately, we can’t use the suspendCoroutineUninterceptedOrReturn function since it only works with non-suspending lambdas. As an alternative, we attempt to launch the lambda within a nested launch function

We notice that the nested launch function returns immediately, and “Download Complete” is printed in parallel with the execution of the nested launch function

This is because we have designed the launch function to start executing the lambda in a fire-and-forget manner, starting it in parallel with the rest of the code and operating independently of its parent coroutine

To address this issue, we introduce a new function called withContext which takes a Dispatcher and a suspending lambda that represents the work to be performed

The implementation is straightforward, we suspend the current coroutine using suspendCoorutineUninterceptedOrReturn and resume its continuation in the completion callback of the newly started coroutine

We can further improve our withContext function by using Kotlin generics, allowing it to return the result of the suspending lambda

Now we can simply replace the inner launch function with withContext and everything works as expected

Let’s look at our current implementation of launch and withContext

That’s a lot of duplicate code. Let’s eliminate redundancy by delegating the dispatch logic to a separate class

This is much better. Our coroutines API has come a long way!

launch is used to start a coroutine that runs in parallel with rest of the code

withContext is used to suspend the current coroutine until the newly launched coroutine, running on a different dispatcher, finishes and optionally returns a result

suspendCoroutineUninterceptedOrReturn is used to suspend current coroutine and manually resume it later. However, it requires special attention when used in code

We must not forget to return COROUTINE_SUSPENDED if we are to resume the continuation ourselves. Additionally, the continuation should always be wrapped with DispatcherContinuation to ensure that the coroutine resumes on the correct thread. Not following these guidelines can lead to issues in the code, so let’s define a simple function that handles these requirements

This function is actually a part of the kotlin.coroutines package

The intercepted() extension function wraps the continuation in DispatchedContinuation, which is then wrapped in a SafeContinuation which ensures that the continuation is not resumed more than once

The official coroutines docs recommend using suspendCoroutine instead of suspendCoroutineUninterceptedOrReturn in everyday code because its less error-prone and automatically manages the required thread dispatch

As a result, functions from the kotlin.coroutines.intrinsics package, including suspendCoroutineUninterceptedOrReturn, are not automatically imported by Android Studio

Additionally, there is another helper function called createCoroutine that can be utilized when creating a new coroutine.

Which can be roughly implemented like this

Here’s our complete code so far

Notice that we’re using our custom createCoroutine and suspendCoroutine implementations instead of the built-in ones, which have a slightly different dispatch mechanism that we’ll discuss later

A Practical Use Case: Copying Files

Let’s consider a common scenario of copying files
We need to write a simple program that copies a list of files to a destination directory but in case there is already a file with the same name in the destination directory, the copy function will take a confirmation from user before proceeding to replace the file

Here’s a sample implementation utilising our coroutine builders

Imagine how complicated the code would become if we were to implement this functionality without utilising coroutines. This is precisely the situation where coroutines demonstrate their true advantage!

Moving on, we attempt to improve performance by starting a new coroutine for each individual file copy operation in a parallel manner

Unfortunately, we quickly encounter an OutOfMemoryException with this approach, when dealing with a large amount of files. Seems like starting a new thread per launch(Dispatcher.Background) invocation is not an ideal choice

We’ll improve upon this further in the next article. Stay tuned!