Kotlinのsuspend関数のバイトコードを読んでみた

kotlin logo

こんにちは。福岡研究所の岩本(@odiak_)です。

みなさん、Kotlinのコルーチンを使っていますか?

私は、最近久しぶりにAndroidのコードを触る機会があり(3年ぶりくらいでしょうか)、以前から存在は知っていたものの詳しく知らなかったコルーチンを少し使ってみました。まずドキュメントを読んでみたのですが、よくデザインされているなと感じました。今回は使っていませんが、ChannelやFlowなども良さそうです。

この記事では、Kotlinのコルーチンを支える言語機能の1つである、suspend修飾子付き関数の動きをバイトコードから読み解いていきます。

対象読者としては、KotlinをAndroidアプリの開発やサーバーサイドで使用していて、言語処理系の挙動にも興味がある方を想定しています。

コルーチンの紹介

ご存知ではない方のために、Kotlinのコルーチンについて簡単に紹介しておきます。

コルーチンは、軽量スレッドのようなものです。コルーチンを起動すると、それはどこかのスレッドで実行されます。デフォルトの動作ではコルーチン毎にスレッドを起動することはなく、プールされたスレッドで実行されるので、大量のコルーチンを一度に起動しても問題ありません。コルーチンは、次のように使います。

import kotlinx.coroutines.*

fun main(args: Array<String>) {
    GlobalScope.launch {
        println("hello")
        delay(1000L)
        println("world")
    }
    Thread.sleep(1200L)
}

ここでは単純に1つのコルーチンを立ち上げて、helloと表示した1秒後にworldと表示しています。より高度な使い方としては、いくつかのコルーチンを立ち上げて並行して何か計算したり、コルーチン同士で通信し合いながら処理を行ったりといったことも可能です。

詳しくは、Kotlinのドキュメントを読んでみてください。

suspend修飾子付き関数

このようなコルーチンの機能の背景にいる登場人物としては大きく分けて、Kotlinの言語自体に備わっている基本的な機能と、コルーチンのライブラリ(kotlinx.coroutines)の2つがあります。前者の言語自体の機能のうち、suspend修飾子付き関数(以降、suspend関数と呼びます)は特に特徴的なものです。この記事では、suspend関数について深く掘り下げていきます。

先ほどの例で挙げたコードでは、関数delayはsuspend関数です。また、GlobalScope.launchに渡しているラムダ式も、明示されてはいませんがsuspend関数です。

suspend関数では、非同期的な処理をまるで同期的な処理のように呼び出すことができます。suspend関数を使わない場合、非同期的な処理を呼び出すにはコールバック関数などを渡す必要がありました。例えば次のように。

fun getPost(id: String, callback: (Content) -> Unit) { /* ... */ }
fun decodeContent(content: String, callback: (String) -> Unit) { /* ... */ }

fun getContent(id: String, callback: (String) -> Unit) {
    getPost(id) { post ->
        decodeContent(post.content) { content ->
            callback(content)
        }
    }
}

コールバックを使うと、ネストが深くなってしまいコードが読みづらくなる上に、条件的に処理を呼び出すなどの複雑なコードが書きづらくなります。これを、コルーチンを使って書くと次のように、まるで同期的な処理のように書くことができます。

suspend fun getPost(id: String): Post { /* ... */ }
suspend fun decodeContent(content: String): String { /* ... */ }

suspend fun getPostContent(id: String): String {
    val post = getPost(id)
    val content = decodeContent(post.content)
    callback(content)
}

suspend関数を呼び出すと、呼び出した側の処理は一旦そこで中断されます。呼び出された関数の処理が終わると、呼び出した側の処理が再開されます。

先ほど関数delayがsuspend関数であると書きましたが、delayを呼び出した場合も同じように一度処理が中断され、指定した時間が経過してから処理が再開されます。注意したいのは、Thread.sleepが処理をブロックするのとは違い、delayのようなsuspend関数は処理をブロックはしないということです。

