Kotlin how suspend works under the hood

How does the compiler translate suspend code so that coroutines can be paused and resumed?

Coroutines in Kotlin are represented by the keyword suspend. Wonder what’s going on inside? How does the compiler translate suspend blocks into code that supports pausing and resuming a coroutine?

Knowing this will help you understand why the suspend function does not return until all running work has completed and how code can suspend execution without blocking threads.

TL;DR; The Kotlin compiler creates a special state machine for each suspend function, this machine takes control of the coroutine!

New to Android? Take a look at these helpful coroutine resources:

For those who prefer video:

Coroutines, a brief introduction

Simply put, coroutines are asynchronous operations in Android. As described in documentationwe can use coroutines to manage asynchronous tasks that could otherwise block the main thread and cause the application UI to hang.

It is also convenient to use coroutines to replace the callback code with imperative code. For example, look at this code using callbacks:

// Simplified code that only considers the happy path
fun loginUser(userId: String, password: String, userResult: Callback<User>) {
  // Async callbacks
  userRemoteDataSource.logUserIn { user ->
    // Successful network request
    userLocalDataSource.logUserIn(user) { userDb ->
      // Result saved in DB
      userResult.success(userDb)
    }
  }
}

We replace these callbacks with successive function calls using coroutines:

suspend fun loginUser(userId: String, password: String): UserDb {
  val user = userRemoteDataSource.logUserIn(userId, password)
  val userDb = userLocalDataSource.logUserIn(user)
  return userDb
}

For functions that are called in coroutines, we have added the keyword suspend. So the compiler knows that these functions are for coroutines. From a developer’s point of view, think of a suspend function as a normal function that can be suspended and resumed at a specific moment.

Unlike callbacks, coroutines offer an easy way to switch between threads and handle exceptions.

But what does the compiler actually do internally when we mark a function as suspend?

Suspend under the hood

Let’s go back to the suspend function loginUsersee if the other functions it calls are also suspend functions:

suspend fun loginUser(userId: String, password: String): UserDb {
  val user = userRemoteDataSource.logUserIn(userId, password)
  val userDb = userLocalDataSource.logUserIn(user)
  return userDb
}

// UserRemoteDataSource.kt
suspend fun logUserIn(userId: String, password: String): User

// UserLocalDataSource.kt
suspend fun logUserIn(userId: String): UserDb

In short, the Kotlin compiler takes suspend functions and converts them into an optimized version of callbacks using state machine (which we’ll talk about later).

Continuation interface

Suspend functions interact with each other using Continuation objects. Continuation an object is a simple generic interface with additional data. Later we will see that the generated state machine for the suspend function will implement this interface.

The interface itself looks like this:

interface Continuation<in T> {
  public val context: CoroutineContext
  public fun resumeWith(value: Result<T>)
}
  • context this is an instance CoroutineContextThe to be used when resuming.

  • resumeWith resumes coroutine execution with resultit can either contain the result of a calculation or an exception.

From Kotlin 1.3 onwards, you can use the extensions function resume(value: T) And resumeWithException(exception: Throwable)these are specialized versions of the method resumeWith.

The compiler replaces the suspend keyword with an additional argument completion (type Continuation) in a function, the argument is used to pass the result of the suspend function to the calling coroutine:

fun loginUser(userId: String, password: String, completion: Continuation<Any?>) {
  val user = userRemoteDataSource.logUserIn(userId, password)
  val userDb = userLocalDataSource.logUserIn(user)
  completion.resume(userDb)
}

To simplify, our example returns Unit instead of an object User.

The suspend function bytecode actually returns Any? since it is a union of types T | COROUTINE_SUSPENDED. This allows a function to return a result synchronously when possible.

If the suspend function does not call other suspend functions, the compiler adds the Continuation argument but does nothing with it, the function bytecode will look like a normal function.

In addition, the interface Continuation can be seen in:

  • When converting callback APIs to coroutines using suspendCoroutine or suspendCancellableCoroutine (preferably used in most cases). You directly interact with the instance Continuationto resume a coroutine suspended after executing a block of code from the suspend function’s arguments.

  • You can start a coroutine with startCoroutine extension functions in the suspend method. She accepts Continuation as an argument to be called when the new coroutine completes with either a result or an exception.

