FLINTERS Engineer's Blog

FLINTERSのエンジニアによる技術ブログ

Lambda SnapStartを完全に理解してみた

社内ブログは初投稿の中途4ヶ月目の脇田です!

先日ミロゴス社との合同勉強会にて「Lambda SnapStartを完全に理解してみた」というタイトルで発表させてもらいましたので今回はその発表内容についての記事になります。(来月にはLambda SnapStartなにも分からない、になっていることでしょう)

AWS re:Invent 2022で発表されたLambda SnapStartですがどうやらjvm実行環境を選んだ際のコールドスタートをSnapshotを利用して改善してくれるなにやらイカしたやつのようです。

といったもののSnapshotってなんですか?といった状態です、なにかしらの状態をキャッシュしていることは文字通り分かりますがよくわからないものはあまり使いたくないですよね。 ということでまずSnapshotについて深掘りをしてみました。

Snapshot深掘り

AWSオフィシャルブログ、その他諸々を巡回したところ....以下のようなことが見えてきました。

  • LambdaのライフサイクルにおけるInit Pahseの処理後の状態をSnapshotとしてCacheしている
  • Snapshotには恐らく2段階存在している(公式で言及はされていないです)
    • LambdaをホストしているVMを対象としたSnapshot
    • ある時点のJavaプロセスを対象としたSnapshot

Init Phase Snapshot

Lambdaには上記スライド内の画像で説明されているようなライフサイクルが存在します。

SnapStartモードではリクエストが発生した場合、INITをスキップしRESTOREが発生、INIT済みの実行環境から処理が再開されることとなっているようです、実行環境自体をキャッシュしているのでコールドスタート改善も納得ですね。

またLambdaの新しいバージョンを公開した時点でINITが発生し実行環境がスナップショットされるのでSnapshotが作成されていなくてコールドスタートに時間を取られてしまうrequestが発生することもなさそうです。

VM Snapshot

Lambdaの裏側ではFirecrackerという仮想マシンが動いています。 https://aws.amazon.com/jp/blogs/opensource/firecracker-open-source-secure-fast-microvm-serverless

このFirecreckerですが、これ自体に自身のプロセスをスナップショットする機能が備わっているのでおそらくSnapStart機能にはこちらが使われていることが考えられます。

開発者はあまり気にするレイヤーではないと思いますので私はそうなんだなぁ程度に受け止めております。

Javaプロセス Snapshot

ところでCoordinated Restore at Checkpoint (CRaC)はみなさんご存知でしょうか?(私は全然知りませんでした。。。)

CRaCとはJavaプロセスをSnapshotするプロジェクトとなるのですが、後述する資料に出てくるSnapStartで使えるJavaプロセスのランタイムフックではこのCRaCの機能が使われています。てことはCRaCでJavaプロセスのスナップショット作成してるんじゃね??、と考えられそうです。(こちらは公式に言及されていないのであくまで推察に過ぎないですが)

こちらはランタイムフックを利用することもあり開発者が関心を持った方がよさそうに思えます。

開発者が利用する際に気をつけた方がよさそうなこと

Snapshotについて解像度が上がったところで開発者が利用する際に気をつけた方がよさそうなことを調べたので主にCRaC周りですが併せて紹介します。

重い処理はstaticに移してパフォーマンスアップ

CRaCはJavaプロセスの任意タイミングでのスナップショットになります、つまりクラスローディングは終わっていそうなのでstaticな処理はコールされてスナップショット対象となるとみてよさそうです、こちらはAWS公式ブログのベストプラクティスでも紹介されています。

SnapStart の利点を最大限に活用するには、関数ハンドラーではなく、初期化コードで起動待機時間の原因となるクラスをプリロードすることをお勧めします。これにより、重いクラスのロードに関連するレイテンシが呼び出しパスから移動し、SnapStart による起動パフォーマンスが最適化されます。 初期化中にクラスをプリロードできない場合は、ダミーの呼び出しでクラスをプリロードすることをお勧めします。これを行うには、次のペット ストア関数の例に示すように、関数ハンドラー コードを更新します。

