The Ultimate Breakdown of Kotlin Coroutines. Part 3. Variable Spilling.

In the previous parts of the breakdown I covered all the essential parts of coroutines, or, as I like to call them - three-headed hydra - suspension, resumption and state-machine. This time I will explain local variables spilling - a process of saving a coroutine state before suspension and restoring them after resumption.

Variables Spilling

All the previous examples did not have local variables, and there is a reason for it. When a coroutine suspends, we should save its local variables. Otherwise, when it resumes, the values of them are lost. So, before the suspension, which can be on each suspend call (more generally, on each suspension point), we save them, and after the resumption, we restore them. There is no reason to restore them right after the call if the call did not return COROUTINE_SUSPENDED: their values are still in local variable slots.

Let us consider a simple example:

import kotlin.coroutines.*

data class A(val i: Int)

var c: Continuation<Unit>? = null

suspend fun suspendMe(): Unit = suspendCoroutine { continuation ->
    c = continuation
}

fun builder(c: suspend () -> Unit) {
    c.startCoroutine(object: Continuation<Unit> {
        override val context = EmptyCoroutineContext
        override fun resumeWith(result: Result<Unit>) {
            result.getOrThrow()
        }
    })
}

suspend operator fun A.plus(a: A) = A(i + a.i)

fun main() {
    val lambda: suspend () -> Unit = {
        val a1 = A(1)
        suspendMe()
        val a2 = A(2)
        println(a1 + a2)
    }
    builder {
        lambda()
    }
    c?.resume(Unit)
}

We should save a1 before suspendMe, and we should restore it after the resumption. Similarly, we should save both a1 and a2 before +, since the compiler does not generally know whether suspend call will suspend, so it assumes that the suspension might happen in each suspension point. So, it spills the locals before each call and unspills them after it.

Thus, the compiler generates the following state machine

fun invokeSuspend($result: Any?): Any? {
    when (this.label) {
        0 -> {
            var a1 = A(1)
            this.L$0 = a1
            this.label = 1
            $result = suspendMe()
            if ($result == COROUTINE_SUSPENDED) return COROUTINE_SUSPENDED
            goto 1_1
        }
        1 -> {
            a1 = this.L$0

            1_1:
            var a2 = A(2)
            this.L$0 = null
            this.label = 2
            $result = plus(a1, a2)
            if ($result == COROUTINE_SUSPENDED) return COROUTINE_SUSPENDED
            goto 2_1
        }
        2 -> {
            2_1:
            println($result)
            return Unit
        }
        else -> {
            throw IllegalStateException("call to 'resume' before 'invoke' with coroutine")
        }
    }
}

As one can see, the generated code does not spill and unspill variables, which are dead, in other words, which are not required afterward. Furthermore, it cleans the field for spilled variables of reference types to avoid memory leaks by pushing null to it so that GC can collect the object. More on that in the following section.

Spilled Variables Cleanup

Since we spill a reference to the continuation object, we now hold an additional reference to the object. Thus, GC cannot clean its memory as long as there is a reference to the continuation. Of course, holding a reference to a not-needed object leads to memory leaks. The compiler clears the fields for reference types up to avoid the leaks.

Consider the following example:

suspend fun blackhole(a: Any?) {}

suspend fun cleanUpExample(a: String, b: String) {
    blackhole(a) // 1
    blackhole(b) // 2
}

After line (1) a is dead, but b is still alive. So, we spill only b. There is no variable alive after line (2), but the continuation object still holds a reference to b in the L$0 field. So, to clean it up and avoid memory leaks, we push null to it.

Generally, the compiler generates spilling and unspilling code so that it uses only the first fields. If there are M fields for references, but we spill only N (where N ≤ M, of course) objects at the suspension point, everything else should be null. However, we do not need to nullify all of them every suspension point. Instead, the compiler checks which of the fields hold references and clears only them.

Additionally, the compiler shrinks and splits LVT records for local variables, so a debugger will not show dead variables as uninitialized.

Stack spilling

In the previous examples, the stack was clean before a call, meaning that there were only call arguments before the call, and only the call result is on the stack after the call.

However, this is not always true. Consider the following example:

val lambda: suspend () -> Unit = {
    val a1 = A(1)
    val a2 = A(2)
    val a3 = A(3)
    a1 + (a2 + a3)
}

and have a closer look at a1 + (a2 + a3) expression. If + were not suspend, the compiler would generate the following bytecode:

