How Suspend Functions Work in Kotlin: Under the hood

The suspension capability is the most essential feature upon which all other Kotlin Coroutines concepts are built. Your main goal in this article will be to gain a solid understanding of how it works internally.

In one line: Suspending a coroutine means stopping it in the middle.

Analogy:

Let’s first understand the actual meaning of suspension with some real-world examples:

1. Let's start playing a video game:
   You are playing well(Suppose).
   You reach a checkpoint.
   You save the current position.You turn off the game and both you and your computer focus on doing different things now.
   You are done with your tasks and you came back and resume from where you saved.

2. Imagine you’re multitasking in the kitchen:
   Start a recipe (coroutine launch): You begin a recipe, by chopping vegetables (coroutine task 1).
   Pause and switch tasks (coroutine suspension): The oven needs preheating (coroutine task 2). You pause chopping (suspend task 1) and set the oven (complete task 2).
   Resume the first task (coroutine resumption): Once the oven is preheated or turned on, you go back to chopping vegetables (resume task 1) where you left off.
   Repeat as needed (multiple coroutines): You can switch between tasks (coroutines) like checking if the water is boiling (another coroutine task) while waiting for something else.

Now, these examples are the best analogy of coroutine suspension.

A coroutine can start executing the function, it can suspend(save the state and leave the thread for others), and later resume, once any suspending task(like a network call) is done executing.

Let’s see a suspending function in action:

suspend fun main() {
  println("Before")
  println("After")
}
// Before
// After

This is a simple program that will print “Before” and “After”. What will happen if we suspend in between these two prints? For that, we can use the suspendCoroutine function provided by the standard Kotlin library.

suspend fun main() {
  println("Before")
  suspendCoroutine<Unit> { }
  println("After")
}
// Before

If you call the above code, you will not see the “After”, and the code will not stop running (as our main function never finished). The coroutine is suspended after “Before”. Our game was stopped and never resumed. So, how can we resume?

This suspendCoroutine invocation ends with a lambda expression ({ }). The function passed as an argument will be invoked before the suspension. This function gets a continuation as an argument.

We can use this continuation to resume our coroutine. This lambda is used to store this continuation somewhere or to plan whether to resume it. We could use it to resume immediately:

suspend fun main() {
  println("Before")
  suspendCoroutine<Unit> { continuation ->
    continuation.resume(Unit)
  }
  println("After")
}
// Before
// After

You might imagine that here we suspend and immediately resume. This is a good intuition, but the truth is that there is an optimization that prevents a suspension if resuming is immediate.

We can use a function named “delay” to suspend a coroutine for a specific amount of time and then resume it.

suspend fun main() {
  println("Before")
  delay(1000)
  println("After")
}
// Before
// (suspended for 1 second)
// After

One thing that needs to be emphasized here is that we suspend a coroutine, not a function. Suspending functions are not coroutines, just functions that can suspend a coroutine.

Notice that this is very different from a thread, which cannot be saved, only blocked.

A coroutine is much more powerful. When suspended

1. It does not consume any resources.

2. A coroutine can be resumed on a different thread

3. A continuation(explained below) can be serialized, deserialized and then resumed.

Internals:

In Kotlin, suspended functions are implemented with the Continuation passing style. This means that continuations are passed from function to function as arguments(Same as Composer in Jetpack Compose).

suspend fun getUser(): User?
suspend fun setUser(user: User)
suspend fun checkAvailability(flight: Flight): Boolean

// under the hood is
fun getUser(continuation: Continuation<*>): Any?
fun setUser(user: User, continuation: Continuation<*>): Any
fun checkAvailability(
flight: Flight,
continuation: Continuation<*>): Any

The returned type is changed to Any because a suspended function can also return a “COROUTINE_SUSPENDED” apart from the defined return type.

Let's take a simple suspending function:

suspend fun myFunction() {
  println("Before")
  delay(1000) // suspending
  println("After")
}

The next thing is that this function needs its continuation to remember its state. Let’s name it “MyFunctionContinuation”.

The function could be started from two places: either from the beginning (in the case of a first call) or from the point after suspension (in the case of resuming from continuation). To identify the current state, we use a field called label. At the start, it is 0, therefore the function will start from the beginning. However, it is set to the next state before each suspension point so that we start from just after the suspension point after a resume.