解説している動画を見てみたが、腑に落ちない

使っていて、ふと疑問が頭に浮かびました。こんな魔法のようなものがどうやって動いているんだろう、と。実行されるのはJVMの上だし、Javaにはこんな機能ありません。

そこで、そういった内部の話を解説しているというYouTube動画を見てみました。

KotlinConf 2017 - Deep Dive into Coroutines on JVM by Roman Elizarov

この動画の前半で、suspend関数はコンパイルされると継続渡しスタイルに変換されて、しかもその継続はステートマシン的なもので表現されるので効率的だよというような話をしています。

最初に見たときは、大まかには理解できたものの、どこか腑に落ちない感覚がありました。

そこで、suspendの付いた関数を使ったコードをJVM向けにコンパイルして、そのバイトコードを見てみることにしました。

以下で行っていることは、先ほどの動画で話されている内容を実際に手を動かして確認してみた、という部分が多いです。
ただ、私はそのステップを踏むことで理解が大幅に進みましたし、その過程は非常に楽しいものでした。

バイトコードを読んでみる

Kotlinのソースコードをコンパイルするといくつかのクラスファイルができます。それをjavapコマンド(JDKに付属しています)でダンプしてみます。

今回は、次のようなソースコードを使いました。

package net.odiak.kotlin_coroutines_experiment

import kotlinx.coroutines.*

suspend fun s1(): Int {
    println("hello")
    delay(1000L)
    println("world")
    return 42
}

fun main(args: Array<String>) {
    runBlocking {
        println(s1())
    }
}

こちらをコンパイルすると、 AppKt, AppKt$s1$1, AppKt$main$1 という3つのクラスができます。

この記事では、そのうちAppKtAppKt$s1$1の2つを見てみます。それぞれを、 javap -c AppKt のように-cオプション付きで表示してみます。すると、次のようになります。

(長くなるのでAppKt$main$1を省略しましたが、読者の皆さんにはぜひ自身でコンパイルして結果を確かめてみていただきたいです)

