負荷テストツール、Gooseのタスクを書きやすくしたい

Rust製負荷テストツール、Gooseを使っているのですが、サンプルを参考に記述すると、かなり書きづらいので、どうにか自分にとって書きやすくならないか？ というのを検証したメモです。

GooseはPython製の負荷テストツール、Locustにインスパイされて作られたものだそうです。負荷テストツール、とは呼んでいますが、実際にはRustのcrateで、ライブラリとして参照し、main.rsにタスクを記述、コンパイルしてテストを実行します。

自分の環境ではコンパイルでも数秒程度なので試行錯誤してシナリオを作るのにも特に問題ありません。テスト対象とするサーバーは実行時の引数で渡せるし、それ以外の実行時に渡したい情報は環境変数などにしてしまえます。

辛いのは、あるリクエストのレスポンスから情報を取得して実行するタスクの記述方法です。

たとえば、負荷をかけるサイトでログインが必要な場合、以下のようなことをするかなと思います。

1. ログインページへリクエストし、CSRF対策のトークンを取得
2. トークンと、ID、パスワードをPOST
3. ログイン後の画面の情報を元に○○する

こういう処理を記述するには、Drupalへの負荷テスト例が参考になるのですが、結構ネストがエグくなります。

ログインだとこんな感じです。

drupal_loadtest.rs

async fn drupal_loadtest_login(user: &GooseUser) -> GooseTaskResult {
    let mut goose = user.get("/user").await?;

    match goose.response {
        Ok(response) => {
            // Copy the headers so we have them for logging if there are errors.
            let headers = &response.headers().clone();
            match response.text().await {
                Ok(html) => {
                    let re = Regex::new(r#"name="form_build_id" value=['"](.*?)['"]"#).unwrap();
                    let form_build_id = match re.captures(&html) {
                        Some(f) => f,
                        None => {
                            // This will automatically get written to the error log if enabled, and will
                            // be displayed to stdout if `-v` is enabled when running the load test.
                            return user.set_failure(
                                "login: no form_build_id on page: /user page",
                                &mut goose.request,
                                Some(&headers),
                                Some(&html),
                            );
                        }
                    };

レスポンスがOkかどうかでmatch、レスポンスのhtmlが取得できるかでまたmatchするため、レスポンスが正常に読み込めた場合の処理まで、4段階インデントします。さらに、リクエストが成功した時、結果を元にもう1リクエストを送ると、以下みたいになります。

drupal_loadtest.rs

/// Post a comment.
async fn drupal_loadtest_post_comment(user: &GooseUser) -> GooseTaskResult {
    // フォームのIDやトークン取得のためにまず1リクエスト
    let mut goose = user.get(&node_path).await?;
    match goose.response {
        Ok(response) => {
            // OkだったらコメントをPOST
            let headers = &response.headers().clone();
            match response.text().await {
                Ok(html) => {
                    // 中略
                    let request_builder = user.goose_post(&comment_path).await?;
                    let mut goose = user.goose_send(request_builder.form(&params), None).await?;

                    // Verify that the comment posted.
                    match goose.response {
                        Ok(response) => {
                            // Copy the headers so we have them for logging if there are errors.
                            let headers = &response.headers().clone();
                            match response.text().await {
                                Ok(html) => {
                                    // コメントが投稿できてるか確認
                                    if !html.contains(&comment_body) {

割とシンプルな処理なのに、行数、インデントともすごいので、長めのシナリオをかける気がしません。例なので、理解しやすいよう単純に書いてあるのだと思いますが、なんとかインデントが必要とされず、読みやすくかけるようにしてみたいなと思い、まず、このようにstructと関数を定義しました。

take_messages関数は、リクエストとレスポンスを元に、Resultを返します。Messagesには、レスポンスから取得したbody部が含まれているので、成功時の処理はMessagesのresponse_bodyを使って書けます。レスポンスがエラー、もしくはbodyが取得できなかった場合は、GooseTaskErrorが返ります。これは、GooseTaskResultのエラー時と同じstructなので、GooseTaskResultを返すfnの中では?演算子が利用できます。

よって、フォーム画面の取得、フォーム画面で取得したトークンを利用してPOST、みたいな処理は以下のように書けます。

example.rs

pub async fn login_as(user: &GooseUser) -> GooseTaskResult {
    // ログインのフォームを取得
    let mut login_form_messages = login_form(user).await?;
    // レスポンスからトークンを取り出し、ログインのPOST
    let mut login_messages = login(user, &mut login_form_messages).await?;
    // 以下でログイン後のレスポンスを使っていろいろする
    // ...
    Ok(())
}

async fn login_form(user: &GooseUser) -> MessagesResult {
    let goose = user.get("/users/sign_in").await?;
    take_messages(user, goose.request, goose.response).await
}

async fn login(user: &GooseUser, messages: &mut Messages) -> MessagesResult {
    let csrf_token = try_csrf_token(user, messages)?;

    let user = get_env(user, messages, "USER")?;
    let password = get_env(user, messages, "PASSWORD")?;

    let params = [
        ("authenticity_token", &csrf_token[..]),
        ("user[email]", &backyard_user),
        ("user[password]", &backyard_password),
    ];

    let request_builder = user.goose_post("/users/sign_in").await?;
    let goose = user.goose_send(request_builder.form(&params), None).await?;

    take_messages(user, goose.request, goose.response).await
}

上記はレスポンスを元にいろいろしたい場合ですが、シンプルに複数のURLへGETを飛ばすだけであれば、こちらの例が参考になります。この例ではURLの配列を元に複数のURLへのGETをタスクとして定義してますが、自分はタスクの名前、ウェイトも指定したかったので、以下のようにしました。

simple_gets.rs

    let path_and_waits = vec![
        ("top", "/top", 100),
        ("posts index", "/api/v1/posts/index", 200),
        // 中略
    ];

    for (name, path, weight) in path_and_waits {
        taskset = taskset.register_task(
            GooseTask::new(Arc::new(move |user| {
                Box::pin(async move {
                    user.get(path).await?;
                    Ok(())
                })
            }))
            .set_name(name)
            .set_weight(weight)?,
        );
    }