Amazon Redshift 用の簡易なDB負荷試験ツールを作ってみた

石川覚
2024.10.16
こんにちは、石川覚(@いしかわん)です。
皆さんはAmazon Redshiftに対して負荷試験をどのように行っているのか気になったことはありますか? 私はあります。
今日は気分を変えて、一度やってみたかった 「のんピ(@non97)」さん構文 の挨拶に変えてみました。（のんピさんのブログはコチラです！）
改めまして、AWS事業本部コンサルティング部の石川です。今日は、先日のブログAmazon Redshift SYSビューによる実行クエリの取得 の出力から Amzon Redshift クエリをリプレイする、簡易な負荷試験するツールを作成したいと思います。
https://dev.classmethod.jp/articles/20241015-amazon-redshift-sys-views-for-exec-queries/
 DB負荷試験ツール 代表的なDB負荷試験私が知ってる、代表的なDB負荷試験を紹介します。ご紹介するツールはそれぞれ素晴らしい特徴があり、用途や環境に応じて選択してください。ザックリ解説すると **Apache JMeterは汎用性が高く、HammerDBはデータベース専用の機能が充実、JdbcRunnerは日本語ドキュメント（これ重要！）**が充実してます。
 Apache JMeterhttps://jmeter.apache.org/usermanual/build-db-test-plan.html
Apache JMeterは汎用的な負荷テストツールとして有名ですが、データベースの負荷試験もサポートしています。
様々なプロトコルとデータベースをサポート
使いやすいGUIインターフェース
拡張性が高く、プラグインで機能追加可能
 HammerDBhttps://www.hammerdb.com/
HammerDBはデータベース専用のベンチマークツールです。
主要なデータベース(SQL Server, Oracle, PostgreSQL, MySQL等)をサポート
GUI、CLI、ブラウザインターフェースに対応
リアルタイムおよび集計されたパフォーマンス分析機能、TPC-CとTPC-Hベンチマークの実装を含む
 JdbcRunnerhttps://dbstudy.info/jdbcrunner/
JdbcRunnerは日本製の汎用データベース負荷テストツールです。
各種RDBMSに対応
スクリプトでトランザクションを定義可能
スループットとレスポンスタイムの測定
 AWSが提供するRedshift負荷試験ツール Test DriveAmazon Redshift の Test Drive は、既存のワークロードに対して最適なRedshiftの構成を見つけるためのオープンソースツールです。
ソースのRedshiftのログからクエリを抽出し、ターゲットのRedshiftデータウェアハウスに対して再生します。異なるRedshiftインスタンスタイプに対して、パフォーマンステストと機能テストを並行して実行します。Redshift SpectrumやCOPYコマンドで使用される外部オブジェクトを複製する機能を提供します。
https://aws.amazon.com/jp/blogs/big-data/find-the-best-amazon-redshift-configuration-for-your-workload-using-redshift-test-drive/
https://github.com/aws/redshift-test-drive
 DB負荷試験ツールをなぜ作るのか？DB負荷試験するツールは他にもあるのに車輪の再開発では？、疑問が湧くと思います。
当初は、AWSが提供するRedshift負荷試験ツール Test Driveを検討していたのですが、様々な検証ができる代わりに重厚長大な印象で、お客様の環境に手をいれることを躊躇してしまいました。
また、実際のDB負荷試験では、細かい動きを調整したかったり、負荷をかけるポイント（読み込み or 書き込み、ショートクエリ or ロングクエリ、それらの複合条件）が異なる場合や、ツールを作りながら動きを深く学べるという利点もあります。
以前、Amazon Redshift のConcurrency Scaling や Amazon Redshift Serverless のプライベートプレビューに参加した際には、JavaでDB負荷試験ツール作成しました。JavaVMは、複数スレッド動かすといい感じでCPUコアを割り当ててくれるので限界性能までブン回ってくれました。
https://press.aboutamazon.com/news-releases/news-release-details/aws-announces-new-serverless-options-three-analytics-services/
 簡易DB負荷試験ツールの作成 どのようなユースケースを想定するのか？現行のDC2インスタンスのクラスタからRA3 インスタンスのクラスタに移行のための検証を想定しています。検証の目的は、A/Bテストの結果を踏まえて移行先の適切なノード数やノードタイプを確定させることです。
そのため、現行のDC2インスタンスのスナップショットから、現行のDC2インスタンスと同スペックのクラスタと、移行先のRA3インスタンスを構築して、負荷試験を実施します。負荷試験に利用するクエリは、現行のDC2インスタンスから実際のクエリを収集して、それぞれのクラスタで実行します。
 簡易DB負荷試験ツールの仕様今回のツールは、本番のデータと参照型クエリを用いる想定で、仕様は以下のとおりです。
