Candy, Vitamin or Painkiller

He's half man and half machine rides the metal monster

ScalaをGraalVMでネイティブイメージにコンパイルしてAWS Lambdaでサクサク動かす

このエントリは ただの集団 Advent Calendar 2018 の18日目です。昨日は ugdark さんの 職種としてのスクラムマスターを調べてみた でした。

普段はScalaアプリケーションをECSやEC2で動かすことが多いのですが、ものによってはPythonやnode.jsでAWS Lambdaの関数ハンドラを書くこともあります。脳内コンテキストスイッチの抑制を考えるとLambda関数ハンドラもScalaで書きたいところですが、コールドスタート時におけるJVM起動のオーバーヘッドなどもあることから、今までにScalaでLambdaという選択は採ることがありませんでした。

そんな中、先日にGraalVMにネイティブイメージをコンパイルする機能があることを知りました。これを使ってScalaコードを実行可能なバイナリにコンパイルし、Lambdaでサクサクと動かしてみようと思います。

GraalVM

https://www.graalvm.org

GraalVMはJSやPython、Ruby、R、JavaやScala、Kotlin、ClojureといったJVM言語、LLVM言語など様々な言語を実行可能で、ある言語から別の言語をオーバーヘッドなしに呼び出せるpolyglotなVMです。
これはこれで大変面白いのですが、今回注目したいのはネイティブイメージのコンパイルを行う機能です。これによってJVM言語のアプリケーションを実行可能なバイナリにコンパイルすることで、起動時間の短縮やメモリフットプリントの削減などが期待できるようです。
ネイティブイメージの生成時にはAOT(ahead-of-time)コンパイルを行うため、動的にクラスロードを行うアプリケーションはコンパイルできないなどの制約もあるようですが、Lambda関数ハンドラで行うような処理では問題になることは少ないかも知れません。

AWS Lambda カスタムランタイム

aws.amazon.com

先日のre:Invent 2018で発表されたのがAWS Lambdaのカスタムランタイムです。これによって、あらゆる言語、あらゆる処理系でLambda関数ハンドラが書けるようになりました。

docs.aws.amazon.com

docs.aws.amazon.com

上記のドキュメントから雑に要約すると、カスタムランタイムを動かすために必要なものは以下のとおりです。

  • エントリーポイントとなるbootstrapの用意
  • 初期化処理の実装
  • イベントループの実装

bootstrap

エントリーポイントとなるbootstrapという名前の実行可能なファイルを用意します。このファイルはLambdaインスタンス作成時に実行されます。
bootstrapは別のファイルを呼び出して実際の処理を委譲してもよいですし、このファイル単体でイベントループを実行しても大丈夫です。

初期化処理

Lambdaインスタンスの実行時に1回だけ実行される処理です。環境変数経由でセッティングを取得したり、DBとのコネクション作成などの初期化処理を必要に応じて行います。

イベントループの実装

Lambdaインスタンスが受け取るイベントを処理するループを実装します。ひとまず動かすためには、

  • next invocation APIにGETし、リクエストIDやリクエストボディなどが含まれたイベントデータを取得する処理の実装
  • イベントを処理するLambda関数ハンドラの実装
  • 処理結果をinvocation response APIにPOSTする処理を実装
  • エラー時にinvocation error APIにPOSTする処理を実装

あたりが必要です。

カスタムランタイムをScalaで実装する

では上で見たようなカスタムランタイムをScalaで実装します。今回はbootstrapファイルの中でイベントループも実装するようにします。

bootstrapの実装

package bootstrap

import play.api.libs.json._
import scalaj.http._

object Main {

  def main(args: Array[String]): Unit = {
    
    // 初期化処理
    val runtime  = System.getenv("AWS_LAMBDA_RUNTIME_API")
    try {
      println(s"runtime: $runtime")
    } catch {
      case e: Exception =>
        val message = Json.obj("errorMessage" -> e.getMessage, "errorType" -> e.getClass.getName).toString
        Http(s"http://$runtime/2018-06-01/runtime/init/error").postData(message).asString
    }

    // イベントループ
    while (true) {
      // イベントデータの取得
      val HttpResponse(body, _, headers) =
        Http(s"http://$runtime/2018-06-01/runtime/invocation/next").asString
      val requestId = headers("lambda-runtime-aws-request-id").head

      try {
        val name = (Json.parse(body) \ "name").get
        val responseJson = Json.obj("message" -> s"Hello, $name!").toString
　　    
        // レスポンスデータの返却
        Http(s"http://$runtime/2018-06-01/runtime/invocation/$requestId/response").postData(responseJson).asString
      } catch {
        case e: Exception =>
          println(e.getClass.getName)
          println(e.getMessage)
          val message = Json.obj("errorMessage" -> e.getMessage, "errorType" -> e.getClass.getName).toString
          Http(s"http://$runtime/2018-06-01/runtime/invocation/$requestId/error").postData(message).asString
      }
    }
  }
}

next invoke APIなどを呼び出すためにScalaJ-HTTPを、JSONのシリアライズ、デシリアライズにはplay-jsonを使用しています。
やっていることは簡単で、リクエストされたJSON {"name": "todokr"}から{"message": "Hello, todokr!"}というレスポンスのJSONを返しているだけです。

sbt-assemblyプラグインの追加