public final class net.odiak.kotlin_coroutines_experiment.AppKt {
  public static final java.lang.Object s1(kotlin.coroutines.Continuation<? super java.lang.Integer>);
    Code:
       0: aload_0
       1: instanceof    #11       // class net/odiak/kotlin_coroutines_experiment/AppKt$s1$1
       4: ifeq          39
       7: aload_0
       8: checkcast     #11       // class net/odiak/kotlin_coroutines_experiment/AppKt$s1$1
      11: astore        4
      13: aload         4
      15: getfield      #15       // Field net/odiak/kotlin_coroutines_experiment/AppKt$s1$1.label:I
      18: ldc           #16       // int -2147483648
      20: iand
      21: ifeq          39
      24: aload         4
      26: dup
      27: getfield      #15       // Field net/odiak/kotlin_coroutines_experiment/AppKt$s1$1.label:I
      30: ldc           #16       // int -2147483648
      32: isub
      33: putfield      #15       // Field net/odiak/kotlin_coroutines_experiment/AppKt$s1$1.label:I
      36: goto          49
      39: new           #11       // class net/odiak/kotlin_coroutines_experiment/AppKt$s1$1
      42: dup
      43: aload_0
      44: invokespecial #20       // Method net/odiak/kotlin_coroutines_experiment/AppKt$s1$1."<init>":(Lkotlin/coroutines/Continuation;)V
      47: astore        4
      49: aload         4
      51: getfield      #24       // Field net/odiak/kotlin_coroutines_experiment/AppKt$s1$1.result:Ljava/lang/Object;
      54: astore_3
      55: invokestatic  #30       // Method kotlin/coroutines/intrinsics/IntrinsicsKt.getCOROUTINE_SUSPENDED:()Ljava/lang/Object;
      58: astore        5
      60: aload         4
      62: getfield      #15       // Field net/odiak/kotlin_coroutines_experiment/AppKt$s1$1.label:I
      65: tableswitch   { // 0 to 1
                     0: 88
                     1: 127
               default: 151
          }
      88: aload_3
      89: invokestatic  #36       // Method kotlin/ResultKt.throwOnFailure:(Ljava/lang/Object;)V
      92: ldc           #38       // String hello
      94: astore_1
      95: iconst_0
      96: istore_2
      97: getstatic     #44       // Field java/lang/System.out:Ljava/io/PrintStream;
     100: aload_1
     101: invokevirtual #49       // Method java/io/PrintStream.println:(Ljava/lang/Object;)V
     104: ldc2_w        #50       // long 1000l
     107: aload         4
     109: aload         4
     111: iconst_1
     112: putfield      #15       // Field net/odiak/kotlin_coroutines_experiment/AppKt$s1$1.label:I
     115: invokestatic  #57       // Method kotlinx/coroutines/DelayKt.delay:(JLkotlin/coroutines/Continuation;)Ljava/lang/Object;
     118: dup
     119: aload         5
     121: if_acmpne     132
     124: aload         5
     126: areturn
     127: aload_3
     128: invokestatic  #36       // Method kotlin/ResultKt.throwOnFailure:(Ljava/lang/Object;)V
     131: aload_3
     132: pop
     133: ldc           #59       // String world
     135: astore_1
     136: iconst_0
     137: istore_2
     138: getstatic     #44       // Field java/lang/System.out:Ljava/io/PrintStream;
     141: aload_1
     142: invokevirtual #49       // Method java/io/PrintStream.println:(Ljava/lang/Object;)V
     145: bipush        42
     147: invokestatic  #65       // Method kotlin/coroutines/jvm/internal/Boxing.boxInt:(I)Ljava/lang/Integer;
     150: areturn
     151: new           #67       // class java/lang/IllegalStateException
     154: dup
     155: ldc           #69       // String call to 'resume' before 'invoke' with coroutine
     157: invokespecial #72       // Method java/lang/IllegalStateException."<init>":(Ljava/lang/String;)V
     160: athrow

  public static final void main(java.lang.String[]);
    Code:
       0: aload_0
       1: ldc           #81       // String args
       3: invokestatic  #87       // Method kotlin/jvm/internal/Intrinsics.checkNotNullParameter:(Ljava/lang/Object;Ljava/lang/String;)V
       6: aconst_null
       7: new           #89       // class net/odiak/kotlin_coroutines_experiment/AppKt$main$1
      10: dup
      11: aconst_null
      12: invokespecial #90       // Method net/odiak/kotlin_coroutines_experiment/AppKt$main$1."<init>":(Lkotlin/coroutines/Continuation;)V
      15: checkcast     #92       // class kotlin/jvm/functions/Function2
      18: iconst_1
      19: aconst_null
      20: invokestatic  #98       // Method kotlinx/coroutines/BuildersKt.runBlocking$default:(Lkotlin/coroutines/CoroutineContext;Lkotlin/jvm/functions/Function2;ILjava/lang/Object;)Ljava/lang/Object;
      23: pop
      24: return
}
final class net.odiak.kotlin_coroutines_experiment.AppKt$s1$1 extends kotlin.coroutines.jvm.internal.ContinuationImpl {
  java.lang.Object result;

  int label;

  public final java.lang.Object invokeSuspend(java.lang.Object);
    Code:
       0: aload_0
       1: aload_1
       2: putfield      #12       // Field result:Ljava/lang/Object;
       5: aload_0
       6: aload_0
       7: getfield      #16       // Field label:I
      10: ldc           #17       // int -2147483648
      12: ior
      13: putfield      #16       // Field label:I
      16: aload_0
      17: invokestatic  #23       // Method net/odiak/kotlin_coroutines_experiment/AppKt.s1:(Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
      20: areturn