Using Dispatchers

You can switch between different dispatchers to run calculations on different threads. How does Kotlin know where to resume suspend computations?

There is a subtype Continuationit is called Dispatched Continuationwhere its method resumemakes a call Dispatcher available in the context of the coroutine CoroutineContext. All dispatchers (Dispatchers) will call the method dispatchexcept for the type Dispatchers.Unconfinedit overrides the method isDispatchNeeded (it is called before calling dispatch) which returns false in this case.

Generated state machine

Clarification: The given code does not fully correspond to the bytecode generated by the compiler. This will be Kotlin code, accurate enough to understand what is really going on inside. This representation was generated by coroutines version 1.3.3 and may change in future versions of the library.

The Kotlin compiler determines when a function can stop inside. Each breakpoint is represented as a separate state in the state machine. The compiler marks such states with labels:

fun loginUser(userId: String, password: String, completion: Continuation<Any?>) {
  
  // Label 0 -> first execution
  val user = userRemoteDataSource.logUserIn(userId, password)
  
  // Label 1 -> resumes from userRemoteDataSource
  val userDb = userLocalDataSource.logUserIn(user)
  
  // Label 2 -> resumes from userLocalDataSource
  completion.resume(userDb)
}

The compiler uses when for states:

fun loginUser(userId: String, password: String, completion: Continuation<Any?>) {
  when(label) {
    0 -> { // Label 0 -> first execution
        userRemoteDataSource.logUserIn(userId, password)
    }
    1 -> { // Label 1 -> resumes from userRemoteDataSource
        userLocalDataSource.logUserIn(user)
    }
    2 -> { // Label 2 -> resumes from userLocalDataSource
        completion.resume(userDb)
    }
    else -> throw IllegalStateException(/* ... */)
  }
}

This code is incomplete because different states cannot exchange information. The compiler uses the same object to exchange Continuation. That’s why the parent type in Continuation this Any? instead of expected return type User.

In doing so, the compiler creates a private class that:

  1. stores the required data

  2. calls a function loginUser recursively to resume calculation

Below is an example of such a generated class:

Comments in the code have been added manually to explain the actions

fun loginUser(userId: String?, password: String?, completion: Continuation<Any?>) {
  
  class LoginUserStateMachine(
    // completion parameter is the callback to the function 
    // that called loginUser
    completion: Continuation<Any?>
  ): CoroutineImpl(completion) {
  
    // Local variables of the suspend function
    var user: User? = null
    var userDb: UserDb? = null
  
    // Common objects for all CoroutineImpls
    var result: Any? = null
    var label: Int = 0
  
    // this function calls the loginUser again to trigger the
    // state machine (label will be already in the next state) and
    // result will be the result of the previous state's computation
    override fun invokeSuspend(result: Any?) {
      this.result = result
      loginUser(null, null, this)
    }
  }
  /* ... */
}

Insofar as invokeSuspend causes loginUser with argument only Continuationother arguments in the function loginUser will be null. At this point, the compiler only needs to add information on how to transition from one state to another.

The compiler needs to know:

  1. The function is called for the first time or

  2. The function was resumed from a previous state To do this, the type of the argument is checked Continuation in function:

fun loginUser(userId: String?, password: String?, completion: Continuation<Any?>) {
  /* ... */
  val continuation = completion as? LoginUserStateMachine ?: LoginUserStateMachine(completion)
  /* ... */
}

If the function is called for the first time, then a new instance is created. LoginUserStateMachineand argument completion is passed to this instance to resume the calculation. Otherwise, the execution of the state machine will continue.

Let’s take a look at the code that the compiler generates to change states and exchange information between them:

fun loginUser(userId: String?, password: String?, completion: Continuation<Any?>) {
    /* ... */

    val continuation = completion as? LoginUserStateMachine ?: LoginUserStateMachine(completion)

    when(continuation.label) {
        0 -> {
            // Checks for failures
            throwOnFailure(continuation.result)
            // Next time this continuation is called, it should go to state 1
            continuation.label = 1
            // The continuation object is passed to logUserIn to resume 
            // this state machine's execution when it finishes
            userRemoteDataSource.logUserIn(userId!!, password!!, continuation)
        }
        1 -> {
            // Checks for failures
            throwOnFailure(continuation.result)
            // Gets the result of the previous state
            continuation.user = continuation.result as User
            // Next time this continuation is called, it should go to state 2
            continuation.label = 2
            // The continuation object is passed to logUserIn to resume 
            // this state machine's execution when it finishes
            userLocalDataSource.logUserIn(continuation.user, continuation)
        }
          /* ... leaving out the last state on purpose */
    }
}