ALOAD 1 // a1
ALOAD 2 // a2
ALOAD 3 // a3
INVOKESTATIC plus
INVOKESTATIC plus
ARETURN

We cannot just make this code suspendable since, after the resumption, the stack has only $result (it is passed to resumewith and is the argument of invokeSuspend). So, there are not enough variables on the stack for the second call. Consequently, we need to save the stack before the call and then restore it after the call. Instead of creating the separate logic of stack into slots spilling, we reuse two already existing ones. One is stack normalization, which is already present in inliner. The inliner spills the stack into locals before the inline call and restores them after the call. So, if we do the same here, the bytecode becomes

ALOAD 1 // a1
ASTORE 4 // a1
ALOAD 2 // a2
ALOAD 3 // a3
INVOKESTATIC plus
ASTORE 5 // a2 + a3
ALOAD 4 // a1
ALOAD 5 // a2 + a3
INVOKESTATIC plus
ARETURN

we need to spill a2 + a3 since we should preserve the order of plus's arguments.

If we look at stack normalization once more, we see that there are now five locals, but, thankfully, we do not spill all of them. a2 + a3 is not alive during both calls and is not present in LVT, so there is no reason for the compiler to spill it. The same applies for slot 4: the variable is dead during the second call, so we spill it only once. a2 and a3 are dead during both calls, and thus they are not spilled, as well as a1 during the second call.

That concludes a short explanation of variables spilling. In the next blogpost I will walk through coroutine intrinsic function - suspendCoroutineUninterceptedOrReturn and a non-intrinsic createCoroutineUnintercepted, which are responsible for accessing coroutine's continuation object and coroutine creation, respectively. I will also explain the machinery around coroutines interception, which drives multithreaded coroutines.

Comments

Post a Comment

Popular posts from this blog

The Ultimate Breakdown of Kotlin Coroutines. Part 1. State-Machines and Suspension.

Introduction Throughout 2020 I wrote a document on coroutines codegen, which aimed mostly at my colleagues at JetBrains and Google to help them understand the pitfalls of the codegen and which design decisions we made. These blog posts are a simplified and trimmed down version of the documentation for a larger audience of experienced Kotlin programmers, who just want to know how the compiler turns sequential code into suspendable and become more proficient at their jobs. The blogposts will not cover kotlinx.coroutines internals. However, I think, with the knowledge of coroutines internals, it will be easier to understand, for example, how Flow's backpressure works and why it is essential for Flow operators to be inline with crossinline lambda parameter (spoilers: this is an optimization). I decided to write the blogposts now, before Kotlin 1.5 release, since Kotlin 1.5 brings new JVM_IR back-end with a lot of improvements in code generations, tail-call optimization of suspend fu
Read more

The Ultimate Breakdown of Kotlin Coroutines. Part 4. Create, Start, Suspend, Intercept.

The previous parts mostly explained what the compiler does with coroutines. However, there is the other side of a coin - runtime support. In Kotlin's case, standard library functions, including intrinsic one, which are responsible for all the essential operations - creation, starting, suspension and interception. There is also a part responsible for resumption, but I already explained it in part 2. Coroutine Intrinsics Previous examples had an elementary coroutine builder. They used so-called empty continuation. Let us now recreate kotlinx.coroutines' async function, which runs a coroutine on another thread and then, upon its completion, returns the result to the main thread. First, we need a class which waits on the main thread: class Async < T > { suspend fun await (): T = TODO() } then the root continuation: class AsyncContinuation < T >: Continuation < T > { override val context = EmptyCoroutineContext override fun resumeWith (r
Read more

The Ultimate Breakdown of Kotlin Coroutines. Part 5. Suspend Functions.

Previous parts covered suspend lambdas and the machinery, which makes coroutines tick - suspend and resume. Also, I covered continuations, state-machines and interception. There is only one part remaining - suspend functions. I will omit coroutines inlining, which despite having complicated implementation, has an easy mental model - inlining does not affect behavior of coroutines. All the inlining is done prior to state-machine transformation. Finally, I will omit optimizations for the same reason - they do not change the behavior. Suspend Functions As explained in the continuation-passing style section, every suspending function's signature is changed: the compiler adds a continuation parameter and changes the return type to Any? . The tricky part becomes when we try to make it suspendable, in other words, when we build a state machine and generate a continuation. Unlike suspend lambdas, we cannot reuse an existing class for the continuation, since the suspend function can be s
Read more