背景

筆者はアプリケーションのバッチ処理を基本的にRustで実装しています。そうしたバッチ処理は以下の記事で書いたようにLambdaで実行する場合が多いです。

ただし、実行に15分以上かかるような処理はLambdaで実行できないので以下の記事で書いたようにECS/AWS Batchで実行します。

バッチ処理の実行時間は事前に予想できない場合が多く、また運用する中で処理時間が長くなっていくこともあります。そのためあらかじめLambdaでもECS/AWS Batchでも実行できるようにしておくと状況に応じて処理基盤を選択できて便利です。

概要

この記事ではRustでLambdaとECS/AWS Batchの両方で実行可能なバッチ処理を実装する方法を解説します。

具体的な内容は以下の通りです:

Cargoを使った複数バイナリの作成: 一つのクレートで複数のバイナリを作成し、処理を共通化します
エントリポイントの実装: それぞれのバイナリにLambdaとECS/AWS Batchで実行するためのエントリポイントを実装します
Dockerのマルチステージビルド: それぞれのバイナリを含めたDockerイメージを作成し、ECRにプッシュします

上記のように実装しておくことで、必要に応じてLambdaとECS/AWS Batchを簡単に使い分けることができます。

クレートの構成

まずRustクレートの構成から詳しく見ていきましょう。サンプル用の新しいクレートを作成します。

cargo new sample --lib

Cargoではsrc/binディレクトリにファイルを配置することでそれぞれに対応したバイナリを作成することができます。ここではsrc/bin/lambda.rsとsrc/bin/batch.rsを作成します。前者がLambdaで、後者がECSやAWS Batchで実行するためのエントリポイントとなります。

sample
├─── src
│ ├── bin
│ │ ├── batch.rs
│ │ └── lambda.rs
│ └── lib.rs
└── Cargo.toml

クレートの中で複数のバイナリを定義するためにCargo.tomlの[[bin]]セクションを以下のように記述する必要があります。また[dependencies]に処理の実行に必要なクレートを追加しておきます。

Cargo.toml

[package]
name = "sample"
version = "0.1.0"
edition = "2021"

[[bin]]
name = "sample-lambda"
path = "src/bin/lambda.rs"

[[bin]]
name = "sample-batch"
path = "src/bin/batch.rs"

[dependencies]
tokio = { version = "1", features = ["full"] }
serde_json = "1.0.96"
lambda_runtime = "0.8.0"
aws_lambda_events = "0.10.0"

各バイナリの実装

src/lib.rsには共通の処理内容を記述します。ここではexecute_process関数を定義しています。バッチ処理内容はここに実装します。

lib.rs

use std::error::Error;

pub type OpaqueError = Box<dyn Error + Send + Sync + 'static>;

pub async fn execute_process(target_id: usize) -> Result<(), OpaqueError> {
    // ここにバッチ処理の実装を書く
    Ok(())
}

Lambdaで実行する場合のエントリポイントとなるsrc/bin/lambda.rsは以下のようになります。ここではSQSからメッセージを受け取ることを想定しています。lambda_runtimeクレートを利用してLambdaのハンドラーを定義します。

bin/lambda.rs

use aws_lambda_events::sqs::SqsEventObj;
use lambda_runtime::{service_fn, LambdaEvent};
use sample::{execute_process, OpaqueError};

struct SQSMessageBody {
    target_id: usize
}

#[tokio::main]
async fn main() -> Result<(), lambda_runtime::Error> {
    lambda_runtime::run(service_fn(lambda_handler)).await?;
    Ok(())
}

async fn lambda_handler(
    event: LambdaEvent<SqsEventObj<SQSMessageBody>>,
) -> Result<(), lambda_runtime::Error> {
    let mut records = event.payload.records;
    let sqs_message_body = match records.pop() {
        Some(record) => record.body,
        None => return Err("message has no record".into()),
    };
    match execute_process(sqs_message_body.target_id).await {
        Ok(_) => Ok(()),
        Err(e) => {
            println!("{:}", e);
            Err(e)
        }
    }
}

ECS/AWS Batchで実行する場合のエントリポイントとなるsrc/bin/batch.rsは以下のようになります。ここでは環境変数から処理のパラメータを取得しています。それ以外は通常の実行バイナリです。

bin/batch.rs

use std::env;

use sample::{execute_process, OpaqueError};

const BATCH_MESSAGE_ENVIRONMENT_VARIABLE_NAME: &str = "MESSAGE";

struct BatchMessage {
    target_id: usize
}

#[tokio::main]
async fn main() -> Result<(), OpaqueError> {
    batch_handler().await?;
    Ok(())
}

async fn batch_handler() -> Result<(), OpaqueError> {
    let message: BatchMessage =
        serde_json::from_str(&env::var(BATCH_MESSAGE_ENVIRONMENT_VARIABLE_NAME)?)?;
    match execute_process(message.target_id).await {
        Ok(_) => Ok(()),
        Err(e) => {
            println!("{:}", e);
            Err(e)
        }
    }
}