  net.odiak.kotlin_coroutines_experiment.AppKt$s1$1(kotlin.coroutines.Continuation);
    Code:
       0: aload_0
       1: aload_1
       2: invokespecial #30       // Method kotlin/coroutines/jvm/internal/ContinuationImpl."<init>":(Lkotlin/coroutines/Continuation;)V
       5: return
}

初めて見る方は何が何やら…という感じだと思いますが、コメントにクラス名やメソッド名が書いてあるのでなんとなく読めるのではと思います。
JVMはスタックマシンなので、命令を呼ぶことでスタックに値を入れたり出したりします。
それぞれの命令がどのような意味かは、リファレンスや本で必要なところだけ見てください。

まず、名前からも推測できますが、3つのクラスがどのようなものかを紹介します。

  • AppKtには、App.ktに含まれるトップレベル関数が、staticメソッドとして定義されている
  • AppKt$s1$1は、関数s1専用の継続クラス
    • 継続クラスとは、ここではkotlin.coroutine.Continuationインタフェースを実装したクラスを指す
      • 簡単に言うと、suspend関数の途中から続きを実行するために用いられるコールバックのようなもの
    • 2つのフィールドを持っている
      • 戻り値のオブジェクトを保持するresult
      • どこに戻るべきかを表すlabel
  • AppKt$main$1は、関数main専用の継続クラス兼、runBlockingに渡すラムダ関数
    • AppKt$main$1とは継承している継続クラスが少し違うが、大きな差はない

まずは、 AppKt.s1 を見ていきます。動画でも紹介されていた通り、引数に継続オブジェクトが追加されています。

また、戻り値はIntegerではなくObjectになっています。戻り値がObjectなのは、メソッドの処理がまだ終わっていない場合にCOROUTINE_SUSPENDEDという特殊なオブジェクトを返すためです。

0-2行目:
第1引数がAppKt$s1$1のインスタンスではない場合は、39行目にジャンプします。

7-11行目:
第1引数をAppKt$s1$1にキャストして変数に入れます。この変数をcontと呼ぶことにしましょう。

13-33行目:
今回はあまり関係ないですが、contのフィールドlabelの最上位ビットを見て、フラグ操作をしています。最上位ビットが1の場合は、AppKt$s1$1.invokeSuspendから呼ばれた場合です。その場合、最上位ビットを0にしてcontlabelに代入し、49行目にジャンプします。最上位ビットが0の場合は、39行目にジャンプします(このケースは、再帰呼び出しの場合です。なので今回は関係ありません)

39-47行目:
AppKt$s1$1のインスタンスを新しく作り、ローカル変数contに入れます。コンストラクタの引数は、s1の第1引数である継続オブジェクトです。
(簡単にいうと、引数に渡された継続オブジェクトをs1用の継続クラスでラップしているということです)

49-54行目:
contのフィールドresultを変数に代入します。resultと呼ぶことにしましょう。

55-58行目:
Kotlinが定義しているCOROUTINE_SUSPENDEDというオブジェクトを取得し、変数に代入します。

62-65行目:
contlabelを読み、その値を元に処理を分岐します。

  • 0の場合:88行目へ
  • 1の場合:127行目へ
  • それ以外:151行目へ(IllegalStateExceptionを投げるだけ)

88-89行目:
cont.resultが例外を表現している場合は例外を投げる関数throwOnFailure(KotlinのResultというinlineクラスのメソッド)を呼びます。ただし、resultはこの時は初期値(null)のままなので、呼び出す意味はあまりないと思われます。

92-101行目:
println("hello")を呼び出します。

111-112行目:
cont.labelに1を設定します。

104-115行目(スタックの関係で行数が前後しています):
関数delayを呼び出します。引数は、1000Lとcontです。

118-126行目:
delayの戻り値がCOROUTINE_SUSPENDEDなら、同じ値をreturnしてs1から抜けます。そうでない場合は、132行目にジャンプします。

127-128行目(labelが1の場合はここにジャンプしてくる):
先ほどと同じくthrowOnFailureを呼び出します。つまり、delayの実行が失敗していないかをチェックするわけです。