// A simplified picture of how myFunction looks under the hood
fun myFunction(continuation: Continuation<Unit>): Any {
  if (continuation.label == 0) { //Starting point
    println("Before")
    continuation.label = 1 //Update just before suspension
    if (delay(1000, continuation) == COROUTINE_SUSPENDED){
      return COROUTINE_SUSPENDED
    }
  }
  //Point after suspension
  if (continuation.label == 1) {
    println("After")
    return Unit
  }
  error("Impossible")
}

When the delay is called, it returns COROUTINE_SUSPENDED, then myFunction returns COROUTINE_SUSPENDED; the same is done by the function that called it, the function that called this function, and all other functions until the top of the call stack. This is how a suspension ends all these functions and leaves the thread available for other runnables (including coroutines) to be used.

What would happen if this “delay ” call didn’t return the COROUTINE_SUSPENDED? What if it just returned the Unit instead? Notice that if the delay just returned a Unit, we would just move to the next state, and the function would behave like any other.

And, Internally continuation class looks something(after many simplifications) like this:

cont = object : ContinuationImpl(continuation) {
  var result: Any? = null
  var label = 0
  override fun invokeSuspend(`$result`: Any?): Any? {
    this.result = `$result`;
    return myFunction(this);
  }
};

Also, if you want to store a state:

suspend fun myFunction() {
  println("Before")
  var counter = 0 //local state
  delay(1000) // suspending
  counter++
  println("Counter: $counter")
  println("After")
}

Here counter is needed in two states (for a label equal to 0 and 1), so it needs to be kept in the continuation. It will be stored right before suspension. Restoring these kinds of properties happens at the beginning of the function. So, this is how the (simplified) function looks under the hood:

fun myFunction(continuation: Continuation<Unit>): Any {
  var counter = continuation.counter //restoring the value at start
  if (continuation.label == 0) {
    println("Before")
    counter = 0 //user-defined
    continuation.counter = counter //saving the value just before suspension
    continuation.label = 1
    if (delay(1000, continuation) == COROUTINE_SUSPENDED){
      return COROUTINE_SUSPENDED
    }
  }
  if (continuation.label == 1) {
    counter = (counter as Int) + 1 //user-defined
    println("Counter: $counter")
    println("After")
    return Unit
  }
  error("Impossible")
}

//Continuation object internal working
class MyFunctionContinuation(val completion: Continuation<Unit>) : Continuation<Unit> {
  override val context: CoroutineContext
    get() = completion.context

  var result: Result<Unit>? = null
  var label = 0
  var counter = 0 //save like a state

  override fun resumeWith(result: Result<Unit>) {
    this.result = result
    val res = try {
      val r = myFunction(this)
      if (r == COROUTINE_SUSPENDED) return
      Result.success(r as Unit)
    } catch (e: Throwable) {
      Result.failure(e)
    }
    completion.resumeWith(res)
  }
}

A similar picture will be there for our most common use cases like making an API call…

suspend fun printUser(token: String) {
  println("Before")
  val userId = getUserId(token) // suspending network call
  println("Got userId: $userId")
  val userName = getUserName(userId, token) // suspending network call
  println(User(userId, userName))
  println("After")
}

Just the difference in this case will be,

1. Storing the results of our functions as shown below…

...
val res = getUserId(token, continuation)
if (res == COROUTINE_SUSPENDED) {
return COROUTINE_SUSPENDED
}
result = Result.success(res)//store the success when when function didn't suspend
...

2. Restore the states at the beginning of our function, like we were doing it for the “counter ” variable:

var counter = continuation.counter// In previous example

var result: Result<Any>? = continuation.result// In this example(when suspend function returns a value)

One thing to note here is that the Continuations will act as our Call Stack and store the states (label, params, variables) of our functions.

If A and B are two suspend functions and A calls B inside of it. Then B will also store the Continuation of A as its Completion state.

BContinuation(
  i = 4,
  label = 1,
  completion = AContinuation(
    i = 4,
    label = 1,
    completion = ...

Internally, these things are implemented in a more complex-looking way, for the sake of Optimizations.

It is implemented with a Loop instead of Recursion.

The key lessons are:

1. Suspending functions are like state machines, with a possible state at the beginning of the function and after each suspending function call.

2. Both the label identifying the state and the local data are kept in the continuation object.

3. Continuation of one function decorates a continuation of its caller function; as a result, all these continuations represent a call stack that is used when we resume or a resumed function completes.

Reference: https://kt.academy/book/coroutines

I hope you got clarity on the internal workings of the suspend function and coroutines. I’ll be posting more detailed Articles on coroutines and Kotlin, so Follow me and subscribe to emails.