複数スレッドで並列実行する
PythonのThreadPoolExecutorで実装する

参照型クエリを対象とする
BIツールがスパイクアクセスの要因であるため
COPYやUNLOADを含めると認証情報の書き換えやS3パスをどうするかなどの対応が必要

クエリ実行の結果キャッシュは無効化する
SET enable_result_cache_for_session TO off;を事前に実行
ノードのNVMe-SSDストレージのキャッシュやコンパイルキャッシュは無効化できないので諦める

ロングクエリを想定してコネクションをキープアライブする
keepalives: 1 (TCP keepaliveを有効にする)
keepalives_idle: 30 (30秒間アイドル状態が続いたらkeepaliveメッセージを送信)
keepalives_interval: 10 (keepaliveメッセージの再送間隔を10秒に設定)
keepalives_count: 5 (最大5回keepaliveメッセージを送信)

ログの出力は、標準出力とファイル
以下、その実装です。
import os
import time
import concurrent.futures
import logging
from typing import List
import psycopg2
from psycopg2.extras import DictCursor

# 並列数
MAX_WORKERS = 10

# Redshift接続情報
redshift_config = {
    'host': '<endpoint>',
    'port': 5439,
    'dbname': '<database_name>',
    'user': '<user>',
    'password': '<password>'
    "keepalives": 1,
    "keepalives_idle": 30,
    "keepalives_interval": 10,
    "keepalives_count": 5
}

# SQLファイルが格納されているディレクトリ
SQL_DIR = 'redshift_queries'

# ログファイルの設定
LOG_FILE = 'stress_buster_results.log'

# ロガーの設定
def setup_logger():
    logger = logging.getLogger('load_test')
    logger.setLevel(logging.INFO)

    file_handler = logging.FileHandler(LOG_FILE)
    file_handler.setLevel(logging.INFO)

    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)

    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    file_handler.setFormatter(formatter)
    console_handler.setFormatter(formatter)

    logger.addHandler(file_handler)
    logger.addHandler(console_handler)

    return logger

# SQLファイルを読み込む関数
def load_sql_files() -> List[str]:
    sql_files = []
    for file in sorted(os.listdir(SQL_DIR)):
        if file.endswith('.sql'):
            with open(os.path.join(SQL_DIR, file), 'r') as f:
                sql_files.append(f.read())
    return sql_files

# クエリを実行する関数
def execute_query(query: str, logger: logging.Logger):
    conn = psycopg2.connect(**redshift_config)

    try:
        with conn.cursor(cursor_factory=DictCursor) as cur:
            start_time = time.time()
            cur.execute('SET enable_result_cache_for_session TO off;') # キャッシュの無効化
            cur.execute(query)
            execution_time = time.time() - start_time

            result = cur.fetchone()
            if result:
                result_str = ', '.join([f"{k}: {v}" for k, v in result.items()])
            else:
                result_str = "No results"

            log_message = f"Query executed in {execution_time:.2f} seconds. Result: {result_str}"
            logger.info(log_message)
    except Exception as e:
        logger.error(f"Error executing query: {str(e)}")
    finally:
        conn.close()
#
# Main
#
def main():
    logger = setup_logger()
    sql_queries = load_sql_files()

    logger.info(f"Starting load test with {len(sql_queries)} queries")

    with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = [executor.submit(execute_query, query, logger) for query in sql_queries]
        concurrent.futures.wait(futures)

    logger.info("Load test completed")

if __name__ == "__main__":
    main()

 補足: パッケージのインストール上記のプログラムを動かすには、psycopg2の導入が必要です。そのためには、必要なパッケージとpsycopg2のインストールが必要です。
$ sudo yum install  libpq-devel gcc python3-devel postgresql-devel　
Last metadata expiration check: 0:00:44 ago on Tue 15 Oct 2024 07:36:30 AM UTC.
Dependencies resolved.
================================================================================
 Package                   Arch   Version                     Repository   Size
================================================================================
Installing:
 gcc                       x86_64 11.4.1-2.amzn2023.0.2       amazonlinux  32 M
 libpq-devel               x86_64 16.4-1.amzn2023.0.1         amazonlinux 105 k
 python3-devel             x86_64 3.9.16-1.amzn2023.0.9       amazonlinux 205 k
Installing dependencies:
 amazon-rpm-config         noarch 228-3.amzn2023.0.2          amazonlinux  74 k
 :
 :
  python-srpm-macros-3.9-41.amzn2023.0.5.noarch                                 
  python3-devel-3.9.16-1.amzn2023.0.9.x86_64                                    
  rust-srpm-macros-21-42.amzn2023.0.2.noarch                                    

