はじめに

とある案件でLocustという負荷試験ツールのテストシナリオを作成しました。
ユーザーがズラッと並んだリストのファイルから、worker間で重複なく一意のユーザーで実行する方式について紹介したいと思います。

Locust自体の説明については詳細に説明している他のブログ記事などがあるため本記事では紹介しません。
弊社のLocust関連の記事はこちら

背景

事前にユーザー数がわからなかったためスケールできる構成を取る必要がありました。
そのため、今回はMaster/Worker構成を選択しました。
シングルホストの場合は、ユーザーリストを上から順番に実行していけば問題ありません。
ですが、Workerが何十台、何百台となった場合になるべく手間のかからない方式で、ユーザーの一意性を担保したいというのが今回のモチベーションになります。

下記のようにいくつか選択肢を取ることができます。

取りうる選択肢

取りうる選択肢は以下のとおりです。

1.ランダムピック
2.ノード間通信
3.あらかじめworkerにファイルを分割して保持しておく
4.ユーザーリストをRDB、redisなどのデータストアに保持しておき、フラグ管理をする

この中から「2.ノード間通信」を選択しました。
理由は、githubにそのまま流用できるコードがあったためです。
また、今回の案件のケースではそれぞれ不採用の理由があり、1はそもそも一意性を担保できない。3では手作業が発生するため手間がかかる。4では利用料金が発生する。
といったデメリットがあったため、消去法で2を選択しました。

コード

実際にテストシナリオで採用したlocustfile.pyは以下のとおりです。
※MasterとWorkerで共通です。

# 必要なモジュールとクラスのインポート
import os
import boto3
from warrant.aws_srp import AWSSRP
import json
from locust import HttpUser, TaskSet, task, constant, events
from locust.runners import MasterRunner, WorkerRunner

# 環境変数からユーザーリストファイルのパス、パスワード、Cognito関連の設定を取得
user_list_file = os.environ['user_list_file']
password = os.environ['password']
pool_region = os.environ['pool_region']
pool_id = os.environ['pool_id']
client_id = os.environ['client_id']
api_endpoint = os.environ['api_endpoint']

# Cognitoクライアントを初期化
cognito = boto3.client('cognito-idp', region_name=pool_region)
# usernamesの初期化
usernames = []

# ワーカーがユーザー名のリストを受信したときのコールバック関数
def setup_test_users(environment, msg, **kwargs):
    global usernames
    usernames = msg.data
    environment.runner.send_message("acknowledge_users", f"Thanks for the {len(msg.data)} users!")

# ユーザー名のリストの受信確認のコールバック関数
def on_acknowledge(msg, **kwargs):
    print(msg.data)

# Locustの初期化時のイベントリスナー
@events.init.add_listener
def on_locust_init(environment, **_kwargs):
    if not isinstance(environment.runner, MasterRunner):
        environment.runner.register_message("test_users", setup_test_users)
    if not isinstance(environment.runner, WorkerRunner):
        environment.runner.register_message("acknowledge_users", on_acknowledge)

# テストの開始時のイベントリスナー
@events.test_start.add_listener
def on_test_start(environment, **_kwargs):
    if not isinstance(environment.runner, WorkerRunner):
        with open(user_list_file, 'r') as file:
            users = file.read().splitlines()

        worker_count = environment.runner.worker_count
        chunk_size = int(len(users) / worker_count)

        for i, worker in enumerate(environment.runner.clients):
            start_index = i * chunk_size
            end_index = start_index + chunk_size if i + 1 < worker_count else len(users)
            data = users[start_index:end_index]
            environment.runner.send_message("test_users", data, worker)

# Cognitoの認証を行うタスクセットクラス
class CognitoAuthTaskSet(TaskSet):
    def on_start(self):
        self.user_name = usernames.pop()
        self.id_token = self.login()
~~~~~Cognitoの認証処理は割愛~~~~~

# APIエンドポイントをテストするタスク
    @task
    def api_test(self):
        headers = {
            'Authorization': f'Bearer {self.id_token}',
            'Content-Type': 'application/json'
        }

        response = self.client.get(api_endpoint, headers=headers)
        print(response.status_code)
        print(response.text)

# Locustのユーザークラス
class CognitoAuthUser(HttpUser):
    tasks = [CognitoAuthTaskSet]
    wait_time = constant(1.0)

ポイント解説

ほぼgithubから流用したコードになりますが、ポイントを解説します。

• テストの開始時に呼ばれる”on_test_start” で
  • Masterはユーザーリストファイルからユーザー名を読み込み、ワーカーの数に基づいて分割します。
  • 各Workerに割り当てられたユーザー名のチャンクを”test_users”メッセージで送信します。
• Cognitoの認証処理は今回の内容とは関係がないため “usernames” をpopしているところまでを記載しています。
• 割愛されていますが、Cognitoで取得したIDトークンをAuthorizationヘッダに設定してGETリクエストを送信しています。

まとめ

ノード間通信を利用すれば、worker間で重複なくユーザーリストを共有することができました。
テストデータの事前配布についてはcsvファイルを用意しておき、MasterからWorkerに配布するという方式がデメリットもなく良さそうです。