132-142行目:
println("world")を呼び出します。

145-150行目:
42という数値をreturnしてs1から抜けます。

151行目以降:
IllegalStateExceptionを投げているだけです。基本的にはここを通りません。

次に、AppKt$s1$1.invokeSuspendのコードを読んでみますが、その前にinvokeSuspendの立ち位置を理解しておきましょう。

  • invokeSuspendは、その祖先クラス(ContinuationImpl, BaseContinuationImpl)またはContinuationインタフェースを見るとその役割が分かる
    • ContinuationImplBaseContinuationImplの実装はこちら
    • Continuationの定義はこちら
  • まず、ContinuationインタフェースはresumeWithというメソッドを持っており、このメソッドは名前の通り中断した処理を再開する
    • 例えば、delay関数に継続オブジェクトを渡した場合、一定時間が経つとその継続オブジェクトのresumeWithが呼ばれる
  • BaseContinuationImplにおけるresumeWithの実装では、自身のinvokeSuspendメソッドを呼び出す
  • そこでCOROUTINE_SUSPENDEDが返ってきたらメソッドは終了
  • それ以外の値が返ってきた場合、内部に持っている継続オブジェクトへと関心を移す
    • s1のコードで見たように、BaseContinuationImplは他の継続オブジェクトをラップする
  • それが同様にBaseContinuationImplであれば、またinvokeSuspendを呼んで、同じことを繰り返す
  • その他のContinuationであれば、resumeWithを呼び出して終了

はい、invokeSuspendBaseContinuationImplにおける重要なメソッドであることが分かったところで、AppKt$s1$1.invokeSuspendを読んでいきましょう。と言っても、大変短いです。

0-2行目:
第1引数をフィールドresultに入れます。

6-13行目:
フィールドlabelの値を取り出し、最上位ビットを1にして代入します。これは先ほども見たように、s1s1自身から再帰呼び出しとして呼び出されたのか、invokeSuspendから呼び出されたのかを区別するためのフラグです。

16-20行目:
自身(this)を引数にしてs1を呼び出し、その戻り値をreturnします。

コードを読んで分かったことのまとめ

関数s1を中心にいろいろ読んでみましたが、ここで少しまとめておきましょう。

  • suspend関数suspend fun s1() -> Intをコンパイルすると、少しシグネチャが変わった関数fun s1(Continuation<Int>) -> Any?と関数s1のための継続クラスAppKt$s1$1ができる
  • 継続クラスには、待ち合わせていた処理の戻り値を保持するresultと、s1内で処理を再開する位置を表すlabelという2つのフィールドがある
  • 継続クラスには、invokeSuspendというメソッドが定義されており、それは中断されていた関数s1の処理を再開するときに呼ばれる
  • コンパイルされた関数s1の動作について
    • 引数に指定された継続オブジェクトは、AppKt$s1$1invokeSuspendから呼び出された場合を除いて継続クラスAppKt$s1$1でラップされる
    • AppKt$s1$1labelによって、コード内の指定の位置にジャンプする
    • labelが0の場合(初期状態)は始めから
      • "hello"と出力する
      • labelを1に設定する
      • 継続オブジェクトを引数に含めてdelay関数を呼ぶ
      • delay関数はCOROUTINE_SUSPENDEDを返すので一旦処理を中断し、COROUTINE_SUSPENDEDをreturnしてs1を抜ける
    • labelが1の場合は途中から
      • delay関数の実行が失敗していた場合は、例外を投げて終了する
      • "world"と出力する
      • 42をreturnしてs1を抜ける