Notice the differences between this and the previous code example:

  • A variable has appeared label from LoginUserStateMachinewhich is transferred to when.

  • Each time a new state is processed, it checks to see if there is an error.

  • Before calling the next suspend function (logUserIn), LoginUserStateMachineupdates a variable label.

  • When another suspend function is called inside the state machine, the instance Continuation (with type LoginUserStateMachine) is passed as an argument. The nested suspend function has also been converted by the compiler with its own state machine. When this internal state machine has completed its work, it will resume the execution of the “parent” state machine.

The last state is to resume execution completion through the challenge continuation.cont.resume (obviously, the input argument completionstored in a variable continuation.cont instance LoginUserStateMachine):

fun loginUser(userId: String?, password: String?, completion: Continuation<Any?>) {
    /* ... */

    val continuation = completion as? LoginUserStateMachine ?: LoginUserStateMachine(completion)

    when(continuation.label) {
        /* ... */
        2 -> {
            // Checks for failures
            throwOnFailure(continuation.result)
            // Gets the result of the previous state
            continuation.userDb = continuation.result as UserDb
            // Resumes the execution of the function that called this one
            continuation.cont.resume(continuation.userDb)
        }
        else -> throw IllegalStateException(/* ... */)
    }
}

The Kotlin compiler does a lot of work under the hood. From suspend features:

suspend fun loginUser(userId: String, password: String): User {
  val user = userRemoteDataSource.logUserIn(userId, password)
  val userDb = userLocalDataSource.logUserIn(user)
  return userDb
}

A large piece of code is generated:

fun loginUser(userId: String?, password: String?, completion: Continuation<Any?>) {

    class LoginUserStateMachine(
        // completion parameter is the callback to the function that called loginUser
        completion: Continuation<Any?>
    ): CoroutineImpl(completion) {
        // objects to store across the suspend function
        var user: User? = null
        var userDb: UserDb? = null

        // Common objects for all CoroutineImpl
        var result: Any? = null
        var label: Int = 0

        // this function calls the loginUser again to trigger the 
        // state machine (label will be already in the next state) and 
        // result will be the result of the previous state's computation
        override fun invokeSuspend(result: Any?) {
            this.result = result
            loginUser(null, null, this)
        }
    }

    val continuation = completion as? LoginUserStateMachine ?: LoginUserStateMachine(completion)

    when(continuation.label) {
        0 -> {
            // Checks for failures
            throwOnFailure(continuation.result)
            // Next time this continuation is called, it should go to state 1
            continuation.label = 1
            // The continuation object is passed to logUserIn to resume 
            // this state machine's execution when it finishes
            userRemoteDataSource.logUserIn(userId!!, password!!, continuation)
        }
        1 -> {
            // Checks for failures
            throwOnFailure(continuation.result)
            // Gets the result of the previous state
            continuation.user = continuation.result as User
            // Next time this continuation is called, it should go to state 2
            continuation.label = 2
            // The continuation object is passed to logUserIn to resume 
            // this state machine's execution when it finishes
            userLocalDataSource.logUserIn(continuation.user, continuation)
        }
        2 -> {
            // Checks for failures
            throwOnFailure(continuation.result)
            // Gets the result of the previous state
            continuation.userDb = continuation.result as UserDb
            // Resumes the execution of the function that called this one
            continuation.cont.resume(continuation.userDb)
        }
        else -> throw IllegalStateException(/* ... */)
    }
}

The Kotlin compiler converts each suspend function to the state machine using callbacks.

By knowing how the compiler works “under the hood”, you better understand:

  • why suspend the function will not return the result until all the work that it started is completed;

  • how the code is suspended without blocking threads (all information about what needs to be done when resuming work is stored in the object Continuation).

Similar Posts

Leave a Reply

Your email address will not be published. Required fields are marked *