» » Kotlin - how suspend works under the hood

Kotlin - how suspend works under the hood

How does the compiler transform  suspend  code so that coroutines can be suspended and resumed?

Coroutines in Kotlin are represented by the suspend keyword  . Wonder what's going on inside? How does the compiler translate  suspend  blocks into code that supports suspending 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:

Coroutines, a brief introduction

Simply put, coroutines are asynchronous operations in Android. As described in the  documentation , we 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 sequential 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 suspend keyword  . 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 is the compiler actually doing internally when we mark a function as  suspend ?

Suspend under the hood

Let's go back to the suspend function  loginUser, see 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 transforms them into an optimized version of callbacks using a  finite 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 the instance  CoroutineContextthat will be used when resuming.

  • resumeWith resumes coroutine execution with  Result , it can either contain the result of the calculation or an exception.

From Kotlin 1.3 onwards, you can use the extensions functions  resume(value: T)  and  resumeWithException(exception: Throwable) , which are specialized versions of the  resumeWith.

The compiler replaces the suspend keyword with an additional argument  completion (type  Continuation) in the 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 a  User.

The suspend function bytecode actually returns  Any? since it is a union of  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.

Additionally, the interface  Continuation can be seen in:

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

  • You can start a coroutine using the  startCoroutine  extension function in the suspend method. It takes  Continuation as an argument, which will 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  Continuationcalled  DispatchedContinuation , where its method  resumemakes the call  Dispatcher available in the coroutine context  CoroutineContext. All dispatchers ( Dispatchers) will call the method  dispatch, except for type  Dispatchers.Unconfined, it overrides the method  isDispatchNeeded (it is called before the call  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 view 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  Continuation. That's why the parent type is in  Continuation this  Any? instead of the expected return type  User.

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

  1. stores the required data

  2. calls the function  loginUser recursively to resume the 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)
    }
  }
  /* ... */
}

Because  it only invokeSuspend calls  loginUser with an argument  Continuation, the rest of the 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  Continuation in the function is checked:

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 an argument  completion is passed to that instance to resume evaluation. 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:

  • There is a variable  label from  LoginUserStateMachine, which is passed 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 the variable  label.

  • When another suspend function is called inside the state machine, the instance  Continuation (of 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 via a call  continuation.cont.resume (obviously, the input argument  completionis stored in an  continuation.cont instance  variable 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  the suspend  function:

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 into a state machine using callbacks.

Knowing how the compiler works “under the hood” will help you better understand:

  • why  the suspend  function won't return until all the work it started has finished;

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

Related Articles

Add Your Comment

reload, if the code cannot be seen

All comments will be moderated before being published.