Complete!
$ sudo pip install psycopg2
Collecting psycopg2
  Downloading psycopg2-2.9.9.tar.gz (384 kB)
     |████████████████████████████████| 384 kB 5.3 MB/s            
  Preparing metadata (setup.py) ... done
Building wheels for collected packages: psycopg2
  Building wheel for psycopg2 (setup.py) ... done
：
：
/root/.cache/pip/wheels/3a/06/25/adb124afd8c8346e45c455f6586f7289cde2b4e339dfbcd9e9
Successfully built psycopg2
Installing collected packages: psycopg2
Successfully installed psycopg2-2.9.9
WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv
 今回のユースケースを想定した手順現行のDC2クラスタのスナップショット取得
現行のDC2クラスタから昨日のツールで、実行したSQLを取得
スナップショットから検証用のDC2クラスタと検証用のRA3クラスタのリストア
検証用のDC2クラスタと検証用のDC2クラスタに対して上記のツールを実行して負荷かける
実行結果を踏まえ、検証用のRA3クラスタのサイジングした後、検証用のRA3クラスタの再検証を繰り返す
 実行結果上記のプログラムを実行すると、実行状態が刻々と出力されます。クエリのエラーはそのまま出力されます。
% python stress_buster.py
2024-10-16 13:17:58,576 - INFO - Starting load test with 18 queries
2024-10-16 13:17:58,914 - INFO - Query executed in 0.06 seconds. Result: seven_days_ago: 2024-10-09 00:00:00
2024-10-16 13:17:58,917 - INFO - Query executed in 0.06 seconds. Result: ?column?: SELECT * FROM table; /* RQEV2-Ym8KY1GKDU */ SELECT column FROM another_table;
2024-10-16 13:17:58,921 - INFO - Query executed in 0.06 seconds. Result: regexp_replace: [hidden] plain A1234 [hidden]
2024-10-16 13:17:58,925 - INFO - Query executed in 0.07 seconds. Result: current_user: root
2024-10-16 13:17:58,926 - INFO - Query executed in 0.07 seconds. Result: ?column?: SELECT * FROM table; /* RQEV2-Ym8KY1GKDU */ SELECT column FROM another_table;
2024-10-16 13:17:58,926 - INFO - Query executed in 0.07 seconds. Result: current_user: root
2024-10-16 13:17:58,928 - INFO - Query executed in 0.07 seconds. Result: seven_days_ago: 2024-10-09 04:17:58.903954+00:00
2024-10-16 13:18:27,024 - ERROR - Error executing query: permission denied for relation stl_querytext
2024-10-16 13:18:27,338 - ERROR - Error executing query: Invalid preceding regular expression prior to repetition operator.  The error occurred while parsing the regular expression fragment: '[^*]*\*+(?>>>HERE>>>:[^/*][^*]'.
DETAIL:  
  -----------------------------------------------
  error:  Invalid preceding regular expression prior to repetition operator.  The error occurred while parsing the regular expression fragment: '[^*]*\*+(?>>>HERE>>>:[^/*][^*]'.
  code:      8002
  context:   T_regexp_init
  query:     -1[child_sequence:1]
  location:  sql_regex_funcs.hpp:191
  process:   padbmaster [pid=1073815989]
  -----------------------------------------------

CONTEXT:  SQL function "regexp_replace" statement 1

2024-10-16 13:18:27,359 - INFO - Load test completed
実行結果は、stress_buster_results.logファイルに追記されますので、後でゆっくり確認できます。
StressBuster % ll
total 40
drwxr-xr-x   6 ishikawa.satoru  staff   192 10 16 13:18 .
drwxr-xr-x  27 ishikawa.satoru  staff   864 10 15 11:43 ..
drwxr-xr-x  20 ishikawa.satoru  staff   640 10 15 17:34 redshift_queries
-rw-r--r--   1 ishikawa.satoru  staff  4026 10 16 16:36 stress_buster.py
-rw-r--r--@  1 ishikawa.satoru  staff  8492 10 16 16:37 stress_buster_results.log
 最後に2つのRedshiftクラスタ間で、簡単にA/Bテストを実施できました。恐らく、思ったよりもコード数が少なく簡単にDB負荷試験ツールが作成できたと思います。コード量が少ないのでカスタマイズも容易です。
今回は、簡易なDB負荷試験ツールを作成しましたが、代表的なDB負荷試験を利用することでも学びがあるはずなので、これらを組み合わせることが理想ではないかと考えております。
なお、昨日のブログは、今日のブログの前座でした。