locust.io でカスタムクライアントを作る

Python

お仕事で自社製の分散 KVS の環境構築をやっている。
運用試験や負荷試験を行うんだけど、負荷試験を行うのにツールを使いたい。
Python 製の locust.io を使う。


locust.io を選んだ理由は

  • クライアントを自分で作れる(負荷試験ツールは HTTP でしかお話できないものが多い気がする)
  • Python 製
  • XML とか書きたくない
  • GUI も要らない(CLI から叩きたい)

というのが主な理由。


というわけで locust.io のカスタムクライアントの作り方を備忘録として残す。
オフィシャルサイトに XML-RPC の Client について書かれているので基本的にこれを真似する。
Testing other systems using custom clients — Locust 0.9.0 documentation


今回はサンプルとして Redis を使ってみる。

# -*- coding: utf-8 -*-
import time
import redis
from locust import Locust, events, task, TaskSet

class RedisClient(object):
    def __init__(self, hosts):
        r = redis.StrictRedis(host='localhost', port=6379, db=0)
        self.client = r

    def set(self, key, value):
        start_time = time.time()
        try:
            ret = self.client.set(key, value)
            total_time = int((time.time() - start_time) * 1000)
            if ret is True:
                events.request_success.fire(request_type='redis',
                                            name='set',
                                            response_time=total_time,
                                            response_length=0)
            else:
                events.request_failure.fire(request_type='redis',
                                            name='set',
                                            response_time=total_time,
                                            exception=None)
        except Exception as e:
            total_time = int((time.time() - start_time) * 1000)
            events.request_failure.fire(request_type='redis',
                                        name='set',
                                        response_time=total_time,
                                        exception=e)

        return ret

    def get(self, key):
        start_time = time.time()
        try:
            ret = self.client.get(key)
            total_time = int((time.time() - start_time) * 1000)
            if ret:
                events.request_success.fire(request_type='redis',
                                            name='get',
                                            response_time=total_time,
                                            response_length=0)
            else:
                events.request_failure.fire(request_type='redis',
                                            name='get',
                                            response_time=total_time,
                                            exception=None)
        except Exception as e:
            events.request_failure.fire(request_type='redis',
                                        name='get',
                                        response_time=total_time,
                                        exception=e)
        return ret


class RedisLocust(Locust):
    def __init__(self, *args, **kwargs):
        super(RedisLocust, self).__init__(*args, **kwargs)
        self.client = RedisClient()


class RedisUser(RedisLocust):
    min_wait = 10
    max_wait = 1000

    class task_set(TaskSet):        
        @task(1)
        def set_key(self):
            self.client.set('foo', 'bar')

        @task(1)
        def get_key(self, **kwargs):
            self.client.get('foo')

locust のレポートで一件あたりのリクエストに掛かった時間を出力する必要があるので、Redis の get や set を独自のクライアントでラップしてやる。
locust.io のドキュメントの XML-RPC はメソッド自体をラップしているが、直感的ではなかったので、こうした。
# 多少のオーバーヘッドは出そうな気はする


基本的に get や set を実行前にタイマーを開始し、メソッドを実行後にタイマーを止めて時間を計測して、それを locust.io のイベントを発火してやれば終わり。


カスタムクライアントを書く事が出来るのは本当に素晴らしい。
因みに自社の分散 KVS に locust.io から繋げるために Python 用のクライアントを作ったけど、これはまた別の機会に。