Kotlinコンパイラはコルーチンをどのように処理するのか?

TL;DR

コルーチンの概念とKotlinコンパイラの動作原理

こんにちは。ONDAでホテル運営管理ソリューションを開発しているバックエンドエンジニア、Gunner(ガナー、チョン・ジェフン)です。

今回はKotlin(コトリン)環境でCoroutine(コルーチン)がどのように動作するかについて解説します。

まず、コルーチンはコールバックで書かれたコードを、順次実行される形に書き換えられます。例を見てみましょう。

// callback版
fun postItem(item: Item) {
	requestTokenAsync { token ->
		createPostAsync(token, item) { post ->
      processPost(post)
    }
  }
}

// coroutine版
suspend fun postItem(item: Item) {
	val token = requestToken()
	val post = createPost(token, item)
	processPost(post)
}

コルーチンを使えば「コールバック地獄」を回避でき、コードの可読性が上がります。では、コルーチンはどうやってこれを実現しているのでしょうか?

1. コルーチンとサスペンションポイント

Wikipediaでは「Coroutines」を次のように定義しています。

"Coroutines are computer program components that allow execution to be suspended and resumed, …" - Wikipedia

定義によれば、コルーチンは処理を一時停止・再開できるプログラムの構成要素です。

Kotlinのコルーチンでも処理を一時停止・再開できます。Kotlinでは処理を一時停止する地点を「suspension point(サスペンションポイント)」と呼びます。IntelliJ IDEを使用している場合、下の画像のようにサスペンションポイントが表示されます。

サスペンションポイントは、宣言時にsuspend修飾子がついた関数を実行するときに表示されます。上の図では、requestToken()、createPost()、processPost()関数がsuspendで宣言されていることがわかります。

コードが実行中にサスペンションポイントに到達すると、元の処理フローを一時停止し、該当関数を実行してから、元の処理を再開します。内部的にどのように一時停止しているのか、サンプルコードで見てみましょう。

suspend fun main() {
	println("Before")

	suspendCoroutine { continuation ->
		thread {
			println("Suspended")
			Thread.sleep(1000)
			continuation.resumeWith(Result.success(Unit))
			println("Resumed")
		}
	}

	println("After")
}

suspendCoroutine関数を呼ぶと、main()関数の処理が一時停止されます。suspendCoroutineに渡したラムダ内でcontinuation.resumeWith()を呼ぶことで、元の関数(この例ではmain)の実行を再開できます。このようにサスペンションポイントで実行を一時停止し、suspend関数を実行した後、元の処理から一時停止した時点を再開できます。

2. Continuation == Callback

コルーチンの動作を説明する記事には、必ず「Continuation」という概念が登場します。「KotlinコルーチンはCPS(Continuation Passing Style)を使ってコルーチンを一時停止・再開する」という説明をよく目にするでしょう。

上の例で、suspendCoroutine関数に渡したラムダが受け取るパラメータ名が「continuation」でした。continuationは以前からある概念で、Scheme言語では「call/cc」という関数でcontinuationを使用します。call/ccは「call-with-current-continuation」の意味です。

continuationを使って元の関数の処理を一時停止させ、suspend関数が完了したときに元の関数の処理を再開できます。では、continuationには何が入っていて、このような動作が可能なのでしょうか?

私たちが何かをしている途中で一時停止して別のことをし、その後、元の一時停止した作業を再開しようとする場合、いくつか覚えておく(またはどこかに書き留めておく)ことがあります。

作業をどこまで進めたか(どこから再開すべきか)の目印
元の作業を中断したときと同じ環境(Excelで作業していたなら、そのExcelファイルと作業時に参照していた他の資料など)

この2つが準備できて初めて、一時停止した作業を再開できます。

関数を一時停止・再開する場合も同じです。関数を再開するには次の2つが必要です。

処理をどこまで実行したか
中断時のcontext変数値(例: varA = 10、varB = "This is string"というcontextの記録)