private static SpringLambdaContainerHandler<AwsProxyRequest, AwsProxyResponse> handler;
  static {
      try {
          handler = SpringLambdaContainerHandler.getAwsProxyHandler(PetStoreSpringAppConfig.class);

          // Use the onStartup method of the handler to register the custom filter
          handler.onStartup(servletContext -> {
              FilterRegistration.Dynamic registration = servletContext.addFilter("CognitoIdentityFilter", CognitoIdentityFilter.class);
              registration.addMappingForUrlPatterns(EnumSet.of(DispatcherType.REQUEST), false, "/*");
          });

          // Send a fake Amazon API Gateway request to the handler to load classes ahead of time
          ApiGatewayRequestIdentity identity = new ApiGatewayRequestIdentity();
          identity.setApiKey("foo");
          identity.setAccountId("foo");
          identity.setAccessKey("foo");

          AwsProxyRequestContext reqCtx = new AwsProxyRequestContext();
          reqCtx.setPath("/pets");
          reqCtx.setStage("default");
          reqCtx.setAuthorizer(null);
          reqCtx.setIdentity(identity);

          AwsProxyRequest req = new AwsProxyRequest();
          req.setHttpMethod("GET");
          req.setPath("/pets");
          req.setBody("");
          req.setRequestContext(reqCtx);

          Context ctx = new TestContext();
          handler.proxy(req, ctx);


      } catch (ContainerInitializationException e) {
          // if we fail here. We re-throw the exception to force another cold start
          e.printStackTrace();
          throw new RuntimeException("Could not initialize Spring framework", e);
      }
  }

上記のコードではSpringでの例となってます。DI対象のFilterの登録、認証情報の取得は通常であればSpring起動時のDIコンテナ初期化時に処理されるためstaticに移して事前に処理しておきSnapshot対象としている、というところでしょうか。

staticに移したら不味いケース

私はすぐstaticに移せばSnapshotになるんなら詰めれるもの片っ端から詰めてやろ。と思いましたが当然制約があるわけでstaticに移したら不味いケースがいくつかあるようです。

ネットワーク接続周り

こちらもAWS公式ブログのベストプラクティスからの引用です。

関数がスナップショットから再開するときは、必ずネットワーク接続を再確立してください。関数ハンドラでネットワーク接続を再確立することをお勧めします。afterRestore または、ランタイム フックを使用できます。

DataBaseConnection,HttpConnectionなどの確立はSnapshotさせない方がよさそうですね。尚AWS SDKが確立したネットワーク接続に関しては自動で再接続されるとのことです。コネクションに関してライブラリの裏側で勝手にやられるケースなどで要注意が必要そうです。

Lambda SnapStartでは開発者が任意のタイミングでSnapshotを削除することはできない

こちらは冒頭でも触れましたがSnapshotは関数公開時に作成され任意のタイミングで作成、削除ができない仕様となっています。

そのためSnapshot作成時に有効なデータがSnapshotの生存期間中に無効になる、と言ったケースが考えられます。 例:寿命の短い認証トークン、SecretManagerから参照するローテートされうるパスワード

これらの取得処理はstaticで行うと予期せぬ不具合につながりそうですね。

初期化時に行いたい処理があるがSnapshotしたくない。といった今回のような場合はCRaCが提供するランタイムカスタムフックを使うと解決できるようです。

afterRestoreランタイムフックは名前の通りJavaプロセスをRestoreした後に呼び出されるものです、Snapshotさせずこちらで対応してしまうのが良いように思えます。 ただしLambdaではタイムアウトが2秒なので注意が必要です。

ただし、afterRestoreはあくまでSnapshot復元後に呼ばれるものなのでRestoreが存在しないコールドスタート後の実行環境をInvokeされた場合実行環境の寿命によっては都度処理が適切な場合も考えられるので判断が難しそうではあるな、、、、といった印象です。

まとめ

  • SnapshotにはVMを対象としたFirecracker、Javaプロセスを対象としたCRaCがおそらくいる
  • 開発者はCRaCのを意識することはありそう
  • Snapshot対象に処理を移すときはstatic部分
  • なんでもstaticに移すのは良くない
  • CRaCのランタイムフックと使い分けましょう

参考

New – Lambda SnapStart で Lambda 関数を高速化 | Amazon Web Services ブログ

Lambda execution environment - AWS Lambda

Runtime hooks for Lambda SnapStart - AWS Lambda

Best practices for working with Lambda SnapStart - AWS Lambda