Dockerイメージのビルド

次に各バイナリを含めたDockerイメージをビルドする手順を見ていきます。マルチステージビルドを活用して、バイナリの作成用イメージとそのバイナリを含めた実行用のイメージを分けて効率よくビルドします。

まずビルダーイメージ(builder)を作成し、それぞれのバイナリをコンパイルします。コンパイルされたバイナリはtarget/releaseディレクトリにCargo.tomlで定義したファイル名で出力されます。ビルドした後はLambda用のイメージ(sample-lambda)とECS/AWS Batch用イメージ(sample-batch)のそれぞれにバイナリをコピーします。

Lambda用のイメージでは、sample-lambdaバイナリを${LAMBDA_RUNTIME_DIR}/bootstrapにコピーします。lambda-handlerコマンドを実行することでLambdaのハンドラーが実行されます。ECS/AWS Batch用のイメージでは、コピーしたsample-batchバイナリを直接実行します。

Dockerfile

FROM amd64/rust:1.79 as builder
WORKDIR /usr/src/backend
COPY . .
RUN cargo build --release

FROM public.ecr.aws/lambda/provided:al2023.2023.11.18.01 as sample-lambda
COPY --from=builder \
    /usr/src/backend/target/release/sample-lambda \
    ${LAMBDA_RUNTIME_DIR}/bootstrap
CMD [ "lambda-handler" ]

FROM amd64/debian:bookworm-slim as sample-batch
RUN apt-get update && \
    apt-get -y install ca-certificates
COPY --from=builder \
    /usr/src/backend/target/release/sample-batch \
    /usr/local/bin/
CMD ["/usr/local/bin/sample-batch"]

このビルドを例えばAWS CodeBuildで行う場合、以下のようなビルドスクリプトを記述します。ソースコードを取得した後、それぞれのイメージをビルドしてECRにプッシュします。対応するECRリポジトリはあらかじめ作成しておきます。CODEBUILD_RESOLVED_SOURCE_VERSION変数を利用してgitのコミットハッシュを取得して、イメージのタグとしています。

buildspec.yml

version: 0.2

# from environment variables
# ECR_NAME={アカウントID}.dkr.ecr.ap-northeast-1.amazonaws.com
# DOCKER_USERNAME={DockerHubのユーザー名}
# DOCKER_TOKEN={DockerHubのトークン}

phases:
    pre_build:
        commands:
            - COMMIT_HASH=$(echo $CODEBUILD_RESOLVED_SOURCE_VERSION | cut -c 1-7)
            - IMAGE_TAG=${COMMIT_HASH}

    build:
        commands:
            - docker login --username $DOCKER_USERNAME --password $DOCKER_TOKEN

            - docker build
              --target sample-lambda
              -t ${ECR_NAME}/sample-lambda:latest
              -t ${ECR_NAME}/sample-lambda:${IMAGE_TAG}
              -f Dockerfile .

            - docker build
              --target sample-batch
              -t ${ECR_NAME}/sample-batch:latest
              -t ${ECR_NAME}/sample-batch:${IMAGE_TAG}
              -f Dockerfile .

    post_build:
        commands:
            - aws ecr get-login-password --region $AWS_DEFAULT_REGION |
              docker login --username AWS --password-stdin $ECR_NAME

            - docker push ${ECR_NAME}/sample-lambda:latest
            - docker push ${ECR_NAME}/sample-lambda:${IMAGE_TAG}

            - docker push ${ECR_NAME}/sample-batch:latest
            - docker push ${ECR_NAME}/sample-batch:${IMAGE_TAG}

この後はECRにプッシュされたイメージをLambdaやECS/AWS Batchにデプロイして実行することができます。

まとめ

この記事ではRustで実装したバッチ処理をLambdaとECS/AWS Batchのどちらでも実行できるように構成する方法を解説しました。ポイントは一つのクレート内でそれぞれのエントリポイントとなるバイナリを作成することです。こうすることで処理を共通で実装し、それぞれの環境に合わせてビルドすることができます。またLambdaでの実行を想定して開発している段階からDockerイメージを利用しておくことで、ECS/AWS Batchへの移行もスムーズに行うことができます。

また実用上はLambdaとECS/AWS Batchのいずれかで実行することになります。その際は両方のためにビルド/デプロイを行う必要はないため、ビルドスクリプトを適宜変更して実行する環境を選んでビルドすると良いでしょう。

We Are Hiring!

VALUESでは「組織の提案力と生産性を最大化する提案ナレッジシェアクラウド」Pitchcraftなどの開発メンバーとしてエンジニアを積極採用中です。

PitchcraftではRustを積極的に採用して開発しています。Rustでの開発に興味がある方はぜひご連絡ください！

参考