AWS Lambda with Rustでレスポンスストリーミング

2023.08.23

Introduction

2023年4月にAWS Lambdaがレスポンスペイロードストリーミングのサポートを発表しました
この機能を使うことで、レスポンスデータが利用可能になった時点で
呼び出し元にデータを送信することができます。
その結果、パフォーマンスの向上やLambdaの6MB制限を超えたデータを送信できます。

今回はレスポンスストリーミングをRustで実装し、
大きめの画像を返してみます。

Environment

  • OS : MacOS 13.0.1
  • rust : 1.71.0

AWSアカウントは設定済みとします。  

Setup

ではまず、AWS LambdaをRustで実装してみましょう。
AWS Lambda with Rustを実装するにはcargo-lambdaを使うのが簡単です。

cargoをつかってcargo-lambdaをインストールしましょう。
通常はHomebrewで問題ないみたいです。

% brew tap cargo-lambda/cargo-lambda
% brew install cargo-lambda

私の場合、cargo installでソースからビルドしました。

% cargo install cargo-lambda

次にプロジェクトの作成です。
newサブコマンドでプロジェクトを作成します。
--httpオプションをつけて関数URLとセットで使えるようにしておきます。

% cargo lambda new large-object-streaming --http
% cd  large-object-streaming

ビルドとデプロイも簡単です。
↓のようにすれば簡単にできます。

% cargo lambda build --release
% cargo lambda deploy

サイズがでかくなりがちなので--releaseつけます。 
デプロイが完了したら、AWSコンソールでLambda→関数→作成した関数と進み、
設定→関数URLと選択します。
とりあえず認証タイプはNONEでよいので、関数URLを割り当てましょう。

「https://<発行された関数URK>/?name=streaming!」に
ブラウザでアクセスすれば、
Hello streaming!, this is an AWS Lambda HTTP request
と画面に表示されます。

また、ローカルで動作を確認したい場合、watchコマンドでLambda関数を待ち受け、
invokeサブコマンドを実行します。

% cargo lambda watch
・・・

% cargo lambda invoke --data-file "payload.json"

{"statusCode":200,"headers":{"content-type":"text/html"},"multiValueHeaders":{"content-type":["text/html"]},"body":"Hello syuta, this is an AWS Lambda HTTP request","isBase64Encoded":false}

payload.jsonはLambdaに渡すパラメータファイルです。
ここでは下記のように定義しました。

{
  "queryStringParameters":{
    "name":"syuta"
  }
}

実際にデプロイされているLambdaを呼び出したい場合、
--remoteを指定して実行します。

% cargo lambda invoke --remote --data-file "payload.json" large-object-streaming

S3にbucketと画像を準備

もう一点、S3に適当なbucketを作成し、
1メガくらいの画像と10メガくらいの画像をアップロードしておきます。

これで準備はOKです。
response payload streamingを試してみましょう。

Try

まずは普通にS3から画像を取得してLambdaで返してみます。
まずはCargo.tomlに依存ライブラリを追加。

[dependencies]
lambda_http = "0.8.1"
lambda_runtime = "0.8.1"
tokio = { version = "1", features = ["macros"] }
tracing = { version = "0.1", features = ["log"] }
tracing-subscriber = { version = "0.3", default-features = false, features = ["fmt"] }

aws-config = "0.56.0"
aws-sdk-config = "0.29.0"
aws-sdk-s3 = "0.29.0"
anyhow = "1.0.75"
serde_json = "1.0.105"

そしてmain.rsにS3の画像を返すコードを記述します。

use aws_config::meta::region::RegionProviderChain;
use lambda_http::{run, service_fn, Body, Error, Request, RequestExt, Response};
use aws_sdk_s3::operation::get_object::GetObjectOutput;
use aws_sdk_s3::Client;

//Get Object from S3
async fn get_object(
    client: &Client,
    bucket_name: &str,
    key: &str,
    ) -> Result<GetObjectOutput, anyhow::Error> {

    let object = client
        .get_object()
        .bucket(bucket_name)
        .key(key)
        .send()
        .await?;
    Ok(object)
}

//lambda function
async fn function(event: Request) -> Result<Response<Body>, Error> {

    let region_provider = RegionProviderChain::default_provider().or_else("ap-northeast-1");
    let config = aws_config::from_env().region(region_provider).load().await;
    let client = Client::new(&config);

    let bucket = "<your bucket>";
    let key = "<image file name>";
    let result = get_object(&client,&bucket,&key).await;

    let image = result.unwrap().body.collect().await; 
    let body = Body::from(image.unwrap().to_vec());

    let resp = Response::builder()
        .status(200)
        .header("content-type", "image/jpeg")
        .body(body)
        .map_err(Box::new)?;
    Ok(resp)
}

#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .with_target(false)
        .without_time()
        .init();

    run(service_fn(function)).await
}

ここではLambda実行時に、S3からオブジェクト(さきほどupした画像)を
取得して返しています。
1メガの画像であれば普通に表示できますが、10メガの画像はInternal Server Errorになります。
これが6MBの壁ですね。

Streamingレスポンスを使う

ではStreamingで6MB以上のデータを返せるようにしてみます。
変更点は2つで、
runではなくlambda_runtime::run_with_streaming_response
を使うのと、handler関数の引数をLambdaEvent<Value>にすることです。

use lambda_runtime::{Error, LambdaEvent};
use serde_json::Value;
use aws_sdk_s3::operation::get_object::GetObjectOutput;
use aws_sdk_s3::Client;
use aws_config::meta::region::RegionProviderChain;
use lambda_http::{service_fn, Body, Response};

async fn function_handler_stream(_event: LambdaEvent<Value>) -> Result<Response<Body>, Error> {
//中身は先程のfunctionと同じ
}

#[tokio::main]
async fn main() -> Result<(), Error> {
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .with_target(false)
        .without_time()
        .init();
    lambda_runtime::run_with_streaming_response(service_fn(function_handler_stream)).await
}

ここで少し注意するのがhandlerの引数です。
べつに引数使わないしLambdaEvent<String>にでもしときゃいいだろ、
とかしてたらURLアクセス時に画像が表示されず0バイトのファイルがダウンロードされる現象になったので、
ちゃんとLambdaEvent<serde_json::Value>を引数に指定しましょう。

そして関数URLの呼び出しモードをRESPONSES_STREAMに設定します。
resp-stream

ビルド&デプロイして再度10メガファイルを取得してみましょう。

% cargo lambda build --release
% cargo lambda deploy

今度はブラウザに画像が表示されます。
これで20メガまでは取得可能です。(ソフトリミット)

Summary

今回はRustでStreaming Responseを試してみました。
これで微妙にサイズの大きいデータでも面倒な回り道せずに
Lambdaだけで処理できますね。

References