GraalVMでネイティブイメージを生成するnative-imageコマンドは入力としてFat Jarを受け取ります。 今回、Jarの生成にはsbt-assemblyを使います。project以下にassembly.sbtというファイルを用意し、sbt-assemblyをsbtプラグインの依存関係に追加します。

addSbtPlugin("com.eed3si9n" % "sbt-assembly" % "0.14.9")

GraalVMのDockerfileを用意する

MacではLambda実行環境であるLinux向けのネイティブイメージを生成できないので、GraalVMをDockerで動かして利用します。
(このあたりは@k2nakamuraさんのこの記事に助けられました。ありがとうございます。)

このようなDockerfileを用意します。

FROM oracle/graalvm-ce:1.0.0-rc10
RUN yum  install -y java-1.8.0-openjdk-headless \
    && update-alternatives --set java $JAVA_HOME/bin/java \
        && mv $JAVA_HOME/jre/lib/security/cacerts $JAVA_HOME/jre/lib/security/cacerts.bak \
    && ln -s /usr/lib/jvm/jre-1.8.0/lib/security/cacerts $JAVA_HOME/jre/lib/security/cacerts

CMD tail -f /dev/null

build.sbtを編集する

ネイティブイメージの生成のため、build.sbtを編集します。

assemblyタスクのセッティング

jarのエントリーポイントの指定や生成されるjarの名前をsettings内で指定します。

mainClass in assembly := Some("bootstrap.Main"),
assemblyJarName in assembly := s"scala-graalvm-lambda_${version.value}.jar",

GraalVMのdocker imageをビルド/起動するsbtカスタムタスクを用意

シェルスクリプトでやってもよいのですが、せっかくなのでsbtのカスタムタスクにしてみます。ビルド/起動済みかの判定はgrepのexit codeで判定するなどちょっと泥臭い感じです。scala.sys.process.__が大活躍してます。
runBuildServerタスクはbuildContainerタスクへの依存を宣言しています。

lazy val buildContainer = taskKey[Unit]("Build docker image of GraalVM")
lazy val runBuildServer = taskKey[Unit]("Run GraalVM server for build native image.")

...

    buildContainer := {
      val exitCode = ("docker images" #| "grep graal-build-img").!(ProcessLogger(_ => ()))
      if (exitCode == 1) {
        println("Build GraalVM container...")
        "docker build -f Dockerfile -t graal-build-img .".!
      } else println("Container is already built.")
    },
    runBuildServer := {
      buildContainer.value
      val exitCode = ("docker ps" #| "grep graal-builder$").!(ProcessLogger(_ => ()))
      if (exitCode == 1) {
        println("Start build server...")
        "docker run --name graal-builder -dt graal-build-img:latest".!
      } else println("Build server is already running.")
    },

GraalVMでネイティブイメージをコンパイルし、zipに固めるカスタムタスクの用意

せっかくなのでこちらもシェルスクリプトではなくsbtのカスタムタスクにします。cleanとassembly、runBuildServerに依存を宣言しています。

    nativeCompile := {
      clean.value
      assembly.value
      runBuildServer.value
      val jarName = (assemblyJarName in assembly).value
      (s"docker cp target/scala-2.12/$jarName graal-builder:server.jar" #&&
       "time docker exec graal-builder native-image -H:+ReportUnsupportedElementsAtRuntime -H:EnableURLProtocols=http,https -J-Xmx3G -J-Xms3G --no-server -jar server.jar" #&&
       "docker cp graal-builder:server target/bootstrap" #&&
       "docker cp graal-builder:/opt/graalvm-ce-1.0.0-rc10/jre/lib/amd64/libsunec.so target/libsunec.so" #&&
       "zip -j target/bundle.zip target/bootstrap target/libsunec.so").!
    }

最終的にbuild.sbtはこのようになりました。

github.com

このsbt nativeCompileを実行すると、target/にbundle.zipというzipファイルが生成されるはずです。

AWS Lambdaへのデプロイ

AWSコンソールからLambda関数を作成し、生成したzipをアップロードします。あとはAPI Gatewayと接続するなどいい感じにやっておきます。

f:id:todokr:20181218000738p:plain

実行速度を見てみる

Cloudwatch LogsのREPORTの行だけを抜き出すとこのような感じです。128MBの環境でもInit込みで150msほど、暖まっていると1.2 ~ 1.5msほど(!)で処理できていることが分かります。

REPORT RequestId: aa66a4c8-020b-11e9-8118-1f1b36c60d1d   Init Duration: 122.14 ms    Duration: 153.49 ms Billed Duration: 300 ms Memory Size: 128 MB Max Memory Used: 52 MB
REPORT RequestId: ab5b8c12-020b-11e9-8b92-a1f863edf3ed  Duration: 1.26 ms   Billed Duration: 100 ms Memory Size: 128 MB Max Memory Used: 52 MB
REPORT RequestId: af655404-020b-11e9-abba-97121f20de35  Duration: 1.56 ms   Billed Duration: 100 ms Memory Size: 128 MB Max Memory Used: 52 MB
REPORT RequestId: c70d15c1-020b-11e9-950c-4b268a6b6d47  Duration: 1.22 ms   Billed Duration: 100 ms Memory Size: 128 MB Max Memory Used: 53 MB

このくらいサクサク動くのであれば、プロダクションでもScala x Lambdaが選択肢に挙がってきそうですね。

サンプルコードの全体はこちらにあります。

github.com