【Kotlin】動的プロキシでsuspend関数を実装する

この記事はKotlin Advent Calendar 2022の5日目の記事になりました。

TL;DR

• インターフェース定義がsuspend関数かつInvocationHandler.invoke関数内でもsuspend関数を呼び出したい場合、素直に書くとinvoke関数が非suspend関数なためコンパイルエラーになる
• suspend () -> ...が(Continuation<*>) -> ...にキャストでき、Continuationもargsから取得できることを利用すれば、この問題は回避できる

前置き

動的プロキシを知っている方は読み飛ばして頂いて大丈夫です。

動的プロキシとは

動的プロキシとは、指定されたインタフェースを実行時に実装するような機能です。

例えば以下のようなことができます。

interface Foo {
    fun foo(): String
}

object FooProxyFactory {
    fun create(): Foo {
        /* 動的にFooを実装して返却 */
    }
}

// 動的に実装したFooを取得
val foo: Foo = FooProxyFactory.create()

これができると、例えば一定の規則に従ったインターフェースを作成するだけで、内部処理は自動実装するようなことができます。

詳しい内容は以下のドキュメントを参照下さい。

動的プロキシを使ってみる

先ほどのFooProxyFactoryを実装してみた内容が以下です。

import java.lang.reflect.InvocationHandler
import java.lang.reflect.Method
import java.lang.reflect.Proxy

class FooProxy(private val text: String) : InvocationHandler {
    override fun invoke(proxy: Any, method: Method, args: Array<out Any>?): Any? {
        // 対象がfoo関数であればtextを返す
        if (method.name == "foo") return text

        // それ以外の場合は一旦エラーにする
        TODO("Not yet implemented")
    }
}

object FooProxyFactory {
    fun create(): Foo {
        val fooClass = Foo::class.java
        val proxy = FooProxy("Hello from dynamic proxy.")

        @Suppress("UNCHECKED_CAST")
        return Proxy.newProxyInstance(
            fooClass.classLoader,
            arrayOf(fooClass),
            proxy
        ) as Foo
    }
}

これを実行すると以下のようになります。
実装クラスを書かずに、動的に実装できていることが分かります。

// 動的に実装したFooを取得
val foo: Foo = FooProxyFactory.create()
println(foo.foo()) // -> Hello from dynamic proxy.がprintされる

軽い解説

InvocationHandlerは、動的プロキシに対する関数呼び出しが有った際に、それに対する処理を実装するクラスです。
FooProxyは、InvocationHandlerを実装しています。

実際に何かを作る際には、invoke関数に処理を書きます。
このサンプルでは、「foo関数が実行された場合、設定された文字列を返却する」という挙動を実装しています。

invoke関数の引数は以下の通りです。

• proxy: メソッドが呼び出されるプロキシ・インスタンス
• method: 呼び出し対象の関数
• args: 引数配列（引数が与えられなかった場合はnullになる(空配列にはならない)）

本題

suspend関数が絡んだ場合に問題になること

インターフェースに定義された関数がsuspend関数で、invoke関数での実装もsuspend関数としたい場合を考えます。
この場合、invoke関数は非suspend関数なため、素直に呼び出すとコンパイルエラーが発生します。

import java.lang.reflect.InvocationHandler

interface Foo {
-   fun foo(): String
+   suspend fun foo(): String // suspend化
}

class FooProxy(private val text: String) : InvocationHandler {
+   // suspend関数で何らかの処理を呼び出したい
+   private suspend fun calcFoo(text: String): String {
+       return text + " suspend."
+   }

    override fun invoke(proxy: Any, method: Method, args: Array<out Any>?): Any? {
-       if (method.name == "foo") return text
+       // calcFooはsuspend関数なため、コンパイルエラーになる！
+       if (method.name == "foo") return calcFoo(text)

        // それ以外の場合は一旦エラーにする
        TODO("Not yet implemented")
    }
}

対策

FooProxyを以下のように変更することで対応できます¹。

import kotlin.coroutines.Continuation

class FooProxy(private val text: String) : InvocationHandler {
    private suspend fun calcFoo(text: String): String {
        return text + " suspend."
    }

    // suspend関数を非suspend関数から呼び出すための関数
    private fun <T> invokeSuspendFunction(
        continuation: Continuation<*>,
        block: suspend () -> T
    ): T = @Suppress("UNCHECKED_CAST") (block as (Continuation<*>) -> T)(continuation)

    override fun invoke(proxy: Any, method: Method, args: Array<out Any>?): Any? {
        if (method.name == "foo")  {
            // args.last()はsuspend関数へ暗黙的に渡されるContinuationの取得
            val continuation = args!!.last() as Continuation<*>
            return invokeSuspendFunction(continuation) {
                calcFoo(text)
            }
        }

        // それ以外の場合は一旦エラーにする
        TODO("Not yet implemented")
    }
}

実行すると以下のようになります。

import kotlinx.coroutines.runBlocking

val foo: Foo = FooProxyFactory.create()
runBlocking { println(foo.foo()) } // -> Hello from dynamic proxy. suspend.がprintされる

軽い解説

FooインターフェースをJavaにデコンパイルすると以下のようになっています。

Fooインターフェースのデコンパイル結果

import kotlin.coroutines.Continuation;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

public interface Foo {
   @Nullable
   Object foo(@NotNull Continuation var1);
}

foo関数は元々引数無し関数でしたが、Java上ではContinuationというパラメータが存在していることが分かります。
Continuationはsuspend関数の制御に利用される値です。
このように、suspend関数は、Java上ではContinuationを引数に取る形になります。

このため、suspend () -> Tは(Continuation<*>) -> Tにキャストすることができます。

suspend () -> Tから(Continuation<*>) -> Tにキャストする様子

private fun <T> invokeSuspendFunction(
    continuation: Continuation<*>,
    block: suspend () -> T
): T = @Suppress("UNCHECKED_CAST") (block as (Continuation<*>) -> T)(continuation)

この値は関数の実行時に渡されるため、invoke関数に渡されるargsから取得することができます。

argsからContinuationを取得する様子

val continuation = args!!.last() as Continuation<*>

これらを組み合わせることで、invoke関数内からのsuspend関数呼び出しを実現しています。

これらの内容は以下のコードを参考に実装しました。

1. 一応自分のユースケースでは動いたのでOKとしていますが、本当にこれで全パターンで大丈夫なのかは確認しきれていません。 ↩