Kotlin Suspend Bytecode

2020-02-08T00:00:00+00:00
Justin Brooks (justin@jzbrooks.com, https://jzbrooks.com/)

The suspend keyword plays an important role in kotlin coroutines and was the only language-level change to support the feature. The team at Jetbrains decided to implement coroutines as a library rather than a language level feature. This means that anyone could write their own implementation of coroutines.

What does the keyword actually do? ¶

The keyword marks regular Kotlin functions as able to call other suspending functions. That translates into bytecode by rewriting the function a little bit.

suspend fun first() {}

Compiling with kotlinc test.kt and decompiling with javap -c TestKt.class produces:

public final class TestKt {
  public static final java.lang.Object first(kotlin.coroutines.Continuation<? super kotlin.Unit>);
    Code:
       0: getstatic     #15                 // Field kotlin/Unit.INSTANCE:Lkotlin/Unit;
       3: areturn
}

The keyword adds a parameter of type Continuation<in T> to functions it modifies. It’s always the last parameter in the argument list and is what provides standard kotlin functions the ability to call other suspend functions.

Since generic arguments can’t be primitive types, return types of suspending functions that would have otherwise been primitive are autoboxed.

Continuation Passing Style ¶

One way to handle asynchronus code is to pass callbacks to methods that execute asynchronusly.

fun main() {
    httpRequest("https://api.github.com") { response ->
        response.readBodyBytes { bodyBytes ->
            // ...
        }
    }
}

Diatribes are aplenty on this strategy’s shortcommings. kotlinx.coroutines actually does something pretty similar under the hood without the syntactic overhead. The aforementioned Continuation<T> parameter is used to keep track of suspension points so that control can be passed around asynchronusly.

class Response
{
    suspend fun readBodyBytes() -> ByteArray {
        delay(100)
        return byteArrayOf(0x23, 0x34, 0x42)
    }
}

suspend fun httpRequest(url: String) -> Response {
    delay(1000)
    return Response()
}

fun main() {
    runBlocking {
        val response = httpRequest("https://api.github.com")
        val bodyBytes = reponse.readBodyBytes()
    }
}

After compiling with kotlinc -cp kotlinx-coroutines-core-1.3.3.jar -jvm-target 12 test.kt the type signatures of the functions become:

public final class TestKt {
  public static final java.lang.Object httpRequest(java.lang.String, kotlin.coroutines.Continuation<? super Response>);
  [...]
}

public final class Response {
  public final java.lang.Object readBodyBytes(kotlin.coroutines.Continuation<? super byte[]>);
  [...]
}

The continuation implementation keeps track of the label where the function should resume execution. When suspnding work completes, execution jumps to the point in the suspended function and resumes. The body of the suspend functions show the label from the continuation being read, then a switch that jumps execuition to the correct location.

Just before any suspending functions are called the location just after the suspending code is involved is saved, the suspending function is invoked, and control is passed back to the caller.

  57: getfield      #15                 // Field Response$readBodyBytes$1.label:I
  60: tableswitch   { // 0 to 1
                 0: 84
                 1: 114
           default: 147
      }
  84: aload_2
  85: invokestatic  #36                 // Method kotlin/ResultKt.throwOnFailure:(Ljava/lang/Object;)V
  88: ldc2_w        #37                 // long 100l
  91: aload_3
  92: aload_3
  93: aload_0
  94: putfield      #41                 // Field Response$readBodyBytes$1.L$0:Ljava/lang/Object;
  97: aload_3
  98: iconst_1
  99: putfield      #15                 // Field Response$readBodyBytes$1.label:I
 102: invokestatic  #47                 // Method kotlinx/coroutines/DelayKt.delay:(JLkotlin/coroutines/Continuation;)Ljava/lang/Object;

Coroutines Off the Main Thread ¶

Suspention points are a natural boundary for context switches. Simply resume the function on a different thread. Let’s see what happens we we change dispatchers.

fun main() {
    runBlocking {
        val response = withContext(Dispatchers.IO) {
            httpRequest("https://api.github.com")
        }
        val bodyBytes = response.readBodyBytes()
    }
}

kotlinc -cp kotlinx-coroutines-core-1.3.3.jar -jvm-target 12 test.kt

javap -c TestKt\$main\$1

After 84 the execution is pretty similar to before the dispatcher change.

Bytecodes 36-84 are responsible for building up the arguments to withContext.

At bytecode 73, withContext is passed a coroutine context, a function with two arguments, and a continuation (since withContext itself is a suspending function).

   9: tableswitch   { // 0 to 2
                 0: 36
                 1: 85
                 2: 131
           default: 161
      }
  36: aload_1
  37: invokestatic  #47                 // Method kotlin/ResultKt.throwOnFailure:(Ljava/lang/Object;)V
  40: aload_0
  41: getfield      #49                 // Field p$:Lkotlinx/coroutines/CoroutineScope;
  44: astore_2
  45: invokestatic  #55                 // Method kotlinx/coroutines/Dispatchers.getIO:()Lkotlinx/coroutines/CoroutineDispatcher;
  48: checkcast     #57                 // class kotlin/coroutines/CoroutineContext
  51: new           #59                 // class TestKt$main$1$response$1
  54: dup
  55: aconst_null
  56: invokespecial #63                 // Method TestKt$main$1$response$1."<init>":(Lkotlin/coroutines/Continuation;)V
  59: checkcast     #7                  // class kotlin/jvm/functions/Function2
  62: aload_0
  63: aload_0
  64: aload_2
  65: putfield      #65                 // Field L$0:Ljava/lang/Object;
  68: aload_0
  69: iconst_1
  70: putfield      #41                 // Field label:I
  73: invokestatic  #71                 // Method kotlinx/coroutines/BuildersKt.withContext:(Lkotlin/coroutines/CoroutineContext;Lkotlin/jvm/functions/Function2;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
  76: dup
  77: aload         5
  79: if_acmpne     98
  82: aload         5
  84: areturn
  // [...] preamble to call the next suspending function
 116: putfield      #41                 // Field label:I
 119: invokevirtual #81                 // Method Response.readBodyBytes:(Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 122: dup
 123: aload         5
 125: if_acmpne     152
 128: aload         5
 130: areturn
 131: aload_0

The function with two arguments that’s passed to withContext can be seen by decompiling another class.

javap -c TestKt\$main\$1\$response\$1

 final class TestKt$main$1$response$1 extends kotlin.coroutines.jvm.internal.SuspendLambda implements kotlin.jvm.functions.Function2<kotlinx.coroutines.CoroutineScope, kotlin.coroutines.Continuation<? super Response>, java.lang.Object> {
  java.lang.Object L$0;

  int label;

  [...]

  public final java.lang.Object invoke(java.lang.Object, java.lang.Object);
    Code:
       0: aload_0
       1: aload_1
       2: aload_2
       3: checkcast     #94                 // class kotlin/coroutines/Continuation
       6: invokevirtual #96                 // Method create:(Ljava/lang/Object;Lkotlin/coroutines/Continuation;)Lkotlin/coroutines/Continuation;
       9: checkcast     #2                  // class TestKt$main$1$response$1
      12: getstatic     #102                // Field kotlin/Unit.INSTANCE:Lkotlin/Unit;
      15: invokevirtual #104                // Method invokeSuspend:(Ljava/lang/Object;)Ljava/lang/Object;
      18: areturn
 }

From kotlinx.coroutines:

public suspend fun <T> withContext(
     context: CoroutineContext,
     block: suspend CoroutineScope.() -> T
): T = suspendCoroutineUninterceptedOrReturn sc@ { uCont ->
  // compute new context
  val oldContext = uCont.context
  val newContext = oldContext + context
  // always check for cancellation of new context
  newContext.checkCompletion()
  // FAST PATH #1 -- new context is the same as the old one
  if (newContext === oldContext) {
     val coroutine = ScopeCoroutine(newContext, uCont)
     return@sc coroutine.startUndispatchedOrReturn(coroutine, block)
  }
  // FAST PATH #2 -- the new dispatcher is the same as the old one (something else changed)
  // `equals` is used by design (see equals implementation is wrapper context like ExecutorCoroutineDispatcher)
  if (newContext[ContinuationInterceptor] == oldContext[ContinuationInterceptor]) {
     val coroutine = UndispatchedCoroutine(newContext, uCont)
     // There are changes in the context, so this thread needs to be updated
     withCoroutineContext(newContext, null) {
         return@sc coroutine.startUndispatchedOrReturn(coroutine, block)
     }
  }
  // SLOW PATH -- use new dispatcher
  val coroutine = DispatchedCoroutine(newContext, uCont)
  coroutine.initParentJob()
  block.startCoroutineCancellable(coroutine, coroutine)
  coroutine.getResult()
}

The lambda passed to withContext is an extension function, which means the compiler generates a function that takes the reciever as the first argument. Since the lambda is also suspending, a Continuation parameter is added to the parameter list by the compiler. Hence, a function with two arguments.