この情報をcontinuationに保存して一時停止すれば、他の作業を終えた後、continuationに入っている情報を使って元の関数がしていた作業を再開できます。普段continuationを使うことがないので実感しにくいかもしれませんが、何か似たコンセプトが思い浮かびませんか?

そう、**「callback」**です。通常callbackは、特定のタイミングで実行したいラムダを渡して実行させるものです。continuationも同様に、特定のタイミングで実行したいラムダを渡して実行させます。

ただし、continuationではラムダを実行する特定のタイミングは「今すぐ」で、元の処理を再開するタイミングは「ラムダが終わった後」になります。callbackの場合、ラムダを実行する特定のタイミングは「callback関数内のどこかで呼び出されたとき」で、元の処理を再開するタイミングは「今すぐ」と見なせます。実行タイミングが異なるだけで、大きく見るとcallbackとcontinuationは類似した概念であることがわかります。

3. Kotlinコンパイラが勝手に変換してくれる

ここまでサスペンションポイントとcontinuationについて見てきました。内部的にはこのように動作するわけですが、誰がこの処理をしてくれるのでしょうか?

Kotlinコンパイラがやってくれます。Kotlinコンパイラはsuspend修飾子がついた関数を見つけると、内部的に次のようにコードを変換します。

// Kotlin
suspend fun createPost(token: Token, item: Item): Post { ... }

// 上記KotlinコードをJavaコードに変換

// Java/JVM
Object createPost(Token token, Item item, Continuation<Post> cont) { ... }

Continuationはインターフェースで、次のように定義されています。

// Kotlin ~v1.2
interface Continuation<in T> {
	val context: CoroutineContext
	fun resume(value: T)
	fun resumeWithException(exception: Throwable)
}

// Kotlin v1.3~
interface Continuation<in T> {
	val context: CoroutineContext
	fun resumeWith(result: Result<T>)
}

resumeWith呼び出しで結果値を渡して元の関数を再開することは確認できるでしょう。しかし、サスペンションポイントごとに一時停止・再開するのに、どこまで実行したかはどうやってわかるのでしょうか?

Kotlinはコード変換時に「label」を置いて示します。コードで例を示します。

suspend fun requestToken() { ... }

suspend fun createPost(token, item) { ... }

suspend fun processPost(post) { ... }

suspend fun postItem(item: Item) {
  // LABEL 0
	val token = requestToken()
  // LABEL 1
	val post = createPost(token, item)
  // LABEL 2
	processPost(post)
}

各サスペンションポイントにコメントでlabel表示をつけました。このようにポイントごとにラベルをつけ、どのラベルかを記憶していれば、次にフローを再開するときに目的のラベル地点から再開できます。

suspend fun postItem(item: Item) {
	switch (label) {
		case 0:
			val token = requestToken()
		case 1:
			val post = createPost(token, item)
		case 2:
			processPost(post)
	}
}

どの段階まで実行したかというlabel情報はどこにあるのでしょうか? 前述の通り、continuationが保持しています。continuationに入っているlabel情報で分岐すれば、一時停止されたコードから実行を再開できます。

fun postItem(item: Item, cont: Continuation) {
	val sm = object : CoroutineImpl { ... }
	when (sm.label) {
		0 -> val token = requestToken(sm)
		1 -> val post = createPost(token, item, sm)
		2 -> processPost(post)
	}
}

ある程度フレームワークが整ってきました。コードを見ると、最初の実行時にcontinuationがなければ作成し、その後サスペンションポイントで呼び出すたびに再利用します。

でも、何か足りない気がしますね。labelを保存するコードがなく、labelが1の場合にcreatePostを呼ぶときに必要なitem値はどこから持ってくるのでしょう? その行が実行されるときはresumeWithで処理が再開されるタイミングなので、元の関数進入時に受け取ったitem値はないはずですが?

その通りです。そのため、上記で述べたようにstateもcontinuationに保存してから、後で取り出して使うコードを追加します。

