[Rust] Momentoで期限付きhtmlをキャッシュしてみる

Introduction

この記事やこの記事では、
AWS Lambda + Momentoの基本的な使い方について紹介してきました。
本稿では、任意のデータを指定時間キャッシュして、
期限付きデータを実現する簡単なサンプルを実装してみます。

Environment

• OS : MacOS 12.4
• rust : 1.61.0
• Momento CLI : 0.20.1

Moment CLIとAWSアカウントはセットアップ済みとします。
また、以前の記事（これとこれ)で作成した
momento-lambdaプロジェクトを元に解説します。

Setup

今回使用するキャッシュ(my_content_cache)は事前にCLIで作成しておきましょう。

% momento cache create --name my_content_cache

Create Momento Example

Cargo.tomlの[dependencies]に下記crateを追加します。

quanta = "0.10.0"
momento = "0.6.0"
once_cell = "1.12.0"
aws-config = "0.15.0"
aws-sdk-secretsmanager = "0.15.0"
serde_json = "1.0.82"
serde_derive = "1.0.138"
serde = "1.0.138"

クライアント初期化

まずはMomentoクライアントの初期化をします。
クライアント初期化は１回ですむように、
static変数としてOnceCellをつかって定義します。
また、MomentoのトークンはAWS Secrets Managerに保存したものを取得します。

const MOMENTO_SECRET_ID: &str = "accounts/MomentoAuthToken";
const SECRET_REGION:&str ="ap-northeast-1";

static momento_client: OnceCell<Mutex<RefCell<SimpleCacheClient>>> = OnceCell::new();

/// クライアントの初期化
async fn init_client() -> Result<SimpleCacheClient,Error> {
    let auth_token = get_momento_auth_token().await?;
    let item_default_ttl_seconds = 60;
    Ok(SimpleCacheClientBuilder::new(
        auth_token,
        NonZeroU64::new(item_default_ttl_seconds).unwrap(),
    )?.build())
}

/// AWS Secrets Managerに保存したトークン取得
async fn get_momento_auth_token() -> Result<String,Error> {
    let shared_config = aws_config::from_env().region(Region::new(SECRET_REGION)).load().await;
    let client = Client::new(&shared_config);
    
    let resp = client.get_secret_value().secret_id(MOMENTO_SECRET_ID).send().await?;
    let json = resp.secret_string().unwrap_or("No value!").to_string();
    let v: Value = serde_json::from_str(&json)?;
    return Ok(v["MOMENTO_AUTH_TOKEN"].as_str().unwrap().to_string());
}

init_client関数はmain関数の中で呼び出します。

#[tokio::main]
async fn main() -> Result<(), Error> {
    let cli:SimpleCacheClient = init_client().await?; 
    let _ = momento_client.set(Mutex::new(RefCell::new(cli)));
    run(service_fn(function_handler)).await
}

キャッシュ登録＆取得の実装

次にキャッシュの登録と取得を実装します。
function_handler内でクエリパラメータactionをチェックし、
saveであればBodyから取得したデータをキャッシュ登録、
getであればクエリパラメータ:keyを使ってキャッシュからデータ取得を行います。

/// キャッシュ保存用構造体
#[derive(Debug,Serialize,Deserialize,Default)]
struct PutParams {
  key:String,
  ttl_sec:u64,
  content_type: String,
  content: String
}

async fn function_handler(event: Request) -> Result<impl IntoResponse, Error> {
　　　　　　　　//キャッシュ名定義
    let cache_name = String::from("my_content_cache");
    
    //Momento Client取得
    let tmp_cache_client = momento_client.get().unwrap().lock().unwrap();
    let mut cache_client = tmp_cache_client.borrow_mut();

    //エラー時のJSON
    let error_json = r#"{"content_type":"text/html","content":"cache Not found"}"#;
　　　　　　　　　//クエリパラメータ
    let param_map = event.query_string_parameters();

    let result = match param_map.first("action") {
        Some("save") => {
            println!("action : save");
            //body(json) to Struct 
            let put_params: PutParams = serde_json::from_slice(event.body().as_ref()).unwrap();
            //struct to json(String)
            let put_params_str: String =serde_json::to_string(&put_params).unwrap();
            // save to momento
            let key = put_params.key;
            match cache_client.set(&cache_name, key.clone(), put_params_str.clone(), NonZeroU64::new(put_params.ttl_sec)).await {
                Ok(_) => {}
                Err(err) => {
                    eprintln!("{}", err);
                }
            };

            Response::builder()
                .status(200)
                .header("content-type", "applicatiom/json")
                .body("{result:'ok'}".to_string())
                .map_err(Box::new)?
        },
        Some("get") => {
            println!("action : get");
            let key = param_map.first("key").unwrap_or("none");
            println!("key : {:?}",key);

           let result_value =  match cache_client.get(&cache_name, key.clone()).await {
                Ok(r) => match r.result {
                    MomentoGetStatus::HIT => String::from_utf8(r.value).unwrap(),
                    _ => error_json.to_string(),
                },
                Err(err) => {
                    eprintln!("{}", err);
                    error_json.to_string()
                }
            };

            let v: Value = serde_json::from_str(&result_value)?;
            Response::builder()
                .status(200)
                .header("content-type", v["content_type"].as_str().unwrap().to_string())
                .body(v["content"].as_str().unwrap().to_string())
                .map_err(Box::new)?
        },
        _ => {
            Response::builder()
                .status(200)
                .header("content-type", "text/html")
                .body("other request!".to_string())
                .map_err(Box::new)?    
        }
    };
    Ok(result)
}

デプロイ＆動作確認

プログラムの修正ができたらビルド＆デプロイをします。

# ビルド
% cargo lambda build 
# デプロイ時にサイズに関連するエラーが出た場合 --releaseをつければ回避できるかも

#アップロード
% cargo lambda deploy --enable-function-url momento-lambda --iam-role <IAM_ROLE_ARN>
? function arn: arn:aws:・・・・・・
? function url: https://yourlambdaurl.lambda-url.your-region.on.aws/

動作確認はcurlを使用します。
適当なキャッシュ登録用json(test.json)を用意します。　　

{"key":"my_key","ttl_sec":30,"content_type":"text/html","content":"<html><head><title>Href Example</title></head><body><h1>Href Example</h1><p><a href='https://dev.classmethod.jp/'> DevelopersIO </a>&nbsp;ClassMethod's 'I've tried it' technology media.</p></body></html>"}

ここではキャッシュのキーを「my_key」、TTLを30秒に設定してHTMLを記述しています。
test.jsonを指定して、LambdaにPOST。

curl -X POST 'https://yourlambdaurl.lambda-url.your-region.on.aws/' -d  @test.json

これで３０秒有効なHTMLコンテンツがキャッシュに登録されました。
登録してから３０秒以内であれば、
↓のURLにアクセスすると有効なHTMLが帰ってきます。

https://yourlambdaurl.lambda-url.your-region.on.aws/?action=get&key=my_key

jsonでcontent typeも設定できるので、
HTML以外でも保存できるはずです。（サイズ制限はありますが）