locustで複数ユーザー実行でのユニークID的なものを割り振る

locustで各ユーザーにユニークなユーザーIDを割り振りつつ再実行する場合は同じユーザーを利用するのにハマったのでやり方を3パターン紹介します。

参考

locust公式GitHub
locustドキュメント(Distributed load generation)

配列を使用したパターン

配列をworker間で共有していないので複数worker使い場合は工夫する必要がありますが、シンプルかつ簡単に実装できます。

from locust import HttpUser, task, between

USERNAMES = [
    "user1",
    "user2",
    "user3",
]

class WebsiteUser(HttpUser):
    wait_time = between(5, 60)

    def __init__(self, parent):
        self.username = ""
        if len(USERNAMES) > 0
            self.username = USERNAMES.pop()

        super().__init__(parent)

    @task
    def task(self):
        print(self.username)

node間通信を使用したパターン

テスト起動時にテストユーザー数からユニークなユーザーを作成して各workerに均等に割り振ります。
Pythonだけで完結しますがWebUIからのユーザーの増加に対応していないのを注意して使用してください。

from locust import HttpUser, task, events, between
from locust.runners import MasterRunner, WorkerRunner

usernames = []

def setup_test_users(environment, msg, **kwargs):
    # 受信したデータをworkerの変数にセット
    usernames.extend(msg.data)


@events.init.add_listener
def on_locust_init(environment, **_kwargs):
    # node間通信の設定
    if not isinstance(environment.runner, MasterRunner):
        environment.runner.register_message("test_users", setup_test_users)


@events.test_start.add_listener
def on_test_start(environment, **_kwargs):
    # テストユーザーの作成と各workerへ送信
    if not isinstance(environment.runner, WorkerRunner):
        users = []
        for i in range(environment.runner.target_user_count):
            users.append(f"User{i}")

        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

            if i + 1 < worker_count:
                end_index = start_index + chunk_size
            else:
                end_index = len(users)

            data = users[start_index:end_index]

            # workerへデータ送信
            environment.runner.send_message("test_users", data, worker)


class WebsiteUser(HttpUser):
    wait_time = between(2, 5)

    def __init__(self, parent):
        self.username = usernames.pop()
        super().__init__(parent)

    @task
    def task(self):
        print(self.username)

Redisを使用したパターン

WebUIでのユーザーの増加にも対応してnode間通信に比べて簡潔に実装できますがRedisを準備する必要があります。

import redis
from locust import HttpUser, task, between, events
from locust.runners import WorkerRunner

USER_COUNTER_KEY = "user_counter"

# Redisクライアントの取得
redis_client = redis.Redis(host = "[REDIS_HOST]", port = [REDIS_PORT], db = 0)

@events.test_start.add_listener
def on_test_start(environment, **_kwargs):
    # masterでkeyを削除してcounterをリセット
    if not isinstance(environment.runner, WorkerRunner):
        redis_client.delete(USER_COUNTER_KEY)


class WebsiteUser(HttpUser):

    wait_time = between(1, 5)

    def on_start(self):
        # インクリメントでカウントアップして値を取得
        user_count = redis_client.incr(USER_COUNTER_KEY)
        self.username = f"user{user_count}"

    @task
    def task(self):
        print(self.username)

おまけ

「Redisを使用したパターン」を「locustとGKEで負荷テストの実行とローカル実行環境」のGitHubを修正したものをGitHubにあげました。

GitHub

使い方

ローカルでの実行は元記事と変わらないので元記事を確認してください。
GKEでの実行方法は「Google Kubernetes Engine を使用した負荷分散テスト」の手順から
gitリポジトリとredisの設定ファイルの読み込みを追加しただけなので差分のみ記載します。

環境の設定

1.GitHub からサンプル リポジトリのクローンを作成します。

git clone https://github.com/marv-kashiwabarak/locust_redis_user_counter

2.作業ディレクトリをクローニングしたリポジトリに変更します。

cd locust_redis_user_counter

Locust のマスターノードとワーカーノードのデプロイ

1.ターゲット ホストとプロジェクト ID を、locust-master-controller.yaml ファイルと locust-worker-controller.yaml ファイル内のデプロイされたエンドポイントとプロジェクト ID に置き換えます。
locust-redis-controller.yamlも置換します。

sed -i -e "s/\[TARGET_HOST\]/$TARGET/g" kubernetes-config/locust-master-controller.yaml
sed -i -e "s/\[TARGET_HOST\]/$TARGET/g" kubernetes-config/locust-worker-controller.yaml
sed -i -e "s/\[PROJECT_ID\]/$PROJECT/g" kubernetes-config/locust-master-controller.yaml
sed -i -e "s/\[PROJECT_ID\]/$PROJECT/g" kubernetes-config/locust-worker-controller.yaml
sed -i -e "s/\[PROJECT_ID\]/$PROJECT/g" kubernetes-config/locust-redis-controller.yaml

2.Locust のマスターノードとワーカーノードをデプロイします。
redisノードのデプロイも追加

kubectl apply -f kubernetes-config/locust-master-controller.yaml
kubectl apply -f kubernetes-config/locust-master-service.yaml
kubectl apply -f kubernetes-config/locust-worker-controller.yaml
kubectl apply -f kubernetes-config/locust-redis-controller.yaml
kubectl apply -f kubernetes-config/locust-redis-service.yaml

手順の変更は以上になります。
これでGKEで使えます。

間違ってるところや他にいい方法がありましたらコメントください。