これを踏まえて、コンパイルされたs1を実行する流れをざっくりとまとめてみます。

  1. 何らかのContinuationインタフェースを実装したオブジェクト(継続オブジェクト)を用意
  2. その継続オブジェクトを引数にしてs1を呼び出す
  3. 継続オブジェクトをAppKt$s1$1でラップする
  4. "hello"を出力する
  5. ラップした継続オブジェクトのlabelを1にする
  6. ラップした継続オブジェクトを引数に含めてdelayを呼び出す
  7. delayCOROUTINE_SUSPENDEDを返すので、s1も同じ値を返して終了
  8. delayに指定した時間が経つ(あるいは実行が失敗する)と、何者かによりラップした継続オブジェクトのresumeWithが呼ばれる
  9. resumeWithが同オブジェクトのinvokeSuspendを呼び出す
  10. invokeSuspendが同オブジェクトを引数に入れてs1を呼び出す
  11. 今度は継続オブジェクトをラップせずにそのまま使う
    • なぜならinvokeSuspends1を呼ぶとき、labelにフラグを立てているから
    • なお、s1でフラグは戻される
  12. 最初に呼ばれた時とはlabelの値が変わっているため、続きから処理が行われる
  13. delayの実行が失敗していた場合は、例外を投げて終了
  14. "world"を出力する
  15. 42を返してs1が終了
    • なお、戻り値はBaseContinuationImplにおけるresumeWithの実装によって、最初にs1へ渡された継続オブジェクトへと渡される

いかがでしたか?
suspend関数がコンパイルされて、同期処理のように書かれたコードがコールバック渡しのように変換されている様子が少しでもお分かりいただけたでしょうか?

この説明ではいろいろなものを省略したので、例えば次のような疑問を抱くかもしれません。

  • 最初の継続オブジェクトはどこで作られるの?
  • delay関数がコールバックを呼ぶ仕組みはどうなっているの?
  • s1関数にループや再帰呼び出しが含まれていた場合はどうなるの?

コードを読んだり同じようにコンパイルしてみたりすれば分かりますが、おまけとして少し触れておきます。

最初の継続オブジェクトはどこで作られるのか

これはkotlinx.couroutinesライブラリの方を読むと何となく分かります。最初の継続オブジェクトは、launchrunBlockingなどの普通の関数からsuspend関数を呼び出すような関数の中で作られています。

あまり詳しくは読んでいませんが、継続オブジェクトなどいろいろな物を用意して、スレッドに実行させたり処理を待ち合わせたりしているようです。

delay関数がコールバックを呼ぶ仕組み

これも軽く読んだ程度ですが、イベントループのような物を使っています。

suspend関数にループや再帰呼び出しが含まれていた場合

ループの場合も、大きくは変わりません。JVMの世界では、ループも単にジャンプを含んだコードになるだけです。
ただし、関数が再び呼び出された際に、ループなどで使用している変数を復元する必要が出てきます。
そこで、継続クラスにフィールドが追加されて、変化する可能性のあるローカル変数をそこに保存しておきます。再び呼び出された際は、適切な場所にジャンプされ、そこで変数を復元することで処理を再開できます。

再帰の場合も特別なことはありません。再帰呼び出しが行われると、継続オブジェクトが同じクラスでどんどんラップされていきます。resumeWithがネストした継続オブジェクトを辿ってくれるので、他のsuspend関数を呼び出す場合と全く同じです。

おわりに

「Kotlinのsuspend関数がバイトコードレベルでどう動いているのか?」という素朴な問いに答えるため、いろいろと調べてみました。
調査するにあたって、コンパイルしたバイトコードを読んだり、関連するライブラリのコードを読んだりしました。やや大変でしたが、程よい難しさで、ソースコード読みの練習としても良かった気がします。
(余談ですが、今回kotlinx.coroutinesなどのコードをGitHub上で読んでいました。今思うと、IDEなど使えばもっと楽だったのでは、という気もします)

今回、「継続」という概念について初めて知りました。まだ雰囲気しか分かっていないので、個人的にもう少し掘り下げてみたいです。

最後までご覧いただきありがとうございました。

ZOZOテクノロジーズでは、各種エンジニアを採用中です。ご興味のある方は以下のリンクからご応募ください。

tech.zozo.com

カテゴリー