fun postItem(item: Item, cont: Continuation) {
	val sm = object : CoroutineImpl { ... }
	when (sm.label) {
		0 -> {
			sm.item = item
			sm.label = 1
			val token = requestToken(sm)
		}
		1 -> {
			val item = sm.item
			val token = sm.result as Token
			sm.label = 2
			val post = createPost(token, item, sm)
		}
		2 -> processPost(post)
	}
}

サスペンションポイントを記憶し、stateを管理する作業をsuspendCoroutine()とcontinuationが行います。Kotlinコンパイラがsuspend修飾子のついたコードを処理する際、このように自動的にコードを変形してくれます。

ただし、実際の変換コードがこのように書かれるわけではありません。実際にコンパイルされたバイトコードを逆コンパイルすると複雑なJavaコードが見えますが、読者の理解を助けるためにKotlinコードで例示しました。大きな流れは同じなので、このように理解しても問題ありません。

4. コルーチンがfuture-likeに提供する拡張機能

上記の例で、callbackスタイルのコードをdirectスタイルで書けるようになることはご理解いただけたと思います。しかし、このように書いたからといって特別なパフォーマンス向上があるようには見えず、callbackを使っても見た目が良くないだけで特に問題ないように見えますが、スタイルの違い以外にメリットはないのでしょうか?

コルーチンの利点は、遊休スレッドを効率的に活用できる点です。サスペンションポイントで一時停止させて別の作業を開始する際、これを特定のスレッドプールで実行すれば、元の作業が再開されるまで一時停止した作業が実行されていた元のスレッドは遊休状態となり、他の必要な場所で使用できます。

このように、多数のasync呼び出しを同時に送って一度に処理しようとする際、async呼び出しがsuspendで呼ばれれば、効率的なスレッド活用が可能です。そして、別の利点があります。通常、JVMの様々なasyncライブラリは独自の「Future-like」を提供します。

Guava: ListenableFuture
RxJava: Observable
JDK8: CompletableFuture
…

クラスは異なりますが、基本的にFutureと動作原理は同じクラスです。Kotlinのコルーチンでは、こうしたfuture-likeクラスと連携する拡張コードを提供しています。

これにより、そのfuture-like拡張コードが提供されているライブラリに限り、統一された形で値を取得できます。例えば、コルーチンでawait()を使えば、対象がListenableFuture、CompletableFuture、Observable、Promiseなど何であっても同じ方法で動作します。これはKotlinの拡張(extension)を通じて可能であり、結果的にコルーチンによってライブラリに関係なく同じ方法で連携コードを書けます。

// Guavaを使ったJava関数
public ListenableFuture<Image> guavaLoadImageAsync(String name) { … }

// RxJavaを使ったJava関数
public Observable<Image> rjLoadImageAsync(String name) { … }

// 2つの画像を合成するJava関数
public Image combineImage(Image image1, Image image2) { … }

// Kotlinで使うコード
// 2つの画像を受け取ってCompletableFuture形式で返す
fun combineImagesAsync(name1: String, name2: String): CompletableFuture<Image> = future {
	val future1 = guavaLoadImageAsync(name1)
	val future2 = rjLoadImageAsync(name2)
	combineImages(future1.await(), future2.await())
}

💡 Kotlin Compiler

Kotlinコンパイラはsuspend修飾子がついた関数を変換する。
変換された関数は、サスペンションポイントとstateを管理するcontinuationをcallbackのように利用して、処理を一時停止・再開できる。
一時停止された元のスレッドは、再開されるまで他の作業で活用できるため効率的。
コルーチンが提供するfuture-likeへの拡張により、統一された形で異なるfuture-likeライブラリを扱える。

ここまで、Coroutineの概念とKotlin Compilerの動作原理について紹介しました。また、様々なfuture-likeに提供される拡張機能のおかげで、コルーチンでは同じ形のコードでfuture-likeを扱えます。